perf: bypass shuffle BatchCoalescer for already-sized batches#5003
perf: bypass shuffle BatchCoalescer for already-sized batches#5003andygrove wants to merge 2 commits into
Conversation
|
Doesn't |
BufBatchWriter pushes every batch through a BatchCoalescer, which by default copies rows into its in-progress builders even when the pushed batch is already exactly target_batch_size and the buffer is empty. PartitionedBatchIterator already emits exactly batch_size-row batches (except the tail), so essentially all shuffle data was copied twice: once by interleave, once by the coalescer. Configure the coalescer's biggest_coalesce_batch_size so batches that are already at least batch_size rows pass straight through without a copy. The limit is batch_size - 1 because the passthrough fires for batches strictly larger than the limit, and our batches are exactly batch_size. This removes one full copy of the shuffle payload with bit-identical output. Part of apache#5002.
051ec4c to
bd91354
Compare
Thanks, updated to use that config. Much simpler. |
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Thanks @andygrove looks good to me
| let block_rows: Vec<usize> = blocks.iter().map(|b| b.num_rows()).collect(); | ||
| assert_eq!( | ||
| block_rows, | ||
| vec![100, 100, 100, 30], |
There was a problem hiding this comment.
Could we add a case with a batch larger than batch_size, like 150 rows with batch_size = 100 ?
I think vec![100, 100, 100, 30] is what the coalescer alone would produce, so this test would still pass without the fast path.
But I'm not really sure about this. Thanks @andygrove
There was a problem hiding this comment.
Thanks for the review @0lai0. I added a test based on your feedback.
The prior test's expected block boundaries [100, 100, 100, 30] are what the coalescer alone would produce, so it passed without the fast path. Add a 150-row batch (> batch_size): the coalescer caps output at target_batch_size and would split it into a 100-row block plus 50 buffered rows, so a single 150-row block can only appear when the passthrough fires.
True for multi-partition and spill. For single-partition, a |
| @@ -28,7 +28,9 @@ use std::io::{Cursor, Seek, SeekFrom, Write}; | |||
| /// | |||
| /// Small batches are coalesced using Arrow's [`BatchCoalescer`] before serialization, | |||
| /// producing exactly `batch_size`-row output batches to reduce per-block IPC schema overhead. | |||
There was a problem hiding this comment.
"producing exactly batch_size-row output batches" is not accurate once a >= batch_size batch is bypassed verbatim, which the single-partition path already does today. Could soften to "at least batch_size" and mention the passthrough.
4 participants
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::writepushes every batch through aBatchCoalescer. In arrow 58.3 the normal push path always doescopy_rowsinto the in-progress builders, even when the pushed batch is exactlytarget_batch_sizeand the buffer is empty. The zero-copy bypass inBatchCoalesceronly activates whenbiggest_coalesce_batch_sizeis explicitly set, which Comet never does.PartitionedBatchIteratoralready emits exactlybatch_size-row batches frominterleave_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::writenow bypasses the coalescer when it holds no buffered rows and the incoming batch is already at leastbatch_sizerows, 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
PartitionedBatchIteratornever emits batches larger thanbatch_size, the bypass fires only for exactly-batch_sizebatches, 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_writerbench, 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-shuffleunit tests continue to pass. A newtest_full_batches_bypass_coalescertest 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.