Conversation
…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%.
andygrove
marked this pull request as draft
September 13, 2026 20:40
andygrove
marked this pull request as ready for review
September 13, 2026 20:43
andygrove
marked this pull request as draft
September 13, 2026 23:03
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
BufBatchWriterper spill or finish cycle, so anything that writer does per instance is multiplied bypartitions x cycles. Three such costs were found in the review:batch_sizechunk was materialized twice.PartitionedBatchIteratoralready emits maximalbatch_sizechunks plus one tail per partition. Full chunks bypass theBatchCoalescer, but the tail iscopy_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.IpcWriteContextper 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.lseekper partition.finish_partitioncalledstream_position()on theBufWriter<File>to record the partition offset, which flushes the buffer, and thenBufBatchWriter::flushflushed it again. Every partition therefore left as its ownwrite(2)plus a seek, so the 1 MiB output buffer never coalesced small partitions.What changes are included in this PR?
BufBatchWritergains a passthrough mode (new_passthrough) that serializes every batch as its own block, and adrainmethod that hands buffered bytes to the underlying writer without flushing it.flushis nowdrainplus a flush. The coalescing mode is unchanged and still used by the single-partition writer, whose inputs can genuinely be small.Vec<u8>becomes aShuffleScratch { buffer, ipc_context }, so the task-scoped recycling that already existed for the byte buffer now also covers the IPC context.LocalPartitionWritermulti-partition mode andSpillWriteruse passthrough writers with the shared scratch.LocalPartitionWritertracks the output offset as a running byte total (spill copy bytes plusBufBatchWriter::bytes_written) instead of asking theBufWriterfor its position, and flushes the output once infinish_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:The main
shuffle_writergroup (16 partitions, all codecs, range partitioning) improved by 3 to 6 percent and the single-partition cases are unchanged.How are these changes tested?
passthrough_writes_each_batch_as_its_own_blockinbuf_batch_writer.rspins the passthrough mode's block boundaries and row order.offsets_match_data_file_with_spilled_and_in_memory_batchesinlocal_partition_writer.rswrites 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.datafusion-comet-shuffletests (127) and clippy pass.CometNativeShuffleSuite(end-to-end native shuffle writes and reads, including spilling cases) passes against the rebuilt native library.