perf: remove per-partition copy, IPC context and flush from the native shuffle writer by andygrove · Pull Request #5908 · apache/datafusion-comet · GitHub
Skip to content

perf: remove per-partition copy, IPC context and flush from the native shuffle writer - #5908

Draft
andygrove wants to merge 3 commits into
apache:mainfrom
andygrove:perf/shuffle-writer-per-partition-overhead
Draft

andygrove wants to merge 3 commits into
apache:mainfrom
andygrove:perf/shuffle-writer-per-partition-overhead

Conversation

@andygrove

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part of #5905 (findings W2 and W3). Does not close it.

Rationale for this change

In a multi-partition native shuffle, every output partition gets its own short-lived BufBatchWriter per spill or finish cycle, so anything that writer does per instance is multiplied by partitions x cycles. Three such costs were found in the review:

  1. Every sub-batch_size chunk was materialized twice. PartitionedBatchIterator already emits maximal batch_size chunks plus one tail per partition. Full chunks bypass the BatchCoalescer, but the tail is copy_rows'd into in-progress builders and re-emitted as the same block. There is never a second batch for the tail to coalesce with in this path, so the copy buys nothing. With many partitions almost every chunk is a tail, so this was a full second copy of the shuffle payload plus one coalescer (with a boxed in-progress builder per column) per partition per cycle.
  2. A fresh IpcWriteContext per writer. arrow-ipc only retains capacity inside the context, so the first block of every partition regrew the body scratch from empty and rebuilt the flatbuffer builder.
  3. A flush and an lseek per partition. finish_partition called stream_position() on the BufWriter<File> to record the partition offset, which flushes the buffer, and then BufBatchWriter::flush flushed it again. Every partition therefore left as its own write(2) plus a seek, so the 1 MiB output buffer never coalesced small partitions.

What changes are included in this PR?

  • BufBatchWriter gains a passthrough mode (new_passthrough) that serializes every batch as its own block, and a drain method that hands buffered bytes to the underlying writer without flushing it. flush is now drain plus a flush. The coalescing mode is unchanged and still used by the single-partition writer, whose inputs can genuinely be small.
  • The per-call scratch Vec<u8> becomes a ShuffleScratch { buffer, ipc_context }, so the task-scoped recycling that already existed for the byte buffer now also covers the IPC context.
  • LocalPartitionWriter multi-partition mode and SpillWriter use passthrough writers with the shared scratch. LocalPartitionWriter tracks the output offset as a running byte total (spill copy bytes plus BufBatchWriter::bytes_written) instead of asking the BufWriter for its position, and flushes the output once in finish_all, where the running total is checked against the actual file position and any mismatch fails the task.

Block boundaries and the file format are unchanged: the multi-partition path wrote one block per chunk before too, just after an extra copy.

Benchmark

shuffle_writer_high_partition (81,920 rows of a 4-column schema, codec None, hash partitioning), Apple Silicon, criterion --baseline:

Partitions Before After Change
200 3.68 ms 2.91 ms -20%
2000 11.95 ms 6.78 ms -43%
8000 26.75 ms 13.25 ms -51%

The main shuffle_writer group (16 partitions, all codecs, range partitioning) improved by 3 to 6 percent and the single-partition cases are unchanged.

How are these changes tested?

  • New unit test passthrough_writes_each_batch_as_its_own_block in buf_batch_writer.rs pins the passthrough mode's block boundaries and row order.
  • New unit test offsets_match_data_file_with_spilled_and_in_memory_batches in local_partition_writer.rs writes two partitions, one with a spilled prefix and in-memory tail, then decodes every block inside each index range and checks the row counts, so a wrong arithmetic offset would fail to parse.
  • Existing datafusion-comet-shuffle tests (127) and clippy pass.
  • CometNativeShuffleSuite (end-to-end native shuffle writes and reads, including spilling cases) passes against the rebuilt native library.

…e shuffle writer

Per-partition writers in a multi-partition shuffle receive maximal batch_size
chunks plus one tail from PartitionedBatchIterator, so the BatchCoalescer they
went through only copied the tail into builders and re-emitted the same block.
Write those batches through verbatim instead.

Thread the IpcWriteContext through the task-scoped scratch alongside the byte
buffer so arrow's body scratch and flatbuffer builder are not regrown from
empty for every partition, and derive partition offsets from a running byte
total instead of stream_position(), which flushed the output BufWriter once
per partition. The total is checked against the file position in finish_all.

shuffle_writer_high_partition: 200 partitions -20%, 2000 -43%, 8000 -51%.
@github-actions github-actions Bot added enhancement New feature or request performance area:shuffle Shuffle (JVM and native) labels Sep 13, 2026
@andygrove
andygrove marked this pull request as draft September 13, 2026 20:40
@andygrove
andygrove marked this pull request as ready for review September 13, 2026 20:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:shuffle Shuffle (JVM and native) enhancement New feature or request performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant