andygrove opened a new pull request, #5003: URL: https://github.com/apache/datafusion-comet/pull/5003
## Which issue does this PR close? Addresses the "Every batch is copied an extra time through `BatchCoalescer`, even when already full" item in #5002. That issue is an umbrella tracking several shuffle-writer optimizations, so this PR does not close it. ## Rationale for this change `BufBatchWriter::write` pushes every batch through a `BatchCoalescer`. In arrow 58.3 the normal push path always does `copy_rows` into the in-progress builders, even when the pushed batch is exactly `target_batch_size` and the buffer is empty. The zero-copy bypass in `BatchCoalescer` only activates when `biggest_coalesce_batch_size` is explicitly set, which Comet never does. `PartitionedBatchIterator` already emits exactly `batch_size`-row batches from `interleave_record_batch` (except the tail), so in the multi-partition spill and finish paths, and the single-partition path, essentially all shuffle data was copied twice: once by interleave, once by the coalescer. ## What changes are included in this PR? `BufBatchWriter::write` now bypasses the coalescer when it holds no buffered rows and the incoming batch is already at least `batch_size` rows, serializing the batch directly. The buffered-rows check preserves output ordering: any rows already buffered in the coalescer must be emitted before the new batch, so the bypass only applies when the buffer is empty. Smaller (tail) batches still flow through the coalescer as before. Since `PartitionedBatchIterator` never emits batches larger than `batch_size`, the bypass fires only for exactly-`batch_size` batches, so block boundaries and output bytes are identical to before. This removes one full copy of the shuffle payload. An end-to-end microbenchmark (`shuffle_writer` bench, 8192-row batches, 16 partitions, Lz4) shows roughly a 3% improvement, with a larger relative win expected for the uncompressed path where the copy is a bigger fraction of total work. ## How are these changes tested? Existing `datafusion-comet-shuffle` unit tests continue to pass. A new `test_full_batches_bypass_coalescer` test feeds a mix of full and partial batches and asserts that full batches are emitted verbatim as their own IPC blocks while a partial batch is still coalesced with the following full batch, and that all rows roundtrip in order. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
