You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/08/23 16:33:19 UTC
(datafusion-comet) branch main updated: perf: reuse Arrow IPC compression context across shuffle blocks (#5038)
This is an automated email from the ASF dual-hosted git repository.
sunchao 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 cee09707a9 perf: reuse Arrow IPC compression context across shuffle blocks (#5038)
cee09707a9 is described below
commit cee09707a97a8272fb18a2f5c733a4336df06634
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Mon Aug 24 00:33:14 2026 +0800
perf: reuse Arrow IPC compression context across shuffle blocks (#5038)
---
native/core/src/execution/operators/shuffle_scan.rs | 17 +++++++++++++++--
native/shuffle/benches/shuffle_writer.rs | 10 ++++++++--
native/shuffle/src/shuffle_writer.rs | 17 +++++++++++++++--
native/shuffle/src/spark_unsafe/row.rs | 5 ++++-
native/shuffle/src/writers/buf_batch_writer.rs | 13 +++++++++----
native/shuffle/src/writers/shuffle_block_writer.rs | 19 ++++++++++++-------
6 files changed, 63 insertions(+), 18 deletions(-)
diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs
index 71cfc29ae3..e2c4a37e2a 100644
--- a/native/core/src/execution/operators/shuffle_scan.rs
+++ b/native/core/src/execution/operators/shuffle_scan.rs
@@ -352,6 +352,7 @@ mod tests {
use crate::execution::shuffle::{CompressionCodec, ShuffleBlockWriter};
use arrow::array::{Int32Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
+ use arrow::ipc::writer::CompressionContext;
use arrow::record_batch::RecordBatch;
use datafusion::physical_plan::metrics::Time;
use std::io::Cursor;
@@ -380,7 +381,14 @@ mod tests {
ShuffleBlockWriter::try_new(&batch.schema(), CompressionCodec::Zstd(1)).unwrap();
let mut buf = Cursor::new(Vec::new());
let ipc_time = Time::new();
- writer.write_batch(&batch, &mut buf, &ipc_time).unwrap();
+ writer
+ .write_batch(
+ &batch,
+ &mut buf,
+ &mut CompressionContext::default(),
+ &ipc_time,
+ )
+ .unwrap();
// Read back (skip 16-byte header: 8 compressed_length + 8 field_count)
let bytes = buf.into_inner();
@@ -446,7 +454,12 @@ mod tests {
let mut buf = Cursor::new(Vec::new());
let ipc_time = Time::new();
writer
- .write_batch(&dict_batch, &mut buf, &ipc_time)
+ .write_batch(
+ &dict_batch,
+ &mut buf,
+ &mut CompressionContext::default(),
+ &ipc_time,
+ )
.unwrap();
let bytes = buf.into_inner();
let body = &bytes[16..];
diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs
index c934226b62..2e82f6fc53 100644
--- a/native/shuffle/benches/shuffle_writer.rs
+++ b/native/shuffle/benches/shuffle_writer.rs
@@ -18,6 +18,7 @@
use arrow::array::builder::{Date32Builder, Decimal128Builder, Int32Builder};
use arrow::array::{builder::StringBuilder, Array, Int32Array, RecordBatch};
use arrow::datatypes::{DataType, Field, Schema};
+use arrow::ipc::writer::CompressionContext;
use arrow::row::{RowConverter, SortField};
use criterion::{criterion_group, criterion_main, Criterion};
use datafusion::datasource::memory::MemorySourceConfig;
@@ -53,10 +54,12 @@ fn criterion_benchmark(c: &mut Criterion) {
let ipc_time = Time::default();
let w =
ShuffleBlockWriter::try_new(&batch.schema(), compression_codec.clone()).unwrap();
+ let mut compression_context = CompressionContext::default();
b.iter(|| {
buffer.clear();
let mut cursor = Cursor::new(&mut buffer);
- w.write_batch(&batch, &mut cursor, &ipc_time).unwrap();
+ w.write_batch(&batch, &mut cursor, &mut compression_context, &ipc_time)
+ .unwrap();
});
});
}
@@ -253,12 +256,15 @@ fn schema_encoding_benchmark(c: &mut Criterion) {
let writer =
ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap();
let ipc_time = Time::default();
+ let mut compression_context = CompressionContext::default();
group.bench_function(format!("write_batch ({name} schema)"), |b| {
let mut buffer = vec![];
b.iter(|| {
buffer.clear();
let mut cursor = Cursor::new(&mut buffer);
- writer.write_batch(&batch, &mut cursor, &ipc_time).unwrap();
+ writer
+ .write_batch(&batch, &mut cursor, &mut compression_context, &ipc_time)
+ .unwrap();
});
});
}
diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs
index 81fc082476..46d3b592bc 100644
--- a/native/shuffle/src/shuffle_writer.rs
+++ b/native/shuffle/src/shuffle_writer.rs
@@ -272,6 +272,7 @@ mod test {
use crate::{read_ipc_compressed, ShuffleBlockWriter};
use arrow::array::{Array, Int64Array, StringArray, StringBuilder};
use arrow::datatypes::{DataType, Field, Schema};
+ use arrow::ipc::writer::CompressionContext;
use arrow::record_batch::RecordBatch;
use arrow::row::{RowConverter, SortField};
use datafusion::datasource::memory::MemorySourceConfig;
@@ -301,8 +302,14 @@ mod test {
let mut cursor = Cursor::new(&mut output);
let writer =
ShuffleBlockWriter::try_new(batch.schema().as_ref(), codec.clone()).unwrap();
+ let mut compression_context = CompressionContext::default();
let length = writer
- .write_batch(&batch, &mut cursor, &Time::default())
+ .write_batch(
+ &batch,
+ &mut cursor,
+ &mut compression_context,
+ &Time::default(),
+ )
.unwrap();
assert_eq!(length, output.len());
@@ -341,8 +348,14 @@ mod test {
let mut output = vec![];
let mut cursor = Cursor::new(&mut output);
let writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone()).unwrap();
+ let mut compression_context = CompressionContext::default();
writer
- .write_batch(&batch, &mut cursor, &Time::default())
+ .write_batch(
+ &batch,
+ &mut cursor,
+ &mut compression_context,
+ &Time::default(),
+ )
.unwrap();
let batch2 = read_ipc_compressed(&output[16..]).unwrap();
diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs
index e0ebf8e7d0..77e44024fb 100644
--- a/native/shuffle/src/spark_unsafe/row.rs
+++ b/native/shuffle/src/spark_unsafe/row.rs
@@ -38,6 +38,7 @@ use arrow::array::{
use arrow::compute::cast;
use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
use arrow::error::ArrowError;
+use arrow::ipc::writer::CompressionContext;
use datafusion::physical_plan::metrics::Time;
use datafusion_comet_jni_bridge::errors::CometError;
use jni::sys::{jint, jlong};
@@ -1387,6 +1388,7 @@ pub fn process_sorted_row_partition(
// Single ipc_time accumulates encode + compression time across all batches.
let ipc_time = Time::default();
+ let mut compression_context = CompressionContext::default();
while current_row < row_num {
let n = std::cmp::min(batch_size, row_num - current_row);
@@ -1420,7 +1422,8 @@ pub fn process_sorted_row_partition(
let mut cursor = Cursor::new(&mut frozen);
let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), codec.clone())?;
- written += block_writer.write_batch(&batch, &mut cursor, &ipc_time)?;
+ written +=
+ block_writer.write_batch(&batch, &mut cursor, &mut compression_context, &ipc_time)?;
if let Some(checksum) = &mut current_checksum {
checksum.update(&mut cursor)?;
diff --git a/native/shuffle/src/writers/buf_batch_writer.rs b/native/shuffle/src/writers/buf_batch_writer.rs
index 72db86529d..55d88a4ba4 100644
--- a/native/shuffle/src/writers/buf_batch_writer.rs
+++ b/native/shuffle/src/writers/buf_batch_writer.rs
@@ -18,6 +18,7 @@
use super::ShuffleBlockWriter;
use arrow::array::RecordBatch;
use arrow::compute::kernels::coalesce::BatchCoalescer;
+use arrow::ipc::writer::CompressionContext;
use datafusion::physical_plan::metrics::Time;
use std::borrow::Borrow;
use std::io::{Cursor, Seek, SeekFrom, Write};
@@ -37,6 +38,7 @@ pub(crate) struct BufBatchWriter<S: Borrow<ShuffleBlockWriter>, W: Write> {
writer: W,
buffer: Vec<u8>,
buffer_max_size: usize,
+ compression_context: CompressionContext,
/// Coalesces small batches into target_batch_size before serialization.
/// Lazily initialized on first write to capture the schema.
coalescer: Option<BatchCoalescer>,
@@ -56,6 +58,7 @@ impl<S: Borrow<ShuffleBlockWriter>, W: Write> BufBatchWriter<S, W> {
writer,
buffer: vec![],
buffer_max_size,
+ compression_context: CompressionContext::default(),
coalescer: None,
batch_size,
}
@@ -107,10 +110,12 @@ impl<S: Borrow<ShuffleBlockWriter>, W: Write> BufBatchWriter<S, W> {
) -> datafusion::common::Result<usize> {
let mut cursor = Cursor::new(&mut self.buffer);
cursor.seek(SeekFrom::End(0))?;
- let bytes_written =
- self.shuffle_block_writer
- .borrow()
- .write_batch(batch, &mut cursor, encode_time)?;
+ let bytes_written = self.shuffle_block_writer.borrow().write_batch(
+ batch,
+ &mut cursor,
+ &mut self.compression_context,
+ encode_time,
+ )?;
let pos = cursor.position();
if pos >= self.buffer_max_size as u64 {
let mut write_timer = write_time.timer();
diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs b/native/shuffle/src/writers/shuffle_block_writer.rs
index 7b6846b3ba..65e7bf4570 100644
--- a/native/shuffle/src/writers/shuffle_block_writer.rs
+++ b/native/shuffle/src/writers/shuffle_block_writer.rs
@@ -142,7 +142,12 @@ impl ShuffleBlockWriter {
}
/// Serialize `batch` as a standalone Arrow IPC stream into `out`.
- fn encode_ipc_stream<W: Write>(&self, batch: &RecordBatch, out: &mut W) -> Result<()> {
+ fn encode_ipc_stream<W: Write>(
+ &self,
+ batch: &RecordBatch,
+ out: &mut W,
+ compression_context: &mut CompressionContext,
+ ) -> Result<()> {
let schema_message = match &self.schema_encoding {
SchemaEncoding::Fallback(schema) => {
// Dictionary encoding requires the schema and record batch to share a dictionary
@@ -159,12 +164,11 @@ impl ShuffleBlockWriter {
// Fast path: reuse the pre-encoded schema message and write the record batch manually.
let data_gen = IpcDataGenerator::default();
let mut dictionary_tracker = DictionaryTracker::new(true);
- let mut compression_context = CompressionContext::default();
let (encoded_dictionaries, encoded_batch) = data_gen.encode(
batch,
&mut dictionary_tracker,
&self.write_options,
- &mut compression_context,
+ compression_context,
)?;
debug_assert!(encoded_dictionaries.is_empty());
@@ -180,6 +184,7 @@ impl ShuffleBlockWriter {
&self,
batch: &RecordBatch,
output: &mut W,
+ compression_context: &mut CompressionContext,
ipc_time: &Time,
) -> Result<usize> {
if batch.num_rows() == 0 {
@@ -194,25 +199,25 @@ impl ShuffleBlockWriter {
match &self.codec {
CompressionCodec::None => {
- self.encode_ipc_stream(batch, output)?;
+ self.encode_ipc_stream(batch, output, compression_context)?;
}
CompressionCodec::Lz4Frame => {
let mut wtr = lz4_flex::frame::FrameEncoder::new(&mut *output);
- self.encode_ipc_stream(batch, &mut wtr)?;
+ self.encode_ipc_stream(batch, &mut wtr, compression_context)?;
wtr.finish().map_err(|e| {
DataFusionError::Execution(format!("lz4 compression error: {e}"))
})?;
}
CompressionCodec::Snappy => {
let mut wtr = snap::write::FrameEncoder::new(&mut *output);
- self.encode_ipc_stream(batch, &mut wtr)?;
+ self.encode_ipc_stream(batch, &mut wtr, compression_context)?;
wtr.into_inner().map_err(|e| {
DataFusionError::Execution(format!("snappy compression error: {e}"))
})?;
}
CompressionCodec::Zstd(level) => {
let mut encoder = zstd::Encoder::new(&mut *output, *level)?;
- self.encode_ipc_stream(batch, &mut encoder)?;
+ self.encode_ipc_stream(batch, &mut encoder, compression_context)?;
encoder.finish()?;
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org