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 68489a7bc6 perf: encode shuffle IPC schema once per writer instead of
per block (#5006)
68489a7bc6 is described below
commit 68489a7bc6cf51fa3afcaefdce5494fc93db4a7b
Author: Andy Grove <[email protected]>
AuthorDate: Fri Jul 24 08:48:17 2026 -0600
perf: encode shuffle IPC schema once per writer instead of per block (#5006)
* perf: encode shuffle IPC schema once per writer instead of per block
ShuffleBlockWriter::write_batch used StreamWriter::try_new for every
block, which re-serializes the schema flatbuffer into each block. Since
every block in a shuffle shares one schema, this repeats identical work,
which matters for wide or deeply nested schemas and for shuffles with
many small blocks (high partition counts).
Pre-encode the schema message once in try_new and write it verbatim at
the start of every block, writing the record batch via
IpcDataGenerator::encode. Schemas containing dictionary types keep using
StreamWriter, whose dictionary-id bookkeeping ties schema and batch
encoding together. Output framing is unchanged.
A new criterion benchmark (shuffle_block_schema_encoding) covers a wide
flat schema and a deeply nested schema. Encoding an 8192-row batch is
~7% (wide flat schema) / ~24% (deeply nested schema) faster with the
pre-encoded schema.
Part of #5002.
* chore: trigger CI
* refactor: fold schema-encoding fields into single Option, document header
size
* refactor: address review on shuffle schema pre-encoding
- Fold schema/schema_message fields into a SchemaEncoding enum stating the
Precoded/Fallback invariant and dropping the fast-path deep schema clone
- Detect dictionary types via Schema::flattened_fields instead of a
hand-maintained recursive match, covering all containers arrow recurses
into
- Store one IpcWriteOptions on the writer, shared by the schema and batch
encoders, and pin it to metadata version V5 (which IPC_EOS depends on)
* fix: share shuffle block writer buffers via Arc to avoid per-partition
deep copy
ShuffleBlockWriter is cloned once per output partition when building
per-partition spill writers. After pre-encoding the IPC schema, each clone
deep-copied the schema message (and header) Vec, so a shuffle with a very
large partition count (e.g. the SPARK-48037 test uses 16M+ partitions)
allocated gigabytes and stalled the executor.
Hold the immutable header and pre-encoded schema behind Arc so cloning is an
O(1) refcount bump. Add a regression test asserting the buffers are shared
across clones.
---
native/shuffle/benches/shuffle_writer.rs | 95 +++++++++-
native/shuffle/src/shuffle_writer.rs | 38 ++++
native/shuffle/src/writers/shuffle_block_writer.rs | 207 +++++++++++++++++----
3 files changed, 299 insertions(+), 41 deletions(-)
diff --git a/native/shuffle/benches/shuffle_writer.rs
b/native/shuffle/benches/shuffle_writer.rs
index ea6ad83fdf..2f1f193cda 100644
--- a/native/shuffle/benches/shuffle_writer.rs
+++ b/native/shuffle/benches/shuffle_writer.rs
@@ -201,6 +201,99 @@ fn create_batch(num_rows: usize, allow_nulls: bool) ->
RecordBatch {
.unwrap()
}
+/// Benchmarks the per-block IPC encoding cost (schema + record batch) in
isolation, using the
+/// `None` codec so that compression does not obscure the schema-encoding
cost. Covers a wide flat
+/// schema and a deeply nested schema, where the schema flatbuffer is largest.
+fn schema_encoding_benchmark(c: &mut Criterion) {
+ let mut group = c.benchmark_group("shuffle_block_schema_encoding");
+
+ for (name, batch) in [
+ ("flat", flat_schema_batch(8192)),
+ ("nested", nested_schema_batch(8192)),
+ ] {
+ let writer =
+ ShuffleBlockWriter::try_new(batch.schema().as_ref(),
CompressionCodec::None).unwrap();
+ let ipc_time = Time::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();
+ });
+ });
+ }
+
+ group.finish();
+}
+
+/// A wide flat schema of primitive columns.
+fn flat_schema_batch(num_rows: usize) -> RecordBatch {
+ let num_cols = 50;
+ let fields: Vec<Field> = (0..num_cols)
+ .map(|i| Field::new(format!("c{i}"), DataType::Int32, false))
+ .collect();
+ let schema = Arc::new(Schema::new(fields));
+ let columns: Vec<Arc<dyn Array>> = (0..num_cols)
+ .map(|i| {
+ let values: Vec<i32> = (0..num_rows as i32).map(|r| r +
i).collect();
+ Arc::new(Int32Array::from(values)) as Arc<dyn Array>
+ })
+ .collect();
+ RecordBatch::try_new(schema, columns).unwrap()
+}
+
+/// A schema of several deeply nested struct columns.
+fn nested_schema_batch(num_rows: usize) -> RecordBatch {
+ let num_cols = 4;
+ let depth = 6;
+
+ let mut fields: Vec<Field> = Vec::with_capacity(num_cols);
+ let mut columns: Vec<Arc<dyn Array>> = Vec::with_capacity(num_cols);
+ for col in 0..num_cols {
+ let array = nested_struct_array(num_rows, depth);
+ fields.push(Field::new(
+ format!("col{col}"),
+ array.data_type().clone(),
+ false,
+ ));
+ columns.push(array);
+ }
+ let schema = Arc::new(Schema::new(fields));
+ RecordBatch::try_new(schema, columns).unwrap()
+}
+
+/// Builds a struct array with a multi-field leaf, wrapped in `depth`
single-field structs.
+fn nested_struct_array(num_rows: usize, depth: usize) -> Arc<dyn Array> {
+ use arrow::array::{Float64Array, Int64Array, StringArray, StructArray};
+
+ // Leaf: struct<a: int64, b: utf8, c: float64>
+ let mut array: Arc<dyn Array> = Arc::new(StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Int64, false)),
+ Arc::new(Int64Array::from(vec![1_i64; num_rows])) as Arc<dyn
Array>,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Utf8, false)),
+ Arc::new(StringArray::from(vec!["x"; num_rows])) as Arc<dyn Array>,
+ ),
+ (
+ Arc::new(Field::new("c", DataType::Float64, false)),
+ Arc::new(Float64Array::from(vec![1.0_f64; num_rows])) as Arc<dyn
Array>,
+ ),
+ ]));
+
+ for level in 0..depth {
+ let field = Arc::new(Field::new(
+ format!("s{level}"),
+ array.data_type().clone(),
+ false,
+ ));
+ array = Arc::new(StructArray::from(vec![(field, array)]));
+ }
+ array
+}
+
fn config() -> Criterion {
Criterion::default()
}
@@ -208,6 +301,6 @@ fn config() -> Criterion {
criterion_group! {
name = benches;
config = config();
- targets = criterion_benchmark
+ targets = criterion_benchmark, schema_encoding_benchmark
}
criterion_main!(benches);
diff --git a/native/shuffle/src/shuffle_writer.rs
b/native/shuffle/src/shuffle_writer.rs
index bf296a13a2..12025fe80d 100644
--- a/native/shuffle/src/shuffle_writer.rs
+++ b/native/shuffle/src/shuffle_writer.rs
@@ -313,6 +313,44 @@ mod test {
}
}
+ /// A dictionary-typed column must roundtrip. Such schemas take the
`StreamWriter` fallback
+ /// (rather than the pre-encoded schema path), because dictionary encoding
requires the schema
+ /// and record batch to share a dictionary tracker.
+ #[test]
+ #[cfg_attr(miri, ignore)] // miri can't call foreign function
`ZSTD_createCCtx`
+ fn roundtrip_ipc_dictionary() {
+ use arrow::array::DictionaryArray;
+ use arrow::datatypes::Int32Type;
+
+ let values: Vec<String> = (0..8192).map(|i| format!("v{}", i %
7)).collect();
+ let dict: DictionaryArray<Int32Type> = values.iter().map(|s|
s.as_str()).collect();
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "d",
+ dict.data_type().clone(),
+ false,
+ )]));
+ let batch =
+ RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(dict) as
Arc<dyn Array>])
+ .unwrap();
+
+ for codec in &[
+ CompressionCodec::None,
+ CompressionCodec::Zstd(1),
+ CompressionCodec::Snappy,
+ CompressionCodec::Lz4Frame,
+ ] {
+ let mut output = vec![];
+ let mut cursor = Cursor::new(&mut output);
+ let writer = ShuffleBlockWriter::try_new(schema.as_ref(),
codec.clone()).unwrap();
+ writer
+ .write_batch(&batch, &mut cursor, &Time::default())
+ .unwrap();
+
+ let batch2 = read_ipc_compressed(&output[16..]).unwrap();
+ assert_eq!(batch, batch2);
+ }
+ }
+
#[test]
#[cfg_attr(miri, ignore)] // miri can't call foreign function
`ZSTD_createCCtx`
fn test_single_partition_shuffle_writer() {
diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs
b/native/shuffle/src/writers/shuffle_block_writer.rs
index 5ed5330e3a..7b6846b3ba 100644
--- a/native/shuffle/src/writers/shuffle_block_writer.rs
+++ b/native/shuffle/src/writers/shuffle_block_writer.rs
@@ -16,12 +16,24 @@
// under the License.
use arrow::array::RecordBatch;
-use arrow::datatypes::Schema;
-use arrow::ipc::writer::StreamWriter;
+use arrow::datatypes::{DataType, Schema, SchemaRef};
+use arrow::ipc::writer::{
+ write_message, CompressionContext, DictionaryTracker, IpcDataGenerator,
IpcWriteOptions,
+ StreamWriter,
+};
+use arrow::ipc::MetadataVersion;
use datafusion::common::DataFusionError;
use datafusion::error::Result;
use datafusion::physical_plan::metrics::Time;
-use std::io::{Cursor, Seek, SeekFrom, Write};
+use std::io::{Seek, SeekFrom, Write};
+use std::sync::Arc;
+
+/// Arrow IPC stream end-of-stream marker: the continuation marker
(`0xFFFFFFFF`) followed by a
+/// zero message length. This is what `StreamWriter::finish` emits for
metadata version V5 with
+/// non-legacy framing; a V4-legacy stream would instead emit four zero bytes
with no continuation
+/// marker. The fast path pins its `IpcWriteOptions` to V5 (see
[`ShuffleBlockWriter::try_new`]), so
+/// this constant is valid; if that assumption ever changes, this must be
revisited.
+const IPC_EOS: [u8; 8] = [0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00];
/// Compression algorithm applied to shuffle IPC blocks.
#[derive(Debug, Clone)]
@@ -32,42 +44,136 @@ pub enum CompressionCodec {
Snappy,
}
+/// How a writer encodes the Arrow IPC schema at the start of each block. The
two arms are mutually
+/// exclusive by schema shape, decided once in
[`ShuffleBlockWriter::try_new`], so the retained
+/// state never carries both a pre-encoded message and a schema at the same
time.
+///
+/// Both arms hold their payload behind an `Arc` because a
`ShuffleBlockWriter` is cloned once per
+/// output partition (see `LocalPartitionWriter`), and a shuffle can request
millions of partitions.
+/// Deep-copying the pre-encoded schema per clone would allocate gigabytes for
a large partition
+/// count; sharing it makes cloning O(1).
+#[derive(Clone)]
+enum SchemaEncoding {
+ /// Dictionary-free schema: the IPC schema message, pre-encoded once, is
written verbatim at the
+ /// start of every block instead of being re-serialized per block.
+ Precoded(Arc<[u8]>),
+ /// Schema containing dictionary types: the schema and record batch must
share a dictionary
+ /// tracker, so each block is encoded with `StreamWriter`, which
re-serializes the schema.
+ Fallback(SchemaRef),
+}
+
/// Writes a record batch as a length-prefixed, compressed Arrow IPC block.
+///
+/// Each block is a self-contained Arrow IPC stream (schema message,
dictionary messages, record
+/// batch message, end-of-stream marker). For the common case of a schema with
no dictionary types,
+/// the schema flatbuffer is encoded once in [`Self::try_new`] and written
verbatim at the start of
+/// every block, rather than being re-serialized per block as
`StreamWriter::try_new` would do.
+/// Schemas that contain dictionary types fall back to `StreamWriter`, whose
dictionary-id
+/// bookkeeping ties schema and batch encoding together.
#[derive(Clone)]
pub struct ShuffleBlockWriter {
codec: CompressionCodec,
- header_bytes: Vec<u8>,
+ /// Shared behind an `Arc` so cloning the writer per output partition
stays O(1); see the note
+ /// on [`SchemaEncoding`].
+ header_bytes: Arc<[u8]>,
+ /// IPC options shared by the schema and record-batch encoders so the two
can never diverge;
+ /// pinned to metadata version V5, which [`IPC_EOS`] depends on.
+ write_options: IpcWriteOptions,
+ schema_encoding: SchemaEncoding,
}
impl ShuffleBlockWriter {
pub fn try_new(schema: &Schema, codec: CompressionCodec) -> Result<Self> {
- let header_bytes = Vec::with_capacity(20);
- let mut cursor = Cursor::new(header_bytes);
+ // Header layout: 8-byte block length placeholder + 8-byte field count
(usize) + 4-byte
+ // codec tag = 20 bytes.
+ let mut header_bytes = Vec::with_capacity(20);
- // leave space for compressed message length
- cursor.seek_relative(8)?;
+ // leave space for compressed message length (filled in per block by
write_batch)
+ header_bytes.extend_from_slice(&[0u8; 8]);
// write number of columns because JVM side needs to know how many
addresses to allocate
let field_count = schema.fields().len();
- cursor.write_all(&field_count.to_le_bytes())?;
+ header_bytes.extend_from_slice(&field_count.to_le_bytes());
// write compression codec to header
- let codec_header = match &codec {
+ let codec_header: &[u8] = match &codec {
CompressionCodec::Snappy => b"SNAP",
CompressionCodec::Lz4Frame => b"LZ4_",
CompressionCodec::Zstd(_) => b"ZSTD",
CompressionCodec::None => b"NONE",
};
- cursor.write_all(codec_header)?;
+ header_bytes.extend_from_slice(codec_header);
- let header_bytes = cursor.into_inner();
+ // Shuffle blocks are always written with metadata version V5. Pin it
explicitly rather than
+ // relying on `IpcWriteOptions::default`, because IPC_EOS is only the
correct end-of-stream
+ // marker for V5. Alignment 64 and non-legacy framing match the arrow
defaults.
+ let write_options = IpcWriteOptions::try_new(64, false,
MetadataVersion::V5)?;
+
+ // `flattened_fields` walks the full nested field tree, so this
catches dictionary types
+ // nested under any container arrow itself recurses into when emitting
dictionaries.
+ let has_dictionaries = schema
+ .flattened_fields()
+ .iter()
+ .any(|f| matches!(f.data_type(), DataType::Dictionary(_, _)));
+
+ // For dictionary-free schemas, pre-encode the IPC schema message once
so it does not have
+ // to be re-serialized per block. Dictionary schemas use the
`StreamWriter` fallback.
+ let schema_encoding = if has_dictionaries {
+ SchemaEncoding::Fallback(Arc::new(schema.clone()))
+ } else {
+ let data_gen = IpcDataGenerator::default();
+ let mut dictionary_tracker = DictionaryTracker::new(true);
+ let encoded_schema =
data_gen.schema_to_bytes_with_dictionary_tracker(
+ schema,
+ &mut dictionary_tracker,
+ &write_options,
+ );
+ let mut buf = Vec::new();
+ write_message(&mut buf, encoded_schema, &write_options)?;
+ SchemaEncoding::Precoded(Arc::from(buf))
+ };
Ok(Self {
codec,
- header_bytes,
+ header_bytes: Arc::from(header_bytes),
+ write_options,
+ schema_encoding,
})
}
+ /// Serialize `batch` as a standalone Arrow IPC stream into `out`.
+ fn encode_ipc_stream<W: Write>(&self, batch: &RecordBatch, out: &mut W) ->
Result<()> {
+ let schema_message = match &self.schema_encoding {
+ SchemaEncoding::Fallback(schema) => {
+ // Dictionary encoding requires the schema and record batch to
share a dictionary
+ // tracker, so `StreamWriter` (which re-encodes the schema per
block) is used here.
+ let mut stream_writer =
+ StreamWriter::try_new_with_options(out, schema,
self.write_options.clone())?;
+ stream_writer.write(batch)?;
+ stream_writer.finish()?;
+ return Ok(());
+ }
+ SchemaEncoding::Precoded(schema_message) => schema_message,
+ };
+
+ // 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,
+ )?;
+ debug_assert!(encoded_dictionaries.is_empty());
+
+ out.write_all(schema_message)?;
+ write_message(&mut *out, encoded_batch, &self.write_options)?;
+ out.write_all(&IPC_EOS)?;
+ Ok(())
+ }
+
/// Writes given record batch as Arrow IPC bytes into given writer.
/// Returns number of bytes written.
pub fn write_batch<W: Write + Seek>(
@@ -86,42 +192,30 @@ impl ShuffleBlockWriter {
// write header
output.write_all(&self.header_bytes)?;
- let output = match &self.codec {
+ match &self.codec {
CompressionCodec::None => {
- let mut arrow_writer = StreamWriter::try_new(output,
&batch.schema())?;
- arrow_writer.write(batch)?;
- arrow_writer.finish()?;
- arrow_writer.into_inner()?
+ self.encode_ipc_stream(batch, output)?;
}
CompressionCodec::Lz4Frame => {
- let mut wtr = lz4_flex::frame::FrameEncoder::new(output);
- let mut arrow_writer = StreamWriter::try_new(&mut wtr,
&batch.schema())?;
- arrow_writer.write(batch)?;
- arrow_writer.finish()?;
+ let mut wtr = lz4_flex::frame::FrameEncoder::new(&mut *output);
+ self.encode_ipc_stream(batch, &mut wtr)?;
wtr.finish().map_err(|e| {
DataFusionError::Execution(format!("lz4 compression error:
{e}"))
- })?
+ })?;
}
-
- CompressionCodec::Zstd(level) => {
- let encoder = zstd::Encoder::new(output, *level)?;
- let mut arrow_writer = StreamWriter::try_new(encoder,
&batch.schema())?;
- arrow_writer.write(batch)?;
- arrow_writer.finish()?;
- let zstd_encoder = arrow_writer.into_inner()?;
- zstd_encoder.finish()?
- }
-
CompressionCodec::Snappy => {
- let mut wtr = snap::write::FrameEncoder::new(output);
- let mut arrow_writer = StreamWriter::try_new(&mut wtr,
&batch.schema())?;
- arrow_writer.write(batch)?;
- arrow_writer.finish()?;
+ let mut wtr = snap::write::FrameEncoder::new(&mut *output);
+ self.encode_ipc_stream(batch, &mut wtr)?;
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)?;
+ encoder.finish()?;
+ }
+ }
// fill ipc length
let end_pos = output.stream_position()?;
@@ -134,7 +228,6 @@ impl ShuffleBlockWriter {
)));
}
- // fill ipc length
output.seek(SeekFrom::Start(start_pos))?;
output.write_all(&ipc_length.to_le_bytes())?;
output.seek(SeekFrom::Start(end_pos))?;
@@ -144,3 +237,37 @@ impl ShuffleBlockWriter {
Ok((end_pos - start_pos) as usize)
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::datatypes::{DataType, Field};
+
+ /// A `ShuffleBlockWriter` is cloned once per output partition (see
`LocalPartitionWriter`), and
+ /// a shuffle can request millions of partitions (e.g. the SPARK-48037
test uses more than 16
+ /// million). Cloning must share the immutable header and pre-encoded
schema buffers rather than
+ /// deep-copying them; otherwise a large partition count allocates
gigabytes and stalls the
+ /// executor.
+ #[test]
+ fn clone_shares_buffers() {
+ let schema = Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]);
+ let writer = ShuffleBlockWriter::try_new(&schema,
CompressionCodec::None).unwrap();
+ let cloned = writer.clone();
+
+ assert!(
+ Arc::ptr_eq(&writer.header_bytes, &cloned.header_bytes),
+ "header bytes should be shared across clones, not deep-copied"
+ );
+
+ match (&writer.schema_encoding, &cloned.schema_encoding) {
+ (SchemaEncoding::Precoded(a), SchemaEncoding::Precoded(b)) =>
assert!(
+ Arc::ptr_eq(a, b),
+ "pre-encoded schema should be shared across clones, not
deep-copied"
+ ),
+ _ => panic!("dictionary-free schema should use the pre-encoded
fast path"),
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
