This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new b75c570748 perf: bypass shuffle BatchCoalescer for already-sized
batches (#5003)
b75c570748 is described below
commit b75c5707489f51485b15fa1d5079948e4d3b9ef2
Author: Andy Grove <[email protected]>
AuthorDate: Thu Jul 23 19:22:17 2026 -0600
perf: bypass shuffle BatchCoalescer for already-sized batches (#5003)
---
native/shuffle/src/shuffle_writer.rs | 79 ++++++++++++++++++++++++++
native/shuffle/src/writers/buf_batch_writer.rs | 27 +++++++--
2 files changed, 100 insertions(+), 6 deletions(-)
diff --git a/native/shuffle/src/shuffle_writer.rs
b/native/shuffle/src/shuffle_writer.rs
index 8d1336dc1c..bf296a13a2 100644
--- a/native/shuffle/src/shuffle_writer.rs
+++ b/native/shuffle/src/shuffle_writer.rs
@@ -924,6 +924,85 @@ mod test {
.sum()
}
+ /// Batches that are already at least `batch_size` rows should bypass the
+ /// coalescer, being written one-block-per-batch, while smaller batches are
+ /// still coalesced. In both cases the data must roundtrip exactly, in
order.
+ #[test]
+ #[cfg_attr(miri, ignore)]
+ fn test_full_batches_bypass_coalescer() {
+ use crate::writers::BufBatchWriter;
+ use arrow::array::Int32Array;
+
+ let schema = Arc::new(Schema::new(vec![Field::new("v",
DataType::Int32, false)]));
+ let batch_size = 100;
+
+ let make_batch = |start: i32, len: i32| {
+ let values: Vec<i32> = (start..start + len).collect();
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(Int32Array::from(values)) as Arc<dyn Array>],
+ )
+ .unwrap()
+ };
+
+ // Sequence: a full batch (bypass), an oversized batch (bypass,
emitted whole),
+ // a partial batch (buffered), then a full batch (must not bypass
because the
+ // coalescer still holds buffered rows).
+ //
+ // The 150-row batch is the case that distinguishes the fast path: the
+ // coalescer alone would split it into a 100-row block plus 50
buffered rows,
+ // so it can only appear as a single 150-row block when the
passthrough fires.
+ let inputs = vec![
+ make_batch(0, batch_size), // rows 0..100 -> bypass (100-row
block)
+ make_batch(batch_size, 150), // rows 100..250 -> oversized bypass
(150-row block)
+ make_batch(250, 30), // rows 250..280 -> buffered
+ make_batch(280, batch_size), // rows 280..380 -> coalesced with
buffered 30
+ ];
+
+ let codec = CompressionCodec::Lz4Frame;
+ let encode_time = Time::default();
+ let write_time = Time::default();
+
+ let mut output = Vec::new();
+ {
+ let mut writer = ShuffleBlockWriter::try_new(schema.as_ref(),
codec.clone()).unwrap();
+ let mut buf_writer = BufBatchWriter::new(
+ &mut writer,
+ Cursor::new(&mut output),
+ 1024 * 1024,
+ batch_size as usize,
+ );
+ for batch in &inputs {
+ buf_writer.write(batch, &encode_time, &write_time).unwrap();
+ }
+ buf_writer.flush(&encode_time, &write_time).unwrap();
+ }
+
+ let blocks = read_all_ipc_batches(&output);
+
+ // The full and oversized batches are emitted verbatim as their own
blocks,
+ // then the partial+full pair coalesces into a full block plus a
30-row tail.
+ let block_rows: Vec<usize> = blocks.iter().map(|b|
b.num_rows()).collect();
+ assert_eq!(
+ block_rows,
+ vec![100, 150, 100, 30],
+ "unexpected block boundaries"
+ );
+
+ // All 380 rows must roundtrip in order (0..380).
+ let mut roundtripped = Vec::new();
+ for block in &blocks {
+ let col = block
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ roundtripped.extend(col.values().iter().copied());
+ }
+ let expected: Vec<i32> = (0..380).collect();
+ assert_eq!(roundtripped, expected, "rows not preserved in order");
+ }
+
#[test]
#[cfg_attr(miri, ignore)]
fn test_empty_schema_shuffle_writer() {
diff --git a/native/shuffle/src/writers/buf_batch_writer.rs
b/native/shuffle/src/writers/buf_batch_writer.rs
index cfddb46539..72db86529d 100644
--- a/native/shuffle/src/writers/buf_batch_writer.rs
+++ b/native/shuffle/src/writers/buf_batch_writer.rs
@@ -26,9 +26,12 @@ use std::io::{Cursor, Seek, SeekFrom, Write};
/// The record batches were first written by ShuffleBlockWriter into an
internal buffer.
/// Once the buffer exceeds the max size, the buffer will be flushed to the
writer.
///
-/// Small batches are coalesced using Arrow's [`BatchCoalescer`] before
serialization,
-/// producing exactly `batch_size`-row output batches to reduce per-block IPC
schema overhead.
-/// The coalescer is lazily initialized on the first write.
+/// Small batches are coalesced using Arrow's [`BatchCoalescer`] before
serialization, reducing
+/// per-block IPC schema overhead. Output batches hold at least `batch_size`
rows, apart from the
+/// remainder emitted on flush. The coalescer is lazily initialized on the
first write and
+/// configured (via `biggest_coalesce_batch_size`) to pass batches that are
already at least
+/// `batch_size` rows straight through, verbatim and without copying them, so
an oversized input
+/// batch is written as a single oversized block.
pub(crate) struct BufBatchWriter<S: Borrow<ShuffleBlockWriter>, W: Write> {
shuffle_block_writer: S,
writer: W,
@@ -64,9 +67,21 @@ impl<S: Borrow<ShuffleBlockWriter>, W: Write>
BufBatchWriter<S, W> {
encode_time: &Time,
write_time: &Time,
) -> datafusion::common::Result<usize> {
- let coalescer = self
- .coalescer
- .get_or_insert_with(|| BatchCoalescer::new(batch.schema(),
self.batch_size));
+ let batch_size = self.batch_size;
+ let coalescer = self.coalescer.get_or_insert_with(|| {
+ // Enable BatchCoalescer's zero-copy passthrough for batches that
are already big
+ // enough, so we don't `copy_rows` the whole batch into the
in-progress builders just
+ // to re-emit a same-sized batch. The passthrough fires for
batches strictly larger
+ // than the limit, so set it to `batch_size - 1` to include
batches of exactly
+ // `batch_size`, which is what `PartitionedBatchIterator` emits
(except the tail),
+ // removing a full copy of the shuffle payload. Block boundaries
are unchanged for
+ // that iterator since it never emits more than `batch_size` rows.
The
+ // single-partition path routes `>= batch_size` batches here
directly, so those now
+ // form one oversized block rather than being split at
`batch_size`. Rows are written
+ // in the same order either way.
+ BatchCoalescer::new(batch.schema(), batch_size)
+
.with_biggest_coalesce_batch_size(Some(batch_size.saturating_sub(1)))
+ });
coalescer.push_batch(batch.clone())?;
// Drain completed batches into a local vec so the coalescer borrow
ends
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
