From 8da6e6b445d8801b9e915c59c778cb7034851a70 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Mon, 31 Aug 2026 09:57:30 +0700 Subject: [PATCH 1/4] perf: reuse zstd compression contexts across shuffle blocks Every shuffle block previously created and destroyed its own zstd context: a fresh CCtx per encoded block and a fresh DCtx per decoded frame. Context setup is pure overhead that scales with block count, so high-partition shuffles with small blocks pay the most. Encode paths now share one context per task, threaded from the task-level owner so codec memory stays bounded regardless of partition count. The remote shuffle path still frees the zstd workspace with each admitted encode, keeping its memory accounting accurate. Decode reuses a per-thread context behind the existing entry points. Wire format is unchanged. On a 4M-row hash shuffle with 10,000 partitions at zstd level 3, encode time drops ~10% and wall time ~7%; larger-block shapes are within noise. lz4 and snappy keep per-block encoders: no reset API in the pinned crates and much smaller setup cost. Part of #5002. --- .../src/execution/operators/shuffle_scan.rs | 9 +- native/shuffle/benches/shuffle_writer.rs | 11 +- native/shuffle/src/codec_context.rs | 102 +++++++ native/shuffle/src/ipc.rs | 77 ++++- native/shuffle/src/lib.rs | 7 +- native/shuffle/src/shuffle_writer.rs | 47 +-- native/shuffle/src/spark_unsafe/row.rs | 9 +- .../shuffle/src/writers/buf_batch_writer.rs | 17 +- .../writers/local/local_partition_writer.rs | 40 ++- native/shuffle/src/writers/local/spill.rs | 21 +- native/shuffle/src/writers/rss/mod.rs | 12 +- .../src/writers/rss/rss_partition_writer.rs | 13 +- .../src/writers/shuffle_block_writer.rs | 269 ++++++++++++++++-- 13 files changed, 533 insertions(+), 101 deletions(-) create mode 100644 native/shuffle/src/codec_context.rs diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs index e5014aa639e..0be1b512769 100644 --- a/native/core/src/execution/operators/shuffle_scan.rs +++ b/native/core/src/execution/operators/shuffle_scan.rs @@ -398,10 +398,9 @@ impl RecordBatchStream for ShuffleScanStream { #[cfg(test)] mod tests { - use crate::execution::shuffle::{CompressionCodec, ShuffleBlockWriter}; + use crate::execution::shuffle::{CompressionCodec, ShuffleBlockWriter, ShuffleCodecContext}; use arrow::array::{Int32Array, RecordBatchOptions, StringArray, UInt32Array}; 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; @@ -416,7 +415,7 @@ mod tests { .write_batch( batch, &mut output, - &mut CompressionContext::default(), + &mut ShuffleCodecContext::default(), &Time::new(), ) .unwrap(); @@ -535,7 +534,7 @@ mod tests { .write_batch( &batch, &mut buf, - &mut CompressionContext::default(), + &mut ShuffleCodecContext::default(), &ipc_time, ) .unwrap(); @@ -607,7 +606,7 @@ mod tests { .write_batch( &dict_batch, &mut buf, - &mut CompressionContext::default(), + &mut ShuffleCodecContext::default(), &ipc_time, ) .unwrap(); diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index 2e82f6fc532..b490f92a349 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -18,7 +18,6 @@ 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; @@ -31,7 +30,7 @@ use datafusion::{ prelude::SessionContext, }; use datafusion_comet_shuffle::{ - CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShuffleWriterExec, + CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShuffleCodecContext, ShuffleWriterExec, }; use itertools::Itertools; use std::io::Cursor; @@ -54,11 +53,11 @@ 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(); + let mut codec_context = ShuffleCodecContext::default(); b.iter(|| { buffer.clear(); let mut cursor = Cursor::new(&mut buffer); - w.write_batch(&batch, &mut cursor, &mut compression_context, &ipc_time) + w.write_batch(&batch, &mut cursor, &mut codec_context, &ipc_time) .unwrap(); }); }); @@ -256,14 +255,14 @@ 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(); + let mut codec_context = ShuffleCodecContext::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, &mut compression_context, &ipc_time) + .write_batch(&batch, &mut cursor, &mut codec_context, &ipc_time) .unwrap(); }); }); diff --git a/native/shuffle/src/codec_context.rs b/native/shuffle/src/codec_context.rs new file mode 100644 index 00000000000..052358ff76b --- /dev/null +++ b/native/shuffle/src/codec_context.rs @@ -0,0 +1,102 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::ipc::writer::CompressionContext; +use std::io; +use zstd::zstd_safe::{CCtx, CParameter, DCtx, ResetDirective}; + +/// Reusable compression state for encoding shuffle blocks. +/// +/// A zstd context costs about a megabyte and real setup time, so a task shares one across all +/// the blocks it encodes instead of paying per block. Keep ownership task-scoped, never +/// per-output-partition -- a shuffle can have thousands of partitions. Local shuffle holds the +/// zstd context for the whole task; the remote (RSS) path frees it after each admitted encode +/// via [`Self::release_zstd`], since its memory accounting only reserves the workspace per +/// invocation. +#[derive(Default)] +pub struct ShuffleCodecContext { + /// Arrow's per-message IPC compression scratch, reused across encodes. + pub(crate) arrow_ipc: CompressionContext, + /// Lazily created, reused across blocks. + zstd: Option>, +} + +impl ShuffleCodecContext { + /// The shared zstd context primed for one frame at `level`, plus the Arrow IPC scratch + /// (returned together because the encoder borrows the context for the whole frame). + /// + /// The session reset and level re-apply happen on every call: writers with different + /// levels can share one context, and a failed encode must not leave state behind. + pub(crate) fn zstd_cctx( + &mut self, + level: i32, + ) -> io::Result<(&mut CCtx<'static>, &mut CompressionContext)> { + let cctx = + match &mut self.zstd { + Some(cctx) => cctx, + none => none.insert(CCtx::try_create().ok_or_else(|| { + io::Error::other("failed to allocate zstd compression context") + })?), + }; + cctx.reset(ResetDirective::SessionOnly) + .map_err(map_zstd_error)?; + cctx.set_parameter(CParameter::CompressionLevel(level)) + .map_err(map_zstd_error)?; + Ok((cctx, &mut self.arrow_ipc)) + } + + /// Drops the cached zstd context, freeing its native workspace. The remote encode path + /// calls this after every admitted encode so the memory lives and dies inside that + /// invocation's reservation; the next zstd encode re-creates it lazily. + pub(crate) fn release_zstd(&mut self) { + self.zstd = None; + } + + /// Test hook for the release-vs-retain contract of the two encode paths. + #[cfg(test)] + pub(crate) fn holds_zstd_cctx(&self) -> bool { + self.zstd.is_some() + } +} + +/// Decode-side counterpart of [`ShuffleCodecContext`]: one context serves every frame a +/// reader decodes instead of allocating a fresh zstd context per frame. +#[derive(Default)] +pub struct ShuffleDecodeContext { + /// Lazily created, reused across frames. + zstd: Option>, +} + +impl ShuffleDecodeContext { + /// The shared zstd decompression context primed for one frame. The session reset clears + /// anything a failed decode (e.g. a truncated fetch) left mid-frame. + pub(crate) fn zstd_dctx(&mut self) -> io::Result<&mut DCtx<'static>> { + let dctx = match &mut self.zstd { + Some(dctx) => dctx, + none => none.insert(DCtx::try_create().ok_or_else(|| { + io::Error::other("failed to allocate zstd decompression context") + })?), + }; + dctx.reset(ResetDirective::SessionOnly) + .map_err(map_zstd_error)?; + Ok(dctx) + } +} + +fn map_zstd_error(code: usize) -> io::Error { + io::Error::other(zstd::zstd_safe::get_error_name(code)) +} diff --git a/native/shuffle/src/ipc.rs b/native/shuffle/src/ipc.rs index 97890f50148..2ccb04d2965 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -15,23 +15,52 @@ // specific language governing permissions and limitations // under the License. +use crate::codec_context::ShuffleDecodeContext; use arrow::array::RecordBatch; use arrow::ipc::reader::StreamReader; use datafusion::common::DataFusionError; use datafusion::error::Result; +use std::cell::RefCell; use std::io::{Error, ErrorKind, Read}; +thread_local! { + /// Backs the entry points below. They're called from many JVM task threads; a + /// thread-local gets each thread context reuse without changing any caller. + static DECODE_CONTEXT: RefCell = + RefCell::new(ShuffleDecodeContext::default()); +} + /// Decode trusted local Comet output without revalidating every Arrow array value or offset. pub fn read_ipc_compressed(bytes: &[u8]) -> Result { - read_ipc_compressed_impl(bytes, false) + DECODE_CONTEXT.with(|context| read_ipc_compressed_impl(&mut context.borrow_mut(), bytes, false)) } /// Decode remotely fetched Comet output, including Arrow buffer and offset validation. pub fn read_ipc_compressed_validated(bytes: &[u8]) -> Result { - read_ipc_compressed_impl(bytes, true) + DECODE_CONTEXT.with(|context| read_ipc_compressed_impl(&mut context.borrow_mut(), bytes, true)) +} + +/// [`read_ipc_compressed`] with a caller-owned decode context. +pub fn read_ipc_compressed_with( + decode_context: &mut ShuffleDecodeContext, + bytes: &[u8], +) -> Result { + read_ipc_compressed_impl(decode_context, bytes, false) +} + +/// [`read_ipc_compressed_validated`] with a caller-owned decode context. +pub fn read_ipc_compressed_validated_with( + decode_context: &mut ShuffleDecodeContext, + bytes: &[u8], +) -> Result { + read_ipc_compressed_impl(decode_context, bytes, true) } -fn read_ipc_compressed_impl(bytes: &[u8], validate: bool) -> Result { +fn read_ipc_compressed_impl( + decode_context: &mut ShuffleDecodeContext, + bytes: &[u8], + validate: bool, +) -> Result { let codec = bytes.get(..4).ok_or_else(|| { DataFusionError::Execution("Failed to decode batch: truncated compression codec".to_owned()) })?; @@ -44,7 +73,10 @@ fn read_ipc_compressed_impl(bytes: &[u8], validate: bool) -> Result )?, // The slice already implements BufRead. Adding another BufReader would let read-ahead // conceal compressed bytes left over after the decoder reaches its end marker. - b"ZSTD" => read_single_batch(zstd::Decoder::with_buffer(&mut encoded)?, validate)?, + b"ZSTD" => read_single_batch( + zstd::Decoder::with_context(&mut encoded, decode_context.zstd_dctx()?), + validate, + )?, b"NONE" => read_single_batch(&mut encoded, validate)?, other => { return Err(DataFusionError::Execution(format!( @@ -114,7 +146,11 @@ fn read_single_batch(input: R, validate: bool) -> Result { #[cfg(test)] mod tests { - use super::{read_ipc_compressed, read_ipc_compressed_validated}; + use super::{ + read_ipc_compressed, read_ipc_compressed_validated, read_ipc_compressed_validated_with, + read_ipc_compressed_with, + }; + use crate::codec_context::ShuffleDecodeContext; use arrow::array::{Int32Array, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::ipc::writer::StreamWriter; @@ -267,4 +303,35 @@ mod tests { assert_eq!(batch, validated); } } + + /// One context across many frames, codecs changing between them, must decode exactly + /// like fresh per-frame decoders. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn decode_context_reused_across_frames_and_codecs() { + let mut ctx = ShuffleDecodeContext::default(); + for _ in 0..3 { + for codec in [b"ZSTD", b"NONE", b"SNAP", b"ZSTD", b"LZ4_", b"ZSTD"] { + let frame = encode(codec, &ipc_stream(1)); + let fresh = read_ipc_compressed(&frame).unwrap(); + let reused = read_ipc_compressed_with(&mut ctx, &frame).unwrap(); + let reused_validated = + read_ipc_compressed_validated_with(&mut ctx, &frame).unwrap(); + assert_eq!(reused, fresh); + assert_eq!(reused_validated, fresh); + } + } + } + + /// A truncated frame must not poison the context for the next valid one. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn decode_context_usable_after_error() { + let mut ctx = ShuffleDecodeContext::default(); + let good = encode(b"ZSTD", &ipc_stream(1)); + let truncated = &good[..good.len() - 7]; + assert!(read_ipc_compressed_with(&mut ctx, truncated).is_err()); + let batch = read_ipc_compressed_with(&mut ctx, &good).unwrap(); + assert_eq!(batch.num_rows(), 3); + } } diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 3adec603aef..15d61bfd65a 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod codec_context; pub(crate) mod comet_partitioning; pub mod ipc; pub(crate) mod metrics; @@ -28,8 +29,12 @@ mod spark_crc32c_hasher; pub mod spark_unsafe; pub(crate) mod writers; +pub use codec_context::{ShuffleCodecContext, ShuffleDecodeContext}; pub use comet_partitioning::CometPartitioning; -pub use ipc::{read_ipc_compressed, read_ipc_compressed_validated}; +pub use ipc::{ + read_ipc_compressed, read_ipc_compressed_validated, read_ipc_compressed_validated_with, + read_ipc_compressed_with, +}; pub use remote_schema::validate_remote_schema; pub use schema_align::SchemaAlignExec; pub use shuffle_writer::{ShuffleWriterDestination, ShuffleWriterExec}; diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 21d71dc1221..a8f138d0bc3 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -392,10 +392,9 @@ fn contextualize_shuffle_error(error: DataFusionError, phase: &str) -> DataFusio #[cfg(test)] mod test { use super::*; - use crate::{read_ipc_compressed, ShuffleBlockWriter}; + use crate::{read_ipc_compressed, ShuffleBlockWriter, ShuffleCodecContext}; 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; @@ -425,14 +424,9 @@ 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 mut codec_context = ShuffleCodecContext::default(); let length = writer - .write_batch( - &batch, - &mut cursor, - &mut compression_context, - &Time::default(), - ) + .write_batch(&batch, &mut cursor, &mut codec_context, &Time::default()) .unwrap(); assert_eq!(length, output.len()); @@ -471,14 +465,9 @@ 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(); + let mut codec_context = ShuffleCodecContext::default(); writer - .write_batch( - &batch, - &mut cursor, - &mut compression_context, - &Time::default(), - ) + .write_batch(&batch, &mut cursor, &mut codec_context, &Time::default()) .unwrap(); let batch2 = read_ipc_compressed(&output[16..]).unwrap(); @@ -1160,6 +1149,7 @@ mod test { let codec = CompressionCodec::Lz4Frame; let encode_time = Time::default(); let write_time = Time::default(); + let mut codec_context = ShuffleCodecContext::default(); // Write with coalescing (batch_size=8192) let mut coalesced_output = Vec::new(); @@ -1172,9 +1162,13 @@ mod test { 8192, ); for batch in &small_batches { - buf_writer.write(batch, &encode_time, &write_time).unwrap(); + buf_writer + .write(batch, &mut codec_context, &encode_time, &write_time) + .unwrap(); } - buf_writer.flush(&encode_time, &write_time).unwrap(); + buf_writer + .flush(&mut codec_context, &encode_time, &write_time) + .unwrap(); } // Write without coalescing (batch_size=1) @@ -1188,9 +1182,13 @@ mod test { 1, ); for batch in &small_batches { - buf_writer.write(batch, &encode_time, &write_time).unwrap(); + buf_writer + .write(batch, &mut codec_context, &encode_time, &write_time) + .unwrap(); } - buf_writer.flush(&encode_time, &write_time).unwrap(); + buf_writer + .flush(&mut codec_context, &encode_time, &write_time) + .unwrap(); } // Coalesced output should be smaller due to fewer IPC schema blocks @@ -1281,6 +1279,7 @@ mod test { let codec = CompressionCodec::Lz4Frame; let encode_time = Time::default(); let write_time = Time::default(); + let mut codec_context = ShuffleCodecContext::default(); let mut output = Vec::new(); { @@ -1292,9 +1291,13 @@ mod test { batch_size as usize, ); for batch in &inputs { - buf_writer.write(batch, &encode_time, &write_time).unwrap(); + buf_writer + .write(batch, &mut codec_context, &encode_time, &write_time) + .unwrap(); } - buf_writer.flush(&encode_time, &write_time).unwrap(); + buf_writer + .flush(&mut codec_context, &encode_time, &write_time) + .unwrap(); } let blocks = read_all_ipc_batches(&output); diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs index 77e44024fbf..1918ce3b18c 100644 --- a/native/shuffle/src/spark_unsafe/row.rs +++ b/native/shuffle/src/spark_unsafe/row.rs @@ -17,6 +17,7 @@ //! Utils for supporting native sort-based columnar shuffle. +use crate::codec_context::ShuffleCodecContext; use crate::spark_unsafe::unsafe_object::{impl_primitive_accessors, SparkUnsafeObject}; use crate::spark_unsafe::{ list::append_list_element, @@ -38,7 +39,6 @@ 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}; @@ -1388,7 +1388,9 @@ 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(); + // One context for every batch this call encodes; the JVM calls in once per sorted + // partition, so there is no wider native scope to hoist it to. + let mut codec_context = ShuffleCodecContext::default(); while current_row < row_num { let n = std::cmp::min(batch_size, row_num - current_row); @@ -1422,8 +1424,7 @@ 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, &mut compression_context, &ipc_time)?; + written += block_writer.write_batch(&batch, &mut cursor, &mut codec_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 55d88a4ba48..74c60747f04 100644 --- a/native/shuffle/src/writers/buf_batch_writer.rs +++ b/native/shuffle/src/writers/buf_batch_writer.rs @@ -16,9 +16,9 @@ // under the License. use super::ShuffleBlockWriter; +use crate::codec_context::ShuffleCodecContext; 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}; @@ -33,12 +33,14 @@ use std::io::{Cursor, Seek, SeekFrom, Write}; /// 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. +/// +/// Encoding methods borrow a [`ShuffleCodecContext`] rather than owning one: these writers +/// are created per output partition, and codec contexts must stay task-scoped. pub(crate) struct BufBatchWriter, W: Write> { shuffle_block_writer: S, writer: W, buffer: Vec, 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, @@ -58,7 +60,6 @@ impl, W: Write> BufBatchWriter { writer, buffer: vec![], buffer_max_size, - compression_context: CompressionContext::default(), coalescer: None, batch_size, } @@ -67,6 +68,7 @@ impl, W: Write> BufBatchWriter { pub(crate) fn write( &mut self, batch: &RecordBatch, + codec_context: &mut ShuffleCodecContext, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result { @@ -96,7 +98,8 @@ impl, W: Write> BufBatchWriter { let mut bytes_written = 0; for batch in &completed { - bytes_written += self.write_batch_to_buffer(batch, encode_time, write_time)?; + bytes_written += + self.write_batch_to_buffer(batch, codec_context, encode_time, write_time)?; } Ok(bytes_written) } @@ -105,6 +108,7 @@ impl, W: Write> BufBatchWriter { fn write_batch_to_buffer( &mut self, batch: &RecordBatch, + codec_context: &mut ShuffleCodecContext, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result { @@ -113,7 +117,7 @@ impl, W: Write> BufBatchWriter { let bytes_written = self.shuffle_block_writer.borrow().write_batch( batch, &mut cursor, - &mut self.compression_context, + codec_context, encode_time, )?; let pos = cursor.position(); @@ -128,6 +132,7 @@ impl, W: Write> BufBatchWriter { pub(crate) fn flush( &mut self, + codec_context: &mut ShuffleCodecContext, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result<()> { @@ -140,7 +145,7 @@ impl, W: Write> BufBatchWriter { } } for batch in &remaining { - self.write_batch_to_buffer(batch, encode_time, write_time)?; + self.write_batch_to_buffer(batch, codec_context, encode_time, write_time)?; } // Flush the byte buffer to the underlying writer diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index 3a9a6484dba..a75626a2fcc 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::codec_context::ShuffleCodecContext; use crate::metrics::ShufflePartitionerMetrics; use crate::writers::local::spill::SpillWriter; use crate::writers::partition_writer::PartitionWriter; @@ -64,6 +65,9 @@ enum DataOutput { pub(crate) struct LocalPartitionWriter { output_index_file: String, data_output: DataOutput, + /// Compression state shared by every block this task writes; the per-partition + /// `BufBatchWriter`s borrow it (see [`ShuffleCodecContext`]). + codec_context: ShuffleCodecContext, /// Start offset of each partition in the data file, plus a trailing entry /// with the total length so partition sizes are simple offset differences. /// Has `num_output_partitions + 1` elements. @@ -121,6 +125,7 @@ impl LocalPartitionWriter { Ok(Self { output_index_file, data_output, + codec_context: ShuffleCodecContext::default(), offsets: vec![0u64; num_output_partitions + 1], batch_size, write_buffer_size, @@ -164,7 +169,12 @@ impl PartitionWriter for LocalPartitionWriter { // `finish_all`. for batch in iter.by_ref() { let batch = batch?; - writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; + writer.write( + &batch, + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; } } DataOutput::Multi { @@ -175,7 +185,7 @@ impl PartitionWriter for LocalPartitionWriter { // Multi-partition output buffers each partition's batches into its own // spill file. `finish_partition` later merges the spill files (and any // remaining in-memory batches) into the shuffle output in partition order. - spill_writers[pid].write(iter, runtime, metrics)?; + spill_writers[pid].write(iter, &mut self.codec_context, runtime, metrics)?; } } @@ -209,7 +219,12 @@ impl PartitionWriter for LocalPartitionWriter { // flushed once in `finish_all`. for batch in iter.by_ref() { let batch = batch?; - writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; + writer.write( + &batch, + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; } } DataOutput::Multi { @@ -243,9 +258,18 @@ impl PartitionWriter for LocalPartitionWriter { ); for batch in iter.by_ref() { let batch = batch?; - buf_batch_writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; + buf_batch_writer.write( + &batch, + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; } - buf_batch_writer.flush(&metrics.encode_time, &metrics.write_time)?; + buf_batch_writer.flush( + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; } } Ok(()) @@ -259,7 +283,11 @@ impl PartitionWriter for LocalPartitionWriter { // single-partition writer this also finalizes the last coalesced batch. let final_offset = match &mut self.data_output { DataOutput::Single(writer) => { - writer.flush(&metrics.encode_time, &metrics.write_time)?; + writer.flush( + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; writer.writer_stream_position()? } DataOutput::Multi { output_writer, .. } => { diff --git a/native/shuffle/src/writers/local/spill.rs b/native/shuffle/src/writers/local/spill.rs index d8534b44564..6ef89c76800 100644 --- a/native/shuffle/src/writers/local/spill.rs +++ b/native/shuffle/src/writers/local/spill.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::codec_context::ShuffleCodecContext; use crate::metrics::ShufflePartitionerMetrics; use crate::writers::BufBatchWriter; use crate::ShuffleBlockWriter; @@ -50,9 +51,13 @@ impl SpillWriter { }) } + /// Stages the batches from `iter` into this partition's spill file. + /// + /// `codec_context` comes from the task-level owner; a `SpillWriter` exists per partition. pub(crate) fn write>>( &mut self, iter: &mut I, + codec_context: &mut ShuffleCodecContext, runtime: &RuntimeEnv, metrics: &ShufflePartitionerMetrics, ) -> datafusion::common::Result<()> { @@ -67,12 +72,22 @@ impl SpillWriter { self.batch_size, ); let initial_position = buf_batch_writer.writer_stream_position()?; - buf_batch_writer.write(&batch?, &metrics.encode_time, &metrics.write_time)?; + buf_batch_writer.write( + &batch?, + codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; for batch in iter.by_ref() { let batch = batch?; - buf_batch_writer.write(&batch, &metrics.encode_time, &metrics.write_time)?; + buf_batch_writer.write( + &batch, + codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; } - buf_batch_writer.flush(&metrics.encode_time, &metrics.write_time)?; + buf_batch_writer.flush(codec_context, &metrics.encode_time, &metrics.write_time)?; let bytes_written = buf_batch_writer .writer_stream_position()? .saturating_sub(initial_position); diff --git a/native/shuffle/src/writers/rss/mod.rs b/native/shuffle/src/writers/rss/mod.rs index a151b6831fc..1e6f6ab27cc 100644 --- a/native/shuffle/src/writers/rss/mod.rs +++ b/native/shuffle/src/writers/rss/mod.rs @@ -22,14 +22,13 @@ mod tests { use super::rss_partition_writer::RssPartitionWriter; use crate::metrics::ShufflePartitionerMetrics; use crate::writers::PartitionWriter; - use crate::{read_ipc_compressed, CompressionCodec, ShuffleBlockWriter}; + use crate::{read_ipc_compressed, CompressionCodec, ShuffleBlockWriter, ShuffleCodecContext}; use arrow::array::{ Array, ArrayRef, DictionaryArray, Int32Array, ListArray, MapArray, StringArray, StructArray, }; use arrow::buffer::OffsetBuffer; use arrow::compute::cast; use arrow::datatypes::{DataType, Field, Int32Type, Schema}; - use arrow::ipc::writer::CompressionContext; use arrow::record_batch::RecordBatch; use datafusion::common::{DataFusionError, Result}; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, Time}; @@ -400,14 +399,9 @@ mod tests { let block_writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); let mut frame = Cursor::new(Vec::new()); - let mut compression_context = CompressionContext::default(); + let mut codec_context = ShuffleCodecContext::default(); block_writer - .write_batch( - batch, - &mut frame, - &mut compression_context, - &Time::default(), - ) + .write_batch(batch, &mut frame, &mut codec_context, &Time::default()) .unwrap() } diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs index 76503011dbf..1d402509679 100644 --- a/native/shuffle/src/writers/rss/rss_partition_writer.rs +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::codec_context::ShuffleCodecContext; use crate::metrics::ShufflePartitionerMetrics; use crate::writers::partition_writer::PartitionWriter; use crate::ShuffleBlockWriter; @@ -26,7 +27,6 @@ use arrow::array::{ }; use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field, Int16Type, Int32Type, Int64Type}; -use arrow::ipc::writer::CompressionContext; use arrow_select::dictionary::garbage_collect_any_dictionary; use datafusion::common::{DataFusionError, Result}; use datafusion_comet_jni_bridge::ShufflePartitionPusher; @@ -49,7 +49,10 @@ pub(crate) struct RssPartitionWriter { pusher: Arc, num_partitions: usize, max_frame_size: usize, - compression_context: CompressionContext, + /// One remote writer serves all of a task's partitions, so the context is task-scoped by + /// construction. Only the Arrow IPC scratch persists between blocks; `write_rss_batch` + /// frees the zstd workspace with each admitted encode. + codec_context: ShuffleCodecContext, next_partition_to_finish: usize, finished: bool, failed: bool, @@ -90,7 +93,7 @@ impl RssPartitionWriter { pusher, num_partitions, max_frame_size, - compression_context: CompressionContext::default(), + codec_context: ShuffleCodecContext::default(), next_partition_to_finish: 0, finished: false, failed: false, @@ -273,7 +276,7 @@ impl RssPartitionWriter { if let Err(error) = self.block_writer.write_rss_batch( compacted_batch, &mut output, - &mut self.compression_context, + &mut self.codec_context, &metrics.encode_time, ) { let exceeded = output.exceeded; @@ -1178,7 +1181,7 @@ mod buffer_tests { .write_rss_batch( &batch, &mut output, - &mut CompressionContext::default(), + &mut ShuffleCodecContext::default(), &Time::default(), ) .unwrap(); diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs b/native/shuffle/src/writers/shuffle_block_writer.rs index c3987003421..3c4b6ae29ec 100644 --- a/native/shuffle/src/writers/shuffle_block_writer.rs +++ b/native/shuffle/src/writers/shuffle_block_writer.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::codec_context::ShuffleCodecContext; use arrow::array::RecordBatch; use arrow::datatypes::{DataType, Schema, SchemaRef}; use arrow::ipc::writer::{ @@ -89,6 +90,9 @@ impl ShuffleBlockWriter { /// Snappy uses a 64 KiB input block, a 76,490-byte output block, and its hash table. 256 KiB /// conservatively covers either encoder. Zstd's streaming estimate includes the C-allocated /// context, window, and input/output buffers; its Rust writer adds a fixed 32 KiB output Vec. + /// + /// Charged and released per admitted invocation, so the zstd context must live and die + /// within that window (see `write_batch_with_codec_limits`). pub(crate) fn rss_codec_workspace(&self) -> Result { match self.codec { CompressionCodec::None => Ok(0), @@ -233,10 +237,10 @@ impl ShuffleBlockWriter { &self, batch: &RecordBatch, output: &mut W, - compression_context: &mut CompressionContext, + codec_context: &mut ShuffleCodecContext, ipc_time: &Time, ) -> Result { - self.write_batch_with_codec_limits(batch, output, compression_context, ipc_time, false) + self.write_batch_with_codec_limits(batch, output, codec_context, ipc_time, false) } /// Encode with the codec settings covered by [`Self::rss_codec_workspace`]. Local shuffle @@ -245,17 +249,17 @@ impl ShuffleBlockWriter { &self, batch: &RecordBatch, output: &mut W, - compression_context: &mut CompressionContext, + codec_context: &mut ShuffleCodecContext, ipc_time: &Time, ) -> Result { - self.write_batch_with_codec_limits(batch, output, compression_context, ipc_time, true) + self.write_batch_with_codec_limits(batch, output, codec_context, ipc_time, true) } fn write_batch_with_codec_limits( &self, batch: &RecordBatch, output: &mut W, - compression_context: &mut CompressionContext, + codec_context: &mut ShuffleCodecContext, ipc_time: &Time, bounded_rss_codec: bool, ) -> Result { @@ -269,9 +273,48 @@ impl ShuffleBlockWriter { // write header output.write_all(&self.header_bytes)?; + let encode_result = + self.compress_ipc_stream(batch, output, codec_context, bounded_rss_codec); + if bounded_rss_codec { + // RSS charges the zstd workspace (rss_codec_workspace) to each admitted encode + // and releases the charge when it ends, success or not. Free the workspace inside + // that window -- kept alive it would be native memory the reservation system no + // longer tracks. Local shuffle keeps the context for the whole task. + codec_context.release_zstd(); + } + encode_result?; + + // fill ipc length + let end_pos = output.stream_position()?; + let ipc_length = end_pos - start_pos - 8; + let max_size = i32::MAX as u64; + if ipc_length > max_size { + return Err(DataFusionError::Execution(format!( + "Shuffle block size {ipc_length} exceeds maximum size of {max_size}. \ + Try reducing batch size or increasing compression level" + ))); + } + + output.seek(SeekFrom::Start(start_pos))?; + output.write_all(&ipc_length.to_le_bytes())?; + output.seek(SeekFrom::Start(end_pos))?; + + timer.stop(); + + Ok((end_pos - start_pos) as usize) + } + + /// Encode `batch` through the configured outer compression codec into `output`. + fn compress_ipc_stream( + &self, + batch: &RecordBatch, + output: &mut W, + codec_context: &mut ShuffleCodecContext, + bounded_rss_codec: bool, + ) -> Result<()> { match &self.codec { CompressionCodec::None => { - self.encode_ipc_stream(batch, output, compression_context)?; + self.encode_ipc_stream(batch, output, &mut codec_context.arrow_ipc)?; } CompressionCodec::Lz4Frame => { let frame_info = if bounded_rss_codec { @@ -282,50 +325,218 @@ impl ShuffleBlockWriter { }; let mut wtr = lz4_flex::frame::FrameEncoder::with_frame_info(frame_info, &mut *output); - self.encode_ipc_stream(batch, &mut wtr, compression_context)?; + self.encode_ipc_stream(batch, &mut wtr, &mut codec_context.arrow_ipc)?; 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, compression_context)?; + self.encode_ipc_stream(batch, &mut wtr, &mut codec_context.arrow_ipc)?; 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, compression_context)?; + let (cctx, arrow_ipc) = codec_context.zstd_cctx(*level)?; + let mut encoder = zstd::Encoder::with_context(&mut *output, cctx); + self.encode_ipc_stream(batch, &mut encoder, arrow_ipc)?; encoder.finish()?; } } - - // fill ipc length - let end_pos = output.stream_position()?; - let ipc_length = end_pos - start_pos - 8; - let max_size = i32::MAX as u64; - if ipc_length > max_size { - return Err(DataFusionError::Execution(format!( - "Shuffle block size {ipc_length} exceeds maximum size of {max_size}. \ - Try reducing batch size or increasing compression level" - ))); - } - - output.seek(SeekFrom::Start(start_pos))?; - output.write_all(&ipc_length.to_le_bytes())?; - output.seek(SeekFrom::Start(end_pos))?; - - timer.stop(); - - Ok((end_pos - start_pos) as usize) + Ok(()) } } #[cfg(test)] mod tests { use super::*; + use crate::codec_context::ShuffleCodecContext; + use crate::read_ipc_compressed; + use arrow::array::{Int64Array, StringArray}; use arrow::datatypes::{DataType, Field}; + use std::io::Cursor; + + fn test_schema() -> Schema { + Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Utf8, false), + ]) + } + + fn test_batch(seed: i64, rows: usize) -> RecordBatch { + let ints: Vec = (0..rows as i64).map(|i| seed * 1_000_000 + i).collect(); + let strings: Vec = (0..rows).map(|i| format!("row-{seed}-{i}")).collect(); + RecordBatch::try_new( + Arc::new(test_schema()), + vec![ + Arc::new(Int64Array::from(ints)), + Arc::new(StringArray::from(strings)), + ], + ) + .unwrap() + } + + /// One long-lived context, a new writer per partition, several blocks per writer -- the + /// same shape as the local finish/spill loops. Every block must decode on its own. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + fn codec_context_reused_across_blocks_and_writers_roundtrips() { + for codec in &[ + CompressionCodec::None, + CompressionCodec::Zstd(1), + CompressionCodec::Snappy, + CompressionCodec::Lz4Frame, + ] { + let mut ctx = ShuffleCodecContext::default(); + let mut blocks: Vec<(RecordBatch, Vec)> = vec![]; + for partition in 0..3i64 { + let writer = ShuffleBlockWriter::try_new(&test_schema(), codec.clone()).unwrap(); + for block in 0..4i64 { + let batch = test_batch(partition * 10 + block, 100); + let mut out = vec![]; + let mut cursor = Cursor::new(&mut out); + writer + .write_batch(&batch, &mut cursor, &mut ctx, &Time::default()) + .unwrap(); + blocks.push((batch, out)); + } + } + for (expected, bytes) in &blocks { + let decoded = read_ipc_compressed(&bytes[16..]).unwrap(); + assert_eq!(&decoded, expected); + } + } + } + + /// Writers with different zstd levels share one context; neither level may stick to the + /// other's blocks. On repetitive data level 19 must compress smaller than level 1 even + /// through the shared context. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + fn codec_context_serves_alternating_zstd_levels() { + let batch = { + let ints: Vec = (0..4096).map(|i| i % 4).collect(); + let strings: Vec = (0..4096).map(|i| format!("padding-{}", i % 8)).collect(); + RecordBatch::try_new( + Arc::new(test_schema()), + vec![ + Arc::new(Int64Array::from(ints)), + Arc::new(StringArray::from(strings)), + ], + ) + .unwrap() + }; + let fast = ShuffleBlockWriter::try_new(&test_schema(), CompressionCodec::Zstd(1)).unwrap(); + let slow = ShuffleBlockWriter::try_new(&test_schema(), CompressionCodec::Zstd(19)).unwrap(); + let mut ctx = ShuffleCodecContext::default(); + let mut sizes = vec![]; + // Interleave so each block re-encounters the other writer's level on the shared context. + for _ in 0..2 { + for writer in [&fast, &slow] { + let mut out = vec![]; + let mut cursor = Cursor::new(&mut out); + writer + .write_batch(&batch, &mut cursor, &mut ctx, &Time::default()) + .unwrap(); + assert_eq!(read_ipc_compressed(&out[16..]).unwrap(), batch); + sizes.push(out.len()); + } + } + // sizes = [fast, slow, fast, slow]; each writer's level must hold on every block. + assert!( + sizes[1] < sizes[0] && sizes[3] < sizes[2], + "level 19 must compress smaller than level 1 through the same reused context: {sizes:?}" + ); + assert_eq!(sizes[0], sizes[2], "same writer, same input, same level"); + assert_eq!(sizes[1], sizes[3], "same writer, same input, same level"); + } + + /// Accepts a fixed number of bytes, then fails every write. + struct FailingSink { + inner: Cursor>, + remaining: usize, + } + + impl Write for FailingSink { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + if buf.len() > self.remaining { + return Err(std::io::Error::other("sink full")); + } + self.remaining -= buf.len(); + self.inner.write(buf) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } + } + + impl Seek for FailingSink { + fn seek(&mut self, pos: SeekFrom) -> std::io::Result { + self.inner.seek(pos) + } + } + + /// A failed write must not poison the context: the next block through the same context + /// has to come out clean. Write-side mirror of `decode_context_usable_after_error`. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + fn codec_context_usable_after_write_error() { + let batch = test_batch(1, 100); + let writer = + ShuffleBlockWriter::try_new(&test_schema(), CompressionCodec::Zstd(1)).unwrap(); + let mut ctx = ShuffleCodecContext::default(); + + // Fits the 20-byte header but not the body: the encoder dies mid-frame. + let mut failing = FailingSink { + inner: Cursor::new(vec![]), + remaining: 64, + }; + assert!(writer + .write_batch(&batch, &mut failing, &mut ctx, &Time::default()) + .is_err()); + + let mut out = vec![]; + let mut cursor = Cursor::new(&mut out); + writer + .write_batch(&batch, &mut cursor, &mut ctx, &Time::default()) + .unwrap(); + assert_eq!(read_ipc_compressed(&out[16..]).unwrap(), batch); + } + + /// RSS encodes free the zstd context each time (its memory is only reserved per + /// invocation); local encodes keep it. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + fn rss_write_releases_zstd_context_local_write_retains_it() { + let batch = test_batch(2, 100); + let writer = + ShuffleBlockWriter::try_new(&test_schema(), CompressionCodec::Zstd(1)).unwrap(); + let mut ctx = ShuffleCodecContext::default(); + + let mut rss_out = vec![]; + let mut cursor = Cursor::new(&mut rss_out); + writer + .write_rss_batch(&batch, &mut cursor, &mut ctx, &Time::default()) + .unwrap(); + assert!( + !ctx.holds_zstd_cctx(), + "remote write must not retain the zstd context past its admitted invocation" + ); + assert_eq!(read_ipc_compressed(&rss_out[16..]).unwrap(), batch); + + let mut local_out = vec![]; + let mut cursor = Cursor::new(&mut local_out); + writer + .write_batch(&batch, &mut cursor, &mut ctx, &Time::default()) + .unwrap(); + assert!( + ctx.holds_zstd_cctx(), + "local write must keep the zstd context for reuse" + ); + assert_eq!(read_ipc_compressed(&local_out[16..]).unwrap(), batch); + } #[test] fn rss_zstd_workspace_accounts_for_compression_level_without_unbounded_estimator_loops() { From abf5807816659e56a7c6fc230d7e6669e6947727 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Mon, 31 Aug 2026 11:44:43 +0700 Subject: [PATCH 2/4] fix: bound retained zstd context memory outside encode bursts SessionOnly reset preserves zstd's allocated window, so a retained context grows to the largest workspace it has seen (~128 MiB for the decoder after one wide-window frame, ~834 MiB for the encoder at level 22) and stays there. Cap retained contexts at 8 MiB -- covering the commonly configured levels -- and drop anything larger after each decode (errors included) and each local block encode; higher levels fall back to per-frame creation, the pre-existing cost. The local writer also releases its context when a spill event or the final flush completes, so nothing is retained between write phases. --- native/shuffle/src/codec_context.rs | 46 +++++++++++-- native/shuffle/src/ipc.rs | 68 +++++++++++++++++++ .../src/partitioners/multi_partition.rs | 1 + native/shuffle/src/shuffle_writer.rs | 50 ++++++++++++++ .../writers/local/local_partition_writer.rs | 17 ++++- .../shuffle/src/writers/partition_writer.rs | 5 ++ .../src/writers/shuffle_block_writer.rs | 46 ++++++++++++- 7 files changed, 227 insertions(+), 6 deletions(-) diff --git a/native/shuffle/src/codec_context.rs b/native/shuffle/src/codec_context.rs index 052358ff76b..37c0c7b5303 100644 --- a/native/shuffle/src/codec_context.rs +++ b/native/shuffle/src/codec_context.rs @@ -19,14 +19,20 @@ use arrow::ipc::writer::CompressionContext; use std::io; use zstd::zstd_safe::{CCtx, CParameter, DCtx, ResetDirective}; +/// Largest zstd workspace worth caching between frames. Covers the commonly configured +/// levels; higher levels (tens to hundreds of MiB of window) fall back to a fresh context +/// per frame, which is what per-block encoding paid anyway. +const MAX_RETAINED_ZSTD_CONTEXT_BYTES: usize = 8 * 1024 * 1024; + /// Reusable compression state for encoding shuffle blocks. /// /// A zstd context costs about a megabyte and real setup time, so a task shares one across all /// the blocks it encodes instead of paying per block. Keep ownership task-scoped, never -/// per-output-partition -- a shuffle can have thousands of partitions. Local shuffle holds the -/// zstd context for the whole task; the remote (RSS) path frees it after each admitted encode -/// via [`Self::release_zstd`], since its memory accounting only reserves the workspace per -/// invocation. +/// per-output-partition -- a shuffle can have thousands of partitions. Local shuffle reuses +/// the zstd context between blocks but bounds what it retains +/// ([`Self::release_zstd_if_oversized`]) and drops it at spill/finish boundaries; the remote +/// (RSS) path frees it after each admitted encode via [`Self::release_zstd`], since its +/// memory accounting only reserves the workspace per invocation. #[derive(Default)] pub struct ShuffleCodecContext { /// Arrow's per-message IPC compression scratch, reused across encodes. @@ -66,6 +72,19 @@ impl ShuffleCodecContext { self.zstd = None; } + /// Drops the cached zstd context when its workspace outgrew + /// [`MAX_RETAINED_ZSTD_CONTEXT_BYTES`] (a session reset keeps the allocation); the next + /// encode re-creates it lazily. + pub(crate) fn release_zstd_if_oversized(&mut self) { + if self + .zstd + .as_ref() + .is_some_and(|cctx| cctx.sizeof() > MAX_RETAINED_ZSTD_CONTEXT_BYTES) + { + self.zstd = None; + } + } + /// Test hook for the release-vs-retain contract of the two encode paths. #[cfg(test)] pub(crate) fn holds_zstd_cctx(&self) -> bool { @@ -95,6 +114,25 @@ impl ShuffleDecodeContext { .map_err(map_zstd_error)?; Ok(dctx) } + + /// Drops the cached zstd context when its workspace outgrew + /// [`MAX_RETAINED_ZSTD_CONTEXT_BYTES`]: one frame advertising a large window grows the + /// context past 100 MiB, and a session reset keeps that allocation. + pub(crate) fn release_zstd_if_oversized(&mut self) { + if self + .zstd + .as_ref() + .is_some_and(|dctx| dctx.sizeof() > MAX_RETAINED_ZSTD_CONTEXT_BYTES) + { + self.zstd = None; + } + } + + /// Test hook for the retained-workspace bound on the decode path. + #[cfg(test)] + pub(crate) fn holds_zstd_dctx(&self) -> bool { + self.zstd.is_some() + } } fn map_zstd_error(code: usize) -> io::Error { diff --git a/native/shuffle/src/ipc.rs b/native/shuffle/src/ipc.rs index 2ccb04d2965..d3b0e9fe6a6 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -60,6 +60,19 @@ fn read_ipc_compressed_impl( decode_context: &mut ShuffleDecodeContext, bytes: &[u8], validate: bool, +) -> Result { + let result = decode_shuffle_frame(decode_context, bytes, validate); + // Decoding a frame with a large advertised window grows the context past 100 MiB, and + // contexts here are long-lived (thread-local per executor thread). Bound what survives + // the frame, whether it decoded or not. + decode_context.release_zstd_if_oversized(); + result +} + +fn decode_shuffle_frame( + decode_context: &mut ShuffleDecodeContext, + bytes: &[u8], + validate: bool, ) -> Result { let codec = bytes.get(..4).ok_or_else(|| { DataFusionError::Execution("Failed to decode batch: truncated compression codec".to_owned()) @@ -323,6 +336,61 @@ mod tests { } } + /// ZSTD shuffle frame whose header advertises the full level-22 window (streaming + /// encode with no pledged source size), so decoding it demands a >100 MiB workspace. + fn encode_zstd_wide_window(payload: &[u8]) -> Vec { + let mut bytes = b"ZSTD".to_vec(); + let mut writer = zstd::Encoder::new(&mut bytes, 22).unwrap(); + writer.write_all(payload).unwrap(); + writer.finish().unwrap(); + bytes + } + + /// Small-window frames keep the context cached for reuse; a frame that inflates the + /// workspace far past its usual size must not leave it pinned in a long-lived context. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn decode_context_drops_oversized_zstd_workspace() { + let mut ctx = ShuffleDecodeContext::default(); + let small = encode(b"ZSTD", &ipc_stream(1)); + assert_eq!( + read_ipc_compressed_with(&mut ctx, &small) + .unwrap() + .num_rows(), + 3 + ); + assert!( + ctx.holds_zstd_dctx(), + "a small-window frame must keep the context cached for reuse" + ); + + let wide = encode_zstd_wide_window(&ipc_stream(1)); + assert_eq!( + read_ipc_compressed_with(&mut ctx, &wide) + .unwrap() + .num_rows(), + 3 + ); + assert!( + !ctx.holds_zstd_dctx(), + "a wide-window frame must not leave its workspace cached" + ); + } + + /// The workspace bound must hold on the error path too: a failed decode of a + /// wide-window frame cannot leave the inflated context behind. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn decode_error_drops_oversized_zstd_workspace() { + let mut ctx = ShuffleDecodeContext::default(); + let wide = encode_zstd_wide_window(&ipc_stream(1)); + assert!(read_ipc_compressed_with(&mut ctx, &wide[..wide.len() - 7]).is_err()); + assert!( + !ctx.holds_zstd_dctx(), + "a failed wide-window decode must not leave its workspace cached" + ); + } + /// A truncated frame must not poison the context for the next valid one. #[test] #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 37ce85d1cd2..1908226e007 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -577,6 +577,7 @@ impl MultiPartitionShuffleRepartitioner { ) }) }; + self.partition_writer.write_burst_complete(); let memory_spilled_bytes = self .reservation diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index a8f138d0bc3..0810bdeca8b 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -566,6 +566,56 @@ mod test { repartitioner.insert_batch(batch.clone()).await.unwrap(); } + /// The zstd context is reused within one encode burst but must not survive past it: a + /// spill event and the final shuffle write each end with the context released. + #[tokio::test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + async fn local_writer_releases_zstd_context_at_burst_boundaries() { + let batch = create_batch(900); + let num_partitions = 2; + let runtime_env = create_runtime(512 * 1024); + let metrics_set = ExecutionPlanMetricsSet::new(); + let dir = tempfile::tempdir().unwrap(); + let shuffle_block_writer = + ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::Zstd(1)) + .unwrap(); + let local_partition_writer = LocalPartitionWriter::try_new( + dir.path().join("data.out").to_str().unwrap().to_string(), + dir.path().join("index.out").to_str().unwrap().to_string(), + shuffle_block_writer, + num_partitions, + 1024, + 1024 * 1024, + Arc::clone(&runtime_env), + ) + .unwrap(); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + local_partition_writer, + CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], num_partitions), + ShufflePartitionerMetrics::new(&metrics_set, 0), + runtime_env, + 1024, + false, + None, + ) + .unwrap(); + + repartitioner.insert_batch(batch.clone()).await.unwrap(); + repartitioner.spill(0).unwrap(); + assert!( + !repartitioner.partition_writer().holds_zstd_cctx(), + "a finished spill burst must not keep the zstd context cached" + ); + + repartitioner.insert_batch(batch.clone()).await.unwrap(); + repartitioner.shuffle_write().unwrap(); + assert!( + !repartitioner.partition_writer().holds_zstd_cctx(), + "finish_all must release the zstd context" + ); + } + #[tokio::test] async fn shuffle_partitioner_charges_shared_buffer_once() { // `insert_batch` slices a large batch into batch_size chunks that all share one backing diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index a75626a2fcc..dc70792974e 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -66,7 +66,8 @@ pub(crate) struct LocalPartitionWriter { output_index_file: String, data_output: DataOutput, /// Compression state shared by every block this task writes; the per-partition - /// `BufBatchWriter`s borrow it (see [`ShuffleCodecContext`]). + /// `BufBatchWriter`s borrow it (see [`ShuffleCodecContext`]). Retention is bounded: + /// released at spill/finish boundaries and whenever its workspace is oversized. codec_context: ShuffleCodecContext, /// Start offset of each partition in the data file, plus a trailing entry /// with the total length so partition sizes are simple offset differences. @@ -134,6 +135,11 @@ impl LocalPartitionWriter { }) } + #[cfg(test)] + pub(crate) fn holds_zstd_cctx(&self) -> bool { + self.codec_context.holds_zstd_cctx() + } + #[cfg(test)] pub(crate) fn get_spill_writers(&self) -> &Vec { match &self.data_output { @@ -319,6 +325,15 @@ impl PartitionWriter for LocalPartitionWriter { output_index.flush()?; write_timer.stop(); + // The shuffle output is complete; nothing else encodes through this context. + self.codec_context.release_zstd(); + Ok(()) } + + fn write_burst_complete(&mut self) { + // A spill burst just ended and the next encode may be a long time coming; the zstd + // workspace is native memory no reservation tracks, so don't sit on it. + self.codec_context.release_zstd(); + } } diff --git a/native/shuffle/src/writers/partition_writer.rs b/native/shuffle/src/writers/partition_writer.rs index 25b0e598df8..51562b86c09 100644 --- a/native/shuffle/src/writers/partition_writer.rs +++ b/native/shuffle/src/writers/partition_writer.rs @@ -68,4 +68,9 @@ pub(crate) trait PartitionWriter: Send + Sync { /// [`finish_partition`](PartitionWriter::finish_partition). fn finish_all(&mut self, metrics: &ShufflePartitionerMetrics) -> datafusion::common::Result<()>; + + /// Marks the end of one burst of [`write`](PartitionWriter::write) calls (a spill + /// event), letting the writer drop transient encode state. Staging more batches + /// afterwards is still allowed. + fn write_burst_complete(&mut self) {} } diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs b/native/shuffle/src/writers/shuffle_block_writer.rs index 3c4b6ae29ec..08abdf9543c 100644 --- a/native/shuffle/src/writers/shuffle_block_writer.rs +++ b/native/shuffle/src/writers/shuffle_block_writer.rs @@ -279,8 +279,12 @@ impl ShuffleBlockWriter { // RSS charges the zstd workspace (rss_codec_workspace) to each admitted encode // and releases the charge when it ends, success or not. Free the workspace inside // that window -- kept alive it would be native memory the reservation system no - // longer tracks. Local shuffle keeps the context for the whole task. + // longer tracks. codec_context.release_zstd(); + } else { + // Local shuffle reuses the context across blocks, but nothing reserves its + // memory: a high-level workspace (hundreds of MiB) must not outlive the block. + codec_context.release_zstd_if_oversized(); } encode_result?; @@ -452,6 +456,46 @@ mod tests { assert_eq!(sizes[1], sizes[3], "same writer, same input, same level"); } + /// Common zstd levels stay cached between local blocks; a high level allocates a + /// workspace of hundreds of MiB that must be dropped as soon as its block is done. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + fn local_write_drops_oversized_zstd_context() { + let batch = test_batch(3, 100); + let mut ctx = ShuffleCodecContext::default(); + + let fast = ShuffleBlockWriter::try_new(&test_schema(), CompressionCodec::Zstd(1)).unwrap(); + let mut fast_out = vec![]; + fast.write_batch( + &batch, + &mut Cursor::new(&mut fast_out), + &mut ctx, + &Time::default(), + ) + .unwrap(); + assert!( + ctx.holds_zstd_cctx(), + "a common-level workspace must stay cached for reuse" + ); + + let slow = ShuffleBlockWriter::try_new(&test_schema(), CompressionCodec::Zstd(22)).unwrap(); + let mut slow_out = vec![]; + slow.write_batch( + &batch, + &mut Cursor::new(&mut slow_out), + &mut ctx, + &Time::default(), + ) + .unwrap(); + assert!( + !ctx.holds_zstd_cctx(), + "a level-22 workspace must not stay cached past its block" + ); + + assert_eq!(read_ipc_compressed(&fast_out[16..]).unwrap(), batch); + assert_eq!(read_ipc_compressed(&slow_out[16..]).unwrap(), batch); + } + /// Accepts a fixed number of bytes, then fails every write. struct FailingSink { inner: Cursor>, From 7c907a84c1df17ea4089c9d8f5b91848d253c389 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Mon, 31 Aug 2026 22:05:33 +0700 Subject: [PATCH 3/4] fix: scope decode context to the scan operator and pin reuse in tests The shuffle scan operator now owns its zstd decode context and passes it through the caller-owned decode entry points, so retained memory dies with the operator instead of living as long as the executor thread; the thread-local remains only for the static JNI decode entry, which has nothing to own a context. The retained-size cap gains a measured level-to-workspace table next to the constant and boundary tests that fail loudly if a zstd upgrade moves levels across the cap, and test-only creation counters pin that a multi-partition burst creates one context rather than one per block. --- .../src/execution/operators/shuffle_scan.rs | 68 +++++++++++---- native/shuffle/src/codec_context.rs | 75 ++++++++++++++-- native/shuffle/src/ipc.rs | 78 +++++++++++++++-- native/shuffle/src/shuffle_writer.rs | 18 ++++ .../writers/local/local_partition_writer.rs | 5 ++ .../src/writers/shuffle_block_writer.rs | 87 +++++++++++++++++++ 6 files changed, 299 insertions(+), 32 deletions(-) diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs index 0be1b512769..a1ddaadd4b2 100644 --- a/native/core/src/execution/operators/shuffle_scan.rs +++ b/native/core/src/execution/operators/shuffle_scan.rs @@ -20,7 +20,10 @@ use crate::{ execution::{ operators::ExecutionError, planner::TEST_EXEC_CONTEXT_ID, - shuffle::{read_ipc_compressed, read_ipc_compressed_validated, validate_remote_schema}, + shuffle::{ + read_ipc_compressed_validated_with, read_ipc_compressed_with, validate_remote_schema, + ShuffleDecodeContext, + }, }, jvm_bridge::{jni_call, JVMClasses}, }; @@ -62,6 +65,10 @@ pub struct ShuffleScanExec { pub schema: SchemaRef, /// The current input batch, populated by get_next_batch() before poll_next(). pub batch: Arc>>, + /// Decode state (zstd workspace) reused across this operator's shuffle blocks. Owning it + /// here scopes any retained workspace to the operator: it is freed when the plan's + /// execution context drops, instead of persisting for the life of a decoding thread. + decode_context: Arc>, /// Cache of plan properties. cache: Arc, /// Metrics collector. @@ -108,6 +115,7 @@ impl ShuffleScanExec { input_source, data_types, batch: Arc::new(Mutex::new(None)), + decode_context: Arc::new(Mutex::new(ShuffleDecodeContext::default())), cache, metrics: metrics_set, baseline_metrics, @@ -137,6 +145,7 @@ impl ShuffleScanExec { self.exec_context_id, self.input_source.as_ref().unwrap().as_obj(), &self.data_types, + &mut self.decode_context.try_lock().unwrap(), &self.decode_time, self.requires_validation, )?; @@ -153,6 +162,7 @@ impl ShuffleScanExec { exec_context_id: i64, iter: &JObject, data_types: &[DataType], + decode_context: &mut ShuffleDecodeContext, decode_time: &Time, requires_validation: bool, ) -> Result { @@ -190,7 +200,12 @@ impl ShuffleScanExec { // Decode the compressed IPC data let mut timer = decode_time.timer(); - let batch = match decode_shuffle_batch(slice, data_types, requires_validation) { + let batch = match decode_shuffle_batch( + decode_context, + slice, + data_types, + requires_validation, + ) { Ok(batch) => batch, Err(failure) => { // Remote inputs must invalidate the failed shuffle generation even when @@ -225,18 +240,22 @@ impl ShuffleScanExec { } fn decode_shuffle_batch( + decode_context: &mut ShuffleDecodeContext, bytes: &[u8], expected_types: &[DataType], requires_validation: bool, ) -> DataFusionResult { if requires_validation { - let batch = read_ipc_compressed_validated(bytes)?; + let batch = read_ipc_compressed_validated_with(decode_context, bytes)?; // Validate before unpack_dictionary or cast_and_stamp_schema can hide an incompatible // wire type by changing values. Keep failures inside get_next's recovery callback. validate_remote_schema(&batch, expected_types)?; Ok(batch) } else { - check_column_count(read_ipc_compressed(bytes)?, expected_types.len()) + check_column_count( + read_ipc_compressed_with(decode_context, bytes)?, + expected_types.len(), + ) } } @@ -398,7 +417,10 @@ impl RecordBatchStream for ShuffleScanStream { #[cfg(test)] mod tests { - use crate::execution::shuffle::{CompressionCodec, ShuffleBlockWriter, ShuffleCodecContext}; + use crate::execution::shuffle::{ + read_ipc_compressed_validated_with, CompressionCodec, ShuffleBlockWriter, + ShuffleCodecContext, ShuffleDecodeContext, + }; use arrow::array::{Int32Array, RecordBatchOptions, StringArray, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; @@ -451,7 +473,9 @@ mod tests { }; payload[signed_flag] = 0; - let decoded = super::read_ipc_compressed_validated(&payload).unwrap(); + let decoded = + read_ipc_compressed_validated_with(&mut ShuffleDecodeContext::default(), &payload) + .unwrap(); assert_eq!(decoded.num_columns(), 1); assert_eq!(decoded.column(0).data_type(), &DataType::UInt32); assert_eq!( @@ -463,16 +487,18 @@ mod tests { .values(), &[u32::MAX] ); - let error = super::decode_shuffle_batch(&payload, &[DataType::Int32], true) - .unwrap_err() - .to_string(); + let mut decode_context = ShuffleDecodeContext::default(); + let error = + super::decode_shuffle_batch(&mut decode_context, &payload, &[DataType::Int32], true) + .unwrap_err() + .to_string(); assert!(error.contains("type mismatch at column 0"), "{error}"); assert!(error.contains("UInt32"), "{error}"); assert!(error.contains("Int32"), "{error}"); // The new logical validation is confined to remote inputs. assert_eq!( - super::decode_shuffle_batch(&payload, &[DataType::Int32], false) + super::decode_shuffle_batch(&mut decode_context, &payload, &[DataType::Int32], false) .unwrap() .column(0) .data_type(), @@ -489,7 +515,9 @@ mod tests { ) .unwrap(); let payload = uncompressed_shuffle_payload(&batch); - let decoded = super::decode_shuffle_batch(&payload, &[], true).unwrap(); + let decoded = + super::decode_shuffle_batch(&mut ShuffleDecodeContext::default(), &payload, &[], true) + .unwrap(); assert_eq!(decoded.num_columns(), 0); assert_eq!(decoded.num_rows(), 3); } @@ -614,8 +642,13 @@ mod tests { let body = &bytes[16..]; // Remote schema validation must preserve the writer's dictionary representation. - let decoded = - super::decode_shuffle_batch(body, &[DataType::Int32, DataType::Utf8], true).unwrap(); + let decoded = super::decode_shuffle_batch( + &mut ShuffleDecodeContext::default(), + body, + &[DataType::Int32, DataType::Utf8], + true, + ) + .unwrap(); assert!( matches!(decoded.column(1).data_type(), DataType::Dictionary(_, _)), "Expected dictionary-encoded column from IPC, got {:?}", @@ -688,8 +721,13 @@ mod tests { let declared = list_of_struct_type(true); let block = RecordBatch::try_from_iter([("payload", block_column)]).unwrap(); let payload = uncompressed_shuffle_payload(&block); - let decoded = - super::decode_shuffle_batch(&payload, std::slice::from_ref(&declared), true).unwrap(); + let decoded = super::decode_shuffle_batch( + &mut ShuffleDecodeContext::default(), + &payload, + std::slice::from_ref(&declared), + true, + ) + .unwrap(); let mut scan = ShuffleScanExec::new( super::super::super::planner::TEST_EXEC_CONTEXT_ID, None, diff --git a/native/shuffle/src/codec_context.rs b/native/shuffle/src/codec_context.rs index 37c0c7b5303..17ceddeadf9 100644 --- a/native/shuffle/src/codec_context.rs +++ b/native/shuffle/src/codec_context.rs @@ -22,6 +22,23 @@ use zstd::zstd_safe::{CCtx, CParameter, DCtx, ResetDirective}; /// Largest zstd workspace worth caching between frames. Covers the commonly configured /// levels; higher levels (tens to hundreds of MiB of window) fall back to a fresh context /// per frame, which is what per-block encoding paid anyway. +/// +/// Workspace sizes measured against zstd-sys 2.0.16+zstd.1.5.7 (`sizeof()` after one +/// streaming frame, no pledged source size). Levels 7/8 sit ~3% under the cap, so a zstd +/// upgrade can silently flip them to release-per-block; re-measure on any dependency bump. +/// +/// | level | CCtx after encode | DCtx after decode | +/// |-------|-------------------|-------------------| +/// | 1 | 1,369,617 | 1,013,552 | +/// | 2 | 2,090,513 | | +/// | 3 | 3,663,377 | 2,586,416 | +/// | 4 | 4,974,097 | | +/// | 5 | 5,498,385 | | +/// | 6 | 5,498,385 | | +/// | 7 | 8,119,825 | | +/// | 8 | 8,119,825 | | +/// | 9 | 15,459,857 (over) | | +/// | 19 | 93,848,207 (over) | 8,877,872 (over) | const MAX_RETAINED_ZSTD_CONTEXT_BYTES: usize = 8 * 1024 * 1024; /// Reusable compression state for encoding shuffle blocks. @@ -39,6 +56,10 @@ pub struct ShuffleCodecContext { pub(crate) arrow_ipc: CompressionContext, /// Lazily created, reused across blocks. zstd: Option>, + /// How many zstd contexts this value has created, so tests can assert that N blocks + /// cost fewer than N creations instead of only observing retained/released state. + #[cfg(test)] + zstd_creations: u32, } impl ShuffleCodecContext { @@ -51,13 +72,18 @@ impl ShuffleCodecContext { &mut self, level: i32, ) -> io::Result<(&mut CCtx<'static>, &mut CompressionContext)> { - let cctx = - match &mut self.zstd { - Some(cctx) => cctx, - none => none.insert(CCtx::try_create().ok_or_else(|| { + let cctx = match &mut self.zstd { + Some(cctx) => cctx, + none => { + #[cfg(test)] + { + self.zstd_creations += 1; + } + none.insert(CCtx::try_create().ok_or_else(|| { io::Error::other("failed to allocate zstd compression context") - })?), - }; + })?) + } + }; cctx.reset(ResetDirective::SessionOnly) .map_err(map_zstd_error)?; cctx.set_parameter(CParameter::CompressionLevel(level)) @@ -90,6 +116,12 @@ impl ShuffleCodecContext { pub(crate) fn holds_zstd_cctx(&self) -> bool { self.zstd.is_some() } + + /// Test hook: zstd contexts created so far, for pinning reuse across blocks. + #[cfg(test)] + pub(crate) fn creation_count(&self) -> u32 { + self.zstd_creations + } } /// Decode-side counterpart of [`ShuffleCodecContext`]: one context serves every frame a @@ -98,6 +130,19 @@ impl ShuffleCodecContext { pub struct ShuffleDecodeContext { /// Lazily created, reused across frames. zstd: Option>, + /// How many zstd contexts this value has created, so tests can assert that N frames + /// cost fewer than N creations instead of only observing retained/released state. + #[cfg(test)] + zstd_creations: u32, +} + +impl std::fmt::Debug for ShuffleDecodeContext { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // DCtx is an opaque FFI handle; report only whether a workspace is cached. + f.debug_struct("ShuffleDecodeContext") + .field("zstd_cached", &self.zstd.is_some()) + .finish() + } } impl ShuffleDecodeContext { @@ -106,9 +151,15 @@ impl ShuffleDecodeContext { pub(crate) fn zstd_dctx(&mut self) -> io::Result<&mut DCtx<'static>> { let dctx = match &mut self.zstd { Some(dctx) => dctx, - none => none.insert(DCtx::try_create().ok_or_else(|| { - io::Error::other("failed to allocate zstd decompression context") - })?), + none => { + #[cfg(test)] + { + self.zstd_creations += 1; + } + none.insert(DCtx::try_create().ok_or_else(|| { + io::Error::other("failed to allocate zstd decompression context") + })?) + } }; dctx.reset(ResetDirective::SessionOnly) .map_err(map_zstd_error)?; @@ -133,6 +184,12 @@ impl ShuffleDecodeContext { pub(crate) fn holds_zstd_dctx(&self) -> bool { self.zstd.is_some() } + + /// Test hook: zstd contexts created so far, for pinning reuse across frames. + #[cfg(test)] + pub(crate) fn creation_count(&self) -> u32 { + self.zstd_creations + } } fn map_zstd_error(code: usize) -> io::Error { diff --git a/native/shuffle/src/ipc.rs b/native/shuffle/src/ipc.rs index d3b0e9fe6a6..d98f170da1e 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -24,8 +24,10 @@ use std::cell::RefCell; use std::io::{Error, ErrorKind, Read}; thread_local! { - /// Backs the entry points below. They're called from many JVM task threads; a - /// thread-local gets each thread context reuse without changing any caller. + /// Backs the context-less entry points below. Their only production caller is the JVM's + /// static decodeShuffleBlock JNI export, which has no native object that could own a + /// context; everything else owns a [`ShuffleDecodeContext`] and uses the `_with` + /// variants, so retention there ends with the owner instead of the thread. static DECODE_CONTEXT: RefCell = RefCell::new(ShuffleDecodeContext::default()); } @@ -63,8 +65,8 @@ fn read_ipc_compressed_impl( ) -> Result { let result = decode_shuffle_frame(decode_context, bytes, validate); // Decoding a frame with a large advertised window grows the context past 100 MiB, and - // contexts here are long-lived (thread-local per executor thread). Bound what survives - // the frame, whether it decoded or not. + // contexts here are long-lived (operator-owned, or thread-local for the JNI entry + // points). Bound what survives the frame, whether it decoded or not. decode_context.release_zstd_if_oversized(); result } @@ -336,16 +338,76 @@ mod tests { } } - /// ZSTD shuffle frame whose header advertises the full level-22 window (streaming - /// encode with no pledged source size), so decoding it demands a >100 MiB workspace. - fn encode_zstd_wide_window(payload: &[u8]) -> Vec { + /// ZSTD shuffle frame written by a streaming encode with no pledged source size, so its + /// header advertises `level`'s full default window and the decoder must allocate it. + fn encode_zstd_at_level(payload: &[u8], level: i32) -> Vec { let mut bytes = b"ZSTD".to_vec(); - let mut writer = zstd::Encoder::new(&mut bytes, 22).unwrap(); + let mut writer = zstd::Encoder::new(&mut bytes, level).unwrap(); writer.write_all(payload).unwrap(); writer.finish().unwrap(); bytes } + /// Level 22's window makes decoding demand a >100 MiB workspace. + fn encode_zstd_wide_window(payload: &[u8]) -> Vec { + encode_zstd_at_level(payload, 22) + } + + /// N small-window frames must cost one context creation, not N; that is the point of + /// carrying a context at all. A level-19 frame decodes through the same context but + /// inflates its workspace to ~8.47 MiB (zstd-sys 2.0.16+zstd.1.5.7), past the 8 MiB + /// retention cap, so it must leave the context released and the next frame pays anew. + #[test] + #[cfg_attr(miri, ignore)] // Miri cannot call Zstd's C FFI. + fn decode_context_creations_track_retention_boundary() { + let mut ctx = ShuffleDecodeContext::default(); + let small = encode(b"ZSTD", &ipc_stream(1)); + for _ in 0..3 { + assert_eq!( + read_ipc_compressed_with(&mut ctx, &small) + .unwrap() + .num_rows(), + 3 + ); + } + assert_eq!( + ctx.creation_count(), + 1, + "small-window frames must share one context" + ); + assert!(ctx.holds_zstd_dctx()); + + let level19 = encode_zstd_at_level(&ipc_stream(1), 19); + assert_eq!( + read_ipc_compressed_with(&mut ctx, &level19) + .unwrap() + .num_rows(), + 3 + ); + assert_eq!( + ctx.creation_count(), + 1, + "the retained context serves the wide frame" + ); + assert!( + !ctx.holds_zstd_dctx(), + "a level-19 frame's workspace exceeds the retention cap and must be dropped" + ); + + assert_eq!( + read_ipc_compressed_with(&mut ctx, &small) + .unwrap() + .num_rows(), + 3 + ); + assert_eq!( + ctx.creation_count(), + 2, + "the dropped context is re-created lazily" + ); + assert!(ctx.holds_zstd_dctx()); + } + /// Small-window frames keep the context cached for reuse; a frame that inflates the /// workspace far past its usual size must not leave it pinned in a long-lived context. #[test] diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 0810bdeca8b..667b107233b 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -603,6 +603,19 @@ mod test { repartitioner.insert_batch(batch.clone()).await.unwrap(); repartitioner.spill(0).unwrap(); + assert!( + repartitioner + .partition_writer() + .get_spill_writers() + .iter() + .all(|writer| writer.has_spill_file()), + "the burst must encode blocks for every partition" + ); + assert_eq!( + repartitioner.partition_writer().zstd_creation_count(), + 1, + "one spill burst across all partitions must create the zstd context exactly once" + ); assert!( !repartitioner.partition_writer().holds_zstd_cctx(), "a finished spill burst must not keep the zstd context cached" @@ -610,6 +623,11 @@ mod test { repartitioner.insert_batch(batch.clone()).await.unwrap(); repartitioner.shuffle_write().unwrap(); + assert_eq!( + repartitioner.partition_writer().zstd_creation_count(), + 2, + "the next burst re-creates the context once, not per block" + ); assert!( !repartitioner.partition_writer().holds_zstd_cctx(), "finish_all must release the zstd context" diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index dc70792974e..870588c1132 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -140,6 +140,11 @@ impl LocalPartitionWriter { self.codec_context.holds_zstd_cctx() } + #[cfg(test)] + pub(crate) fn zstd_creation_count(&self) -> u32 { + self.codec_context.creation_count() + } + #[cfg(test)] pub(crate) fn get_spill_writers(&self) -> &Vec { match &self.data_output { diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs b/native/shuffle/src/writers/shuffle_block_writer.rs index 08abdf9543c..066235bcbf2 100644 --- a/native/shuffle/src/writers/shuffle_block_writer.rs +++ b/native/shuffle/src/writers/shuffle_block_writer.rs @@ -496,6 +496,93 @@ mod tests { assert_eq!(read_ipc_compressed(&slow_out[16..]).unwrap(), batch); } + /// Retention is only worth its complexity if consecutive blocks actually share one + /// context: two level-6 blocks must cost a single context creation. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + fn zstd_context_created_once_for_retained_level() { + let batch = test_batch(4, 100); + let writer = + ShuffleBlockWriter::try_new(&test_schema(), CompressionCodec::Zstd(6)).unwrap(); + let mut ctx = ShuffleCodecContext::default(); + for _ in 0..2 { + let mut out = vec![]; + writer + .write_batch( + &batch, + &mut Cursor::new(&mut out), + &mut ctx, + &Time::default(), + ) + .unwrap(); + assert_eq!(read_ipc_compressed(&out[16..]).unwrap(), batch); + } + assert_eq!( + ctx.creation_count(), + 1, + "the second block must reuse the first block's context" + ); + assert!(ctx.holds_zstd_cctx()); + } + + /// Level 9's workspace measures 15,459,857 bytes (zstd-sys 2.0.16+zstd.1.5.7), past the + /// 8 MiB retention cap, so each block pays its own context creation and release. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + fn zstd_context_recreated_per_block_past_retention_cap() { + let batch = test_batch(5, 100); + let writer = + ShuffleBlockWriter::try_new(&test_schema(), CompressionCodec::Zstd(9)).unwrap(); + let mut ctx = ShuffleCodecContext::default(); + for _ in 0..2 { + let mut out = vec![]; + writer + .write_batch( + &batch, + &mut Cursor::new(&mut out), + &mut ctx, + &Time::default(), + ) + .unwrap(); + assert_eq!(read_ipc_compressed(&out[16..]).unwrap(), batch); + assert!( + !ctx.holds_zstd_cctx(), + "a level-9 workspace must be released after every block" + ); + } + assert_eq!(ctx.creation_count(), 2); + } + + /// Level 8's workspace measures 8,119,825 bytes (zstd-sys 2.0.16+zstd.1.5.7) -- about 3% + /// under the retention cap. A zstd bump that grows it past the cap would turn off reuse + /// at the highest still-retained level with no other symptom; fail loudly here instead. + #[test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + fn zstd_context_retained_at_level_eight_near_cap() { + let batch = test_batch(6, 100); + let writer = + ShuffleBlockWriter::try_new(&test_schema(), CompressionCodec::Zstd(8)).unwrap(); + let mut ctx = ShuffleCodecContext::default(); + for _ in 0..2 { + let mut out = vec![]; + writer + .write_batch( + &batch, + &mut Cursor::new(&mut out), + &mut ctx, + &Time::default(), + ) + .unwrap(); + assert_eq!(read_ipc_compressed(&out[16..]).unwrap(), batch); + } + assert_eq!( + ctx.creation_count(), + 1, + "level 8 must stay under the retention cap and keep reusing one context" + ); + assert!(ctx.holds_zstd_cctx()); + } + /// Accepts a fixed number of bytes, then fails every write. struct FailingSink { inner: Cursor>, From 386bcd04244e4d3d0905f02da7cf18c7d61b8c85 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Tue, 1 Sep 2026 09:17:10 +0700 Subject: [PATCH 4/4] bench: add decoder-reuse microbenchmark for shuffle frames Prebuilt shuffle frames (numeric plus nullable string data) decoded through one reused context and through a fresh context per frame: repeated small zstd frames, a large-frame control, a sequence where a wide-window frame pushes the retained workspace past the cap before small frames resume, and an uncompressed control. Decoded results are asserted identical across variants before anything is measured. --- native/shuffle/Cargo.toml | 4 + native/shuffle/benches/ipc_decode.rs | 168 +++++++++++++++++++++++++++ 2 files changed, 172 insertions(+) create mode 100644 native/shuffle/benches/ipc_decode.rs diff --git a/native/shuffle/Cargo.toml b/native/shuffle/Cargo.toml index 71be932422c..3e1b5e1d18c 100644 --- a/native/shuffle/Cargo.toml +++ b/native/shuffle/Cargo.toml @@ -75,6 +75,10 @@ required-features = ["shuffle-bench"] name = "shuffle_writer" harness = false +[[bench]] +name = "ipc_decode" +harness = false + [[bench]] name = "row_columnar" harness = false diff --git a/native/shuffle/benches/ipc_decode.rs b/native/shuffle/benches/ipc_decode.rs new file mode 100644 index 00000000000..ce64f6f1d4b --- /dev/null +++ b/native/shuffle/benches/ipc_decode.rs @@ -0,0 +1,168 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Measures shuffle frame decoding with a reused [`ShuffleDecodeContext`] against a fresh +//! context per frame (the cost profile of creating decode state for every block). + +use arrow::array::{Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::metrics::Time; +use datafusion_comet_shuffle::{ + read_ipc_compressed_with, CompressionCodec, ShuffleBlockWriter, ShuffleCodecContext, + ShuffleDecodeContext, +}; +use std::hint::black_box; +use std::io::{Cursor, Write}; +use std::sync::Arc; + +const SMALL_ROWS: usize = 8192; +const SMALL_FRAMES: usize = 64; +const LARGE_ROWS: usize = SMALL_ROWS * SMALL_FRAMES; +/// One wide-window frame is inserted after every this many small frames in the +/// over-cap recovery scenario. +const WIDE_FRAME_INTERVAL: usize = 16; + +fn test_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + ])) +} + +fn make_batch(start: usize, rows: usize) -> RecordBatch { + let ids = Int64Array::from_iter_values((start..start + rows).map(|i| i as i64)); + let names = StringArray::from_iter( + (start..start + rows).map(|i| (i % 7 != 0).then(|| format!("row-{i}-payload"))), + ); + RecordBatch::try_new(test_schema(), vec![Arc::new(ids), Arc::new(names)]).unwrap() +} + +/// Encodes `batch` with the crate's block writer and strips the 16-byte block header +/// (length + field count), leaving the codec tag + payload that the decoder consumes. +fn make_frame(batch: &RecordBatch, codec: CompressionCodec) -> Vec { + let writer = ShuffleBlockWriter::try_new(batch.schema().as_ref(), codec).unwrap(); + let mut buffer = Vec::new(); + writer + .write_batch( + batch, + &mut Cursor::new(&mut buffer), + &mut ShuffleCodecContext::default(), + &Time::default(), + ) + .unwrap(); + buffer.split_off(16) +} + +/// A zstd frame written by a streaming level-19 encode with no pledged source size, so its +/// header advertises the level's full default window. Decoding it grows the context's +/// workspace past the retention cap, forcing the context to be dropped and re-created. +fn make_wide_window_frame(batch: &RecordBatch) -> Vec { + let uncompressed = make_frame(batch, CompressionCodec::None); + let ipc_payload = &uncompressed[4..]; + let mut bytes = b"ZSTD".to_vec(); + let mut encoder = zstd::Encoder::new(&mut bytes, 19).unwrap(); + encoder.write_all(ipc_payload).unwrap(); + encoder.finish().unwrap(); + bytes +} + +fn small_frames(codec: CompressionCodec) -> Vec<(Vec, RecordBatch)> { + (0..SMALL_FRAMES) + .map(|i| { + let batch = make_batch(i * SMALL_ROWS, SMALL_ROWS); + (make_frame(&batch, codec.clone()), batch) + }) + .collect() +} + +fn scenario_small_zstd() -> Vec<(Vec, RecordBatch)> { + small_frames(CompressionCodec::Zstd(3)) +} + +fn scenario_large_frame() -> Vec<(Vec, RecordBatch)> { + let batch = make_batch(0, LARGE_ROWS); + vec![(make_frame(&batch, CompressionCodec::Zstd(3)), batch)] +} + +fn scenario_over_cap_recovery() -> Vec<(Vec, RecordBatch)> { + let wide_batch = make_batch(0, SMALL_ROWS); + let wide_frame = make_wide_window_frame(&wide_batch); + let mut frames = Vec::new(); + for (i, entry) in scenario_small_zstd().into_iter().enumerate() { + frames.push(entry); + if (i + 1) % WIDE_FRAME_INTERVAL == 0 { + frames.push((wide_frame.clone(), wide_batch.clone())); + } + } + frames +} + +fn scenario_none() -> Vec<(Vec, RecordBatch)> { + small_frames(CompressionCodec::None) +} + +/// Both decode variants must produce the exact batches the frames were built from. +fn assert_variants_decode_identically(frames: &[(Vec, RecordBatch)]) { + let mut reused = ShuffleDecodeContext::default(); + for (frame, expected) in frames { + assert_eq!( + &read_ipc_compressed_with(&mut reused, frame).unwrap(), + expected + ); + let mut fresh = ShuffleDecodeContext::default(); + assert_eq!( + &read_ipc_compressed_with(&mut fresh, frame).unwrap(), + expected + ); + } +} + +fn criterion_benchmark(c: &mut Criterion) { + let scenarios = [ + ("small_frames_zstd3", scenario_small_zstd()), + ("large_frame_zstd3", scenario_large_frame()), + ("over_cap_recovery_zstd", scenario_over_cap_recovery()), + ("small_frames_none", scenario_none()), + ]; + + let mut group = c.benchmark_group("ipc_decode"); + for (name, frames) in &scenarios { + assert_variants_decode_identically(frames); + + group.bench_function(format!("{name}/reused"), |b| { + let mut context = ShuffleDecodeContext::default(); + b.iter(|| { + for (frame, _) in frames { + black_box(read_ipc_compressed_with(&mut context, frame).unwrap()); + } + }); + }); + group.bench_function(format!("{name}/fresh"), |b| { + b.iter(|| { + for (frame, _) in frames { + let mut context = ShuffleDecodeContext::default(); + black_box(read_ipc_compressed_with(&mut context, frame).unwrap()); + } + }); + }); + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches);