From 6fc7f7cbd186f97dbd5c86680c5738e52787651a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 23 Jul 2026 10:21:50 -0600 Subject: [PATCH 1/5] 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. --- native/shuffle/benches/shuffle_writer.rs | 95 ++++++++++- native/shuffle/src/shuffle_writer.rs | 38 +++++ .../src/writers/shuffle_block_writer.rs | 150 +++++++++++++----- 3 files changed, 243 insertions(+), 40 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 = (0..num_cols) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, false)) + .collect(); + let schema = Arc::new(Schema::new(fields)); + let columns: Vec> = (0..num_cols) + .map(|i| { + let values: Vec = (0..num_rows as i32).map(|r| r + i).collect(); + Arc::new(Int32Array::from(values)) as Arc + }) + .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 = Vec::with_capacity(num_cols); + let mut columns: Vec> = 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 { + use arrow::array::{Float64Array, Int64Array, StringArray, StructArray}; + + // Leaf: struct + let mut array: Arc = Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("a", DataType::Int64, false)), + Arc::new(Int64Array::from(vec![1_i64; num_rows])) as Arc, + ), + ( + Arc::new(Field::new("b", DataType::Utf8, false)), + Arc::new(StringArray::from(vec!["x"; num_rows])) as Arc, + ), + ( + Arc::new(Field::new("c", DataType::Float64, false)), + Arc::new(Float64Array::from(vec![1.0_f64; num_rows])) as Arc, + ), + ])); + + 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 8d1336dc1c..d5bacea8f3 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 = (0..8192).map(|i| format!("v{}", i % 7)).collect(); + let dict: DictionaryArray = 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]) + .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..84ee45a1ef 100644 --- a/native/shuffle/src/writers/shuffle_block_writer.rs +++ b/native/shuffle/src/writers/shuffle_block_writer.rs @@ -16,12 +16,20 @@ // 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 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, matching what `StreamWriter::finish` emits for metadata version V5. +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 +40,119 @@ pub enum CompressionCodec { Snappy, } +/// Returns true if `data_type` is, or nests, a dictionary type. +fn contains_dictionary(data_type: &DataType) -> bool { + match data_type { + DataType::Dictionary(_, _) => true, + DataType::List(f) + | DataType::LargeList(f) + | DataType::FixedSizeList(f, _) + | DataType::Map(f, _) + | DataType::RunEndEncoded(_, f) => contains_dictionary(f.data_type()), + DataType::Struct(fields) => fields.iter().any(|f| contains_dictionary(f.data_type())), + DataType::Union(fields, _) => fields + .iter() + .any(|(_, f)| contains_dictionary(f.data_type())), + _ => false, + } +} + /// 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, + schema: SchemaRef, + /// Pre-encoded Arrow IPC schema message, written at the start of every block. Only used when + /// the schema has no dictionary types. + schema_message: Vec, + /// Whether the schema contains any dictionary types (see [`Self::encode_ipc_stream`]). + has_dictionaries: bool, } impl ShuffleBlockWriter { pub fn try_new(schema: &Schema, codec: CompressionCodec) -> Result { - let header_bytes = Vec::with_capacity(20); - let mut cursor = Cursor::new(header_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)?; - - let header_bytes = cursor.into_inner(); + header_bytes.extend_from_slice(codec_header); + + // Pre-encode the IPC schema message once so it does not have to be re-serialized per block. + let options = IpcWriteOptions::default(); + 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, + &options, + ); + let mut schema_message = Vec::new(); + write_message(&mut schema_message, encoded_schema, &options)?; + + let has_dictionaries = schema + .fields() + .iter() + .any(|f| contains_dictionary(f.data_type())); Ok(Self { codec, header_bytes, + schema: Arc::new(schema.clone()), + schema_message, + has_dictionaries, }) } + /// Serialize `batch` as a standalone Arrow IPC stream into `out`. + fn encode_ipc_stream(&self, batch: &RecordBatch, out: &mut W) -> Result<()> { + if self.has_dictionaries { + // 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(out, &self.schema)?; + stream_writer.write(batch)?; + stream_writer.finish()?; + return Ok(()); + } + + // Fast path: reuse the pre-encoded schema message and write the record batch manually. + let options = IpcWriteOptions::default(); + 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, + &options, + &mut compression_context, + )?; + debug_assert!(encoded_dictionaries.is_empty()); + + out.write_all(&self.schema_message)?; + write_message(&mut *out, encoded_batch, &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( @@ -86,42 +171,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 +207,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))?; From 0b1d5aeb41cf298cb11b0ae40cb87f542eacab65 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 23 Jul 2026 11:35:49 -0600 Subject: [PATCH 2/5] chore: trigger CI From 2d6d05dc03b6c6edf1c9d3a9c1a018bde60cecbd Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 23 Jul 2026 12:22:29 -0600 Subject: [PATCH 3/5] refactor: fold schema-encoding fields into single Option, document header size --- .../src/writers/shuffle_block_writer.rs | 51 +++++++++++-------- 1 file changed, 30 insertions(+), 21 deletions(-) diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs b/native/shuffle/src/writers/shuffle_block_writer.rs index 84ee45a1ef..7ea407ee60 100644 --- a/native/shuffle/src/writers/shuffle_block_writer.rs +++ b/native/shuffle/src/writers/shuffle_block_writer.rs @@ -70,15 +70,18 @@ pub struct ShuffleBlockWriter { codec: CompressionCodec, header_bytes: Vec, schema: SchemaRef, - /// Pre-encoded Arrow IPC schema message, written at the start of every block. Only used when - /// the schema has no dictionary types. - schema_message: Vec, - /// Whether the schema contains any dictionary types (see [`Self::encode_ipc_stream`]). - has_dictionaries: bool, + /// Pre-encoded Arrow IPC schema message, written verbatim at the start of every block. + /// + /// `None` indicates the schema contains dictionary types, whose dictionary-id bookkeeping ties + /// schema and batch encoding together, so the schema cannot be reused across blocks and + /// [`Self::encode_ipc_stream`] falls back to `StreamWriter`. + schema_message: Option>, } impl ShuffleBlockWriter { pub fn try_new(schema: &Schema, codec: CompressionCodec) -> Result { + // 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 (filled in per block by write_batch) @@ -97,42 +100,48 @@ impl ShuffleBlockWriter { }; header_bytes.extend_from_slice(codec_header); - // Pre-encode the IPC schema message once so it does not have to be re-serialized per block. - let options = IpcWriteOptions::default(); - 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, - &options, - ); - let mut schema_message = Vec::new(); - write_message(&mut schema_message, encoded_schema, &options)?; - let has_dictionaries = schema .fields() .iter() .any(|f| contains_dictionary(f.data_type())); + // 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 and + // leave this `None`. + let schema_message = if has_dictionaries { + None + } else { + let options = IpcWriteOptions::default(); + 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, + &options, + ); + let mut buf = Vec::new(); + write_message(&mut buf, encoded_schema, &options)?; + Some(buf) + }; + Ok(Self { codec, header_bytes, schema: Arc::new(schema.clone()), schema_message, - has_dictionaries, }) } /// Serialize `batch` as a standalone Arrow IPC stream into `out`. fn encode_ipc_stream(&self, batch: &RecordBatch, out: &mut W) -> Result<()> { - if self.has_dictionaries { + let Some(schema_message) = &self.schema_message else { // 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(out, &self.schema)?; stream_writer.write(batch)?; stream_writer.finish()?; return Ok(()); - } + }; // Fast path: reuse the pre-encoded schema message and write the record batch manually. let options = IpcWriteOptions::default(); @@ -147,7 +156,7 @@ impl ShuffleBlockWriter { )?; debug_assert!(encoded_dictionaries.is_empty()); - out.write_all(&self.schema_message)?; + out.write_all(schema_message)?; write_message(&mut *out, encoded_batch, &options)?; out.write_all(&IPC_EOS)?; Ok(()) From ca61876bb3dbc0d0ff747f8b047f61a2fab76bbc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 23 Jul 2026 14:51:56 -0600 Subject: [PATCH 4/5] 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) --- .../src/writers/shuffle_block_writer.rs | 95 ++++++++++--------- 1 file changed, 50 insertions(+), 45 deletions(-) diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs b/native/shuffle/src/writers/shuffle_block_writer.rs index 7ea407ee60..a6d48e1825 100644 --- a/native/shuffle/src/writers/shuffle_block_writer.rs +++ b/native/shuffle/src/writers/shuffle_block_writer.rs @@ -21,6 +21,7 @@ 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; @@ -28,7 +29,10 @@ 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, matching what `StreamWriter::finish` emits for metadata version V5. +/// 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. @@ -40,21 +44,17 @@ pub enum CompressionCodec { Snappy, } -/// Returns true if `data_type` is, or nests, a dictionary type. -fn contains_dictionary(data_type: &DataType) -> bool { - match data_type { - DataType::Dictionary(_, _) => true, - DataType::List(f) - | DataType::LargeList(f) - | DataType::FixedSizeList(f, _) - | DataType::Map(f, _) - | DataType::RunEndEncoded(_, f) => contains_dictionary(f.data_type()), - DataType::Struct(fields) => fields.iter().any(|f| contains_dictionary(f.data_type())), - DataType::Union(fields, _) => fields - .iter() - .any(|(_, f)| contains_dictionary(f.data_type())), - _ => false, - } +/// 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. +#[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(Vec), + /// 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. @@ -69,13 +69,10 @@ fn contains_dictionary(data_type: &DataType) -> bool { pub struct ShuffleBlockWriter { codec: CompressionCodec, header_bytes: Vec, - schema: SchemaRef, - /// Pre-encoded Arrow IPC schema message, written verbatim at the start of every block. - /// - /// `None` indicates the schema contains dictionary types, whose dictionary-id bookkeeping ties - /// schema and batch encoding together, so the schema cannot be reused across blocks and - /// [`Self::encode_ipc_stream`] falls back to `StreamWriter`. - schema_message: Option>, + /// 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 { @@ -100,64 +97,72 @@ impl ShuffleBlockWriter { }; header_bytes.extend_from_slice(codec_header); + // 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 - .fields() + .flattened_fields() .iter() - .any(|f| contains_dictionary(f.data_type())); + .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 and - // leave this `None`. - let schema_message = if has_dictionaries { - None + // 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 options = IpcWriteOptions::default(); 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, - &options, + &write_options, ); let mut buf = Vec::new(); - write_message(&mut buf, encoded_schema, &options)?; - Some(buf) + write_message(&mut buf, encoded_schema, &write_options)?; + SchemaEncoding::Precoded(buf) }; Ok(Self { codec, header_bytes, - schema: Arc::new(schema.clone()), - schema_message, + write_options, + schema_encoding, }) } /// Serialize `batch` as a standalone Arrow IPC stream into `out`. fn encode_ipc_stream(&self, batch: &RecordBatch, out: &mut W) -> Result<()> { - let Some(schema_message) = &self.schema_message else { - // 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(out, &self.schema)?; - stream_writer.write(batch)?; - stream_writer.finish()?; - return Ok(()); + 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 options = IpcWriteOptions::default(); 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, - &options, + &self.write_options, &mut compression_context, )?; debug_assert!(encoded_dictionaries.is_empty()); out.write_all(schema_message)?; - write_message(&mut *out, encoded_batch, &options)?; + write_message(&mut *out, encoded_batch, &self.write_options)?; out.write_all(&IPC_EOS)?; Ok(()) } From 0a6ef3c94d43279a2093e5aa3010a9109e6a8d1a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 06:20:19 -0600 Subject: [PATCH 5/5] 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. --- .../src/writers/shuffle_block_writer.rs | 49 +++++++++++++++++-- 1 file changed, 45 insertions(+), 4 deletions(-) diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs b/native/shuffle/src/writers/shuffle_block_writer.rs index a6d48e1825..7b6846b3ba 100644 --- a/native/shuffle/src/writers/shuffle_block_writer.rs +++ b/native/shuffle/src/writers/shuffle_block_writer.rs @@ -47,11 +47,16 @@ pub enum CompressionCodec { /// 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(Vec), + 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), @@ -68,7 +73,9 @@ enum SchemaEncoding { #[derive(Clone)] pub struct ShuffleBlockWriter { codec: CompressionCodec, - header_bytes: Vec, + /// 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, @@ -123,12 +130,12 @@ impl ShuffleBlockWriter { ); let mut buf = Vec::new(); write_message(&mut buf, encoded_schema, &write_options)?; - SchemaEncoding::Precoded(buf) + SchemaEncoding::Precoded(Arc::from(buf)) }; Ok(Self { codec, - header_bytes, + header_bytes: Arc::from(header_bytes), write_options, schema_encoding, }) @@ -230,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"), + } + } +}