diff --git a/native/shuffle/src/ipc.rs b/native/shuffle/src/ipc.rs index 81ee41332a..d0af69eef5 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -19,34 +19,47 @@ use arrow::array::RecordBatch; use arrow::ipc::reader::StreamReader; use datafusion::common::DataFusionError; use datafusion::error::Result; +use std::cell::RefCell; +use std::io::Read; pub fn read_ipc_compressed(bytes: &[u8]) -> Result { match &bytes[0..4] { - b"SNAP" => { - let decoder = snap::read::FrameDecoder::new(&bytes[4..]); - let mut reader = - unsafe { StreamReader::try_new(decoder, None)?.with_skip_validation(true) }; - reader.next().unwrap().map_err(|e| e.into()) - } - b"LZ4_" => { - let decoder = lz4_flex::frame::FrameDecoder::new(&bytes[4..]); - let mut reader = - unsafe { StreamReader::try_new(decoder, None)?.with_skip_validation(true) }; - reader.next().unwrap().map_err(|e| e.into()) - } - b"ZSTD" => { - let decoder = zstd::Decoder::new(&bytes[4..])?; - let mut reader = - unsafe { StreamReader::try_new(decoder, None)?.with_skip_validation(true) }; - reader.next().unwrap().map_err(|e| e.into()) - } - b"NONE" => { - let mut reader = - unsafe { StreamReader::try_new(&bytes[4..], None)?.with_skip_validation(true) }; - reader.next().unwrap().map_err(|e| e.into()) - } + b"SNAP" => decode_ipc_stream(snap::read::FrameDecoder::new(&bytes[4..])), + b"LZ4_" => decode_ipc_stream(lz4_flex::frame::FrameDecoder::new(&bytes[4..])), + b"ZSTD" => ZSTD_DECODER.with(|decoder| decoder.borrow_mut().read_batch(&bytes[4..])), + b"NONE" => decode_ipc_stream(&bytes[4..]), other => Err(DataFusionError::Execution(format!( "Failed to decode batch: invalid compression codec: {other:?}" ))), } } + +/// Read the single record batch from an uncompressed Arrow IPC stream. +fn decode_ipc_stream(reader: R) -> Result { + let mut reader = unsafe { StreamReader::try_new(reader, None)?.with_skip_validation(true) }; + reader.next().unwrap().map_err(|e| e.into()) +} + +/// Per-thread reusable zstd decompression context, so decoding does not allocate a fresh +/// `ZSTD_DCtx` for every block. +struct ZstdBlockDecoder { + context: zstd::zstd_safe::DCtx<'static>, +} + +impl ZstdBlockDecoder { + fn read_batch(&mut self, frame: &[u8]) -> Result { + // The reader may stop before draining the whole frame, so reset the session state to leave + // the reused context clean for the next block (as a fresh context would be). + self.context + .reset(zstd::zstd_safe::ResetDirective::SessionOnly) + .map_err(|_| DataFusionError::Execution("failed to reset zstd context".to_string()))?; + let decoder = zstd::stream::read::Decoder::with_context(frame, &mut self.context); + decode_ipc_stream(decoder) + } +} + +thread_local! { + static ZSTD_DECODER: RefCell = RefCell::new(ZstdBlockDecoder { + context: zstd::zstd_safe::DCtx::create(), + }); +} diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 8d1336dc1c..116d5017e8 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -313,6 +313,80 @@ mod test { } } + /// Writing and reading several distinct zstd blocks on one thread exercises reuse of the + /// thread-local zstd compression/decompression contexts across frames: a context left dirty by + /// one block would corrupt the next. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + fn roundtrip_ipc_multiple_zstd_blocks() { + use arrow::array::Int32Array; + + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); + let writer = + ShuffleBlockWriter::try_new(schema.as_ref(), CompressionCodec::Zstd(1)).unwrap(); + + // Distinct batches (varying length and values) written as consecutive blocks. + let batches: Vec = (0..5) + .map(|b| { + let values: Vec = (0..(100 + b * 37)).map(|i| i + b * 1000).collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values)) as Arc], + ) + .unwrap() + }) + .collect(); + + let mut output = vec![]; + let mut cursor = Cursor::new(&mut output); + for batch in &batches { + writer + .write_batch(batch, &mut cursor, &Time::default()) + .unwrap(); + } + + let decoded = read_all_ipc_batches(&output); + assert_eq!(decoded, batches); + } + + /// The pre-encoded schema message and the per-block record batch are encoded with separate + /// dictionary trackers, so a dictionary-typed column must still roundtrip: the dictionary + /// batch's id has to match the id assigned while encoding the schema. + #[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..1525fa5029 100644 --- a/native/shuffle/src/writers/shuffle_block_writer.rs +++ b/native/shuffle/src/writers/shuffle_block_writer.rs @@ -16,12 +16,21 @@ // 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::cell::RefCell; +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 +41,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. 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)?; + header_bytes.extend_from_slice(codec_header); - let header_bytes = cursor.into_inner(); + // 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 +172,35 @@ 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) => { + // Reuse the zstd compression context (ZSTD_CCtx + workspace) across blocks on this + // thread instead of allocating a fresh one per block. The stream is compressed + // incrementally (not buffered whole) so large blocks are not copied an extra time. + ZSTD_ENCODER.with(|encoder| { + encoder + .borrow_mut() + .compress_stream(*level, output, |w| self.encode_ipc_stream(batch, w)) + })?; + } + } // fill ipc length let end_pos = output.stream_position()?; @@ -134,7 +213,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 +222,47 @@ impl ShuffleBlockWriter { Ok((end_pos - start_pos) as usize) } } + +/// Per-thread reusable zstd compression context, so the streaming shuffle writer does not allocate +/// a fresh `ZSTD_CCtx` (and workspace) for every block. +struct ZstdBlockEncoder { + context: zstd::zstd_safe::CCtx<'static>, + /// The compression level currently configured on the context; the level is (re)applied only + /// when it changes. + level: Option, +} + +impl ZstdBlockEncoder { + /// Compress the IPC stream produced by `build` into `output`, reusing this thread's + /// `ZSTD_CCtx`. Compression is streamed incrementally, so the uncompressed stream is never + /// fully buffered. + fn compress_stream( + &mut self, + level: i32, + output: &mut W, + build: impl FnOnce(&mut dyn Write) -> Result<()>, + ) -> Result<()> { + if self.level != Some(level) { + self.context + .set_parameter(zstd::zstd_safe::CParameter::CompressionLevel(level)) + .map_err(|_| { + DataFusionError::Execution(format!("failed to set zstd level {level}")) + })?; + self.level = Some(level); + } + + let mut encoder = zstd::stream::write::Encoder::with_context(output, &mut self.context); + build(&mut encoder)?; + encoder + .finish() + .map_err(|e| DataFusionError::Execution(format!("zstd compression error: {e}")))?; + Ok(()) + } +} + +thread_local! { + static ZSTD_ENCODER: RefCell = RefCell::new(ZstdBlockEncoder { + context: zstd::zstd_safe::CCtx::create(), + level: None, + }); +}