diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs index 8778ce36998..b5a282f5a39 100644 --- a/native/core/src/execution/operators/shuffle_scan.rs +++ b/native/core/src/execution/operators/shuffle_scan.rs @@ -408,10 +408,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::IpcWriteContext; use arrow::record_batch::RecordBatch; use datafusion::physical_plan::metrics::Time; use std::io::Cursor; @@ -426,7 +425,7 @@ mod tests { .write_batch( batch, &mut output, - &mut IpcWriteContext::default(), + &mut ShuffleCodecContext::default(), &Time::new(), ) .unwrap(); @@ -542,7 +541,12 @@ mod tests { let mut buf = Cursor::new(Vec::new()); let ipc_time = Time::new(); writer - .write_batch(&batch, &mut buf, &mut IpcWriteContext::default(), &ipc_time) + .write_batch( + &batch, + &mut buf, + &mut ShuffleCodecContext::default(), + &ipc_time, + ) .unwrap(); // Read back (skip 16-byte header: 8 compressed_length + 8 field_count) @@ -612,7 +616,7 @@ mod tests { .write_batch( &dict_batch, &mut buf, - &mut IpcWriteContext::default(), + &mut ShuffleCodecContext::default(), &ipc_time, ) .unwrap(); diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index c9c088e2802..c535ccbb137 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::IpcWriteContext; 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 = IpcWriteContext::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(); }); }); @@ -285,14 +284,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 = IpcWriteContext::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(); }); }); @@ -322,7 +321,7 @@ fn ipc_context_reuse_benchmark(c: &mut Criterion) { let lifetime = if reuse { "reused" } else { "fresh" }; group.bench_function(format!("{name}/{rows}/{codec:?}/{lifetime}"), |b| { let ipc_time = Time::default(); - let mut context = IpcWriteContext::default(); + let mut context = ShuffleCodecContext::default(); let mut buffer = Vec::new(); // Warm the output buffer and the retained context before timing. writer @@ -335,7 +334,7 @@ fn ipc_context_reuse_benchmark(c: &mut Criterion) { .unwrap(); b.iter(|| { if !reuse { - context = IpcWriteContext::default(); + context = ShuffleCodecContext::default(); } buffer.clear(); writer diff --git a/native/shuffle/src/codec_context.rs b/native/shuffle/src/codec_context.rs new file mode 100644 index 00000000000..7f5e2e23e79 --- /dev/null +++ b/native/shuffle/src/codec_context.rs @@ -0,0 +1,129 @@ +// 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::IpcWriteContext; +use std::io; +use zstd::zstd_safe::{CCtx, CParameter, ResetDirective}; + +/// Largest zstd workspace worth caching between blocks. Covers the commonly configured +/// levels; higher levels (tens to hundreds of MiB of window) fall back to a fresh context +/// per block, 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 | +/// |-------|-------------------| +/// | 1 | 1,369,617 | +/// | 2 | 2,090,513 | +/// | 3 | 3,663,377 | +/// | 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) | +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 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. + pub(crate) arrow_ipc: IpcWriteContext, + /// 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 { + /// 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 IpcWriteContext)> { + 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)) + .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; + } + + /// 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 { + 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 + } +} + +fn map_zstd_error(code: usize) -> io::Error { + io::Error::other(zstd::zstd_safe::get_error_name(code)) +} diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 766634eb71e..938c13e0fda 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; @@ -30,6 +31,7 @@ mod spark_crc32c_hasher; pub mod spark_unsafe; pub(crate) mod writers; +pub use codec_context::ShuffleCodecContext; pub use comet_partitioning::CometPartitioning; pub use ipc::{read_ipc_compressed, read_ipc_compressed_validated}; pub use remote_schema::{decode_remote_shuffle_batch, validate_remote_schema}; diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 31986a71f27..37f67eacbd6 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -535,6 +535,7 @@ impl MultiPartitionShuffleRepartitioner { ) }) }; + self.partition_writer.write_burst_complete(); // Count the input capacity released from buffering by this spill, including a // rejected reservation. Shared allocations are charged once within a spill, but diff --git a/native/shuffle/src/remote_schema_tests.rs b/native/shuffle/src/remote_schema_tests.rs index 52fafa0cf6c..9d445a0241c 100644 --- a/native/shuffle/src/remote_schema_tests.rs +++ b/native/shuffle/src/remote_schema_tests.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::ShuffleCodecContext; use crate::{decode_remote_shuffle_batch, CompressionCodec, ShuffleBlockWriter}; use arrow::array::{ Array, ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, FixedSizeListArray, @@ -27,7 +28,6 @@ use arrow::datatypes::{ ArrowDictionaryKeyType, DataType, Field, Fields, Int16Type, Int32Type, Int64Type, Int8Type, Schema, UInt16Type, UInt32Type, UInt64Type, UInt8Type, }; -use arrow::ipc::writer::IpcWriteContext; use datafusion::physical_plan::metrics::Time; use std::io::Cursor; use std::sync::Arc; @@ -40,7 +40,7 @@ fn encoded_batch(batch: &RecordBatch, codec: CompressionCodec, rss: bool) -> Vec } .unwrap(); let mut output = Cursor::new(Vec::new()); - let mut context = IpcWriteContext::default(); + let mut context = ShuffleCodecContext::default(); if rss { writer .write_rss_batch(batch, &mut output, &mut context, &Time::default()) diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 8d668de7725..5c85168221b 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -408,10 +408,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::IpcWriteContext; use arrow::record_batch::RecordBatch; use arrow::row::{RowConverter, SortField}; use datafusion::datasource::memory::MemorySourceConfig; @@ -441,14 +440,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 = IpcWriteContext::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()); @@ -487,14 +481,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 = IpcWriteContext::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(); @@ -593,6 +582,74 @@ 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() + .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" + ); + + 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" + ); + } + #[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 @@ -1227,6 +1284,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(); @@ -1241,11 +1299,17 @@ mod test { let mut scratch = Vec::new(); for batch in &small_batches { buf_writer - .write(batch, &mut scratch, &encode_time, &write_time) + .write( + batch, + &mut scratch, + &mut codec_context, + &encode_time, + &write_time, + ) .unwrap(); } buf_writer - .flush(&mut scratch, &encode_time, &write_time) + .flush(&mut scratch, &mut codec_context, &encode_time, &write_time) .unwrap(); } @@ -1262,11 +1326,17 @@ mod test { let mut scratch = Vec::new(); for batch in &small_batches { buf_writer - .write(batch, &mut scratch, &encode_time, &write_time) + .write( + batch, + &mut scratch, + &mut codec_context, + &encode_time, + &write_time, + ) .unwrap(); } buf_writer - .flush(&mut scratch, &encode_time, &write_time) + .flush(&mut scratch, &mut codec_context, &encode_time, &write_time) .unwrap(); } @@ -1358,6 +1428,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(); { @@ -1371,11 +1442,17 @@ mod test { let mut scratch = Vec::new(); for batch in &inputs { buf_writer - .write(batch, &mut scratch, &encode_time, &write_time) + .write( + batch, + &mut scratch, + &mut codec_context, + &encode_time, + &write_time, + ) .unwrap(); } buf_writer - .flush(&mut scratch, &encode_time, &write_time) + .flush(&mut scratch, &mut codec_context, &encode_time, &write_time) .unwrap(); } diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs index a3f3b3d36fc..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::IpcWriteContext; 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 = IpcWriteContext::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 1719cc105f9..e6d98c70c39 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::IpcWriteContext; use datafusion::physical_plan::metrics::Time; use std::borrow::Borrow; use std::io::{Cursor, Seek, SeekFrom, Write}; @@ -38,11 +38,13 @@ 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_max_size: usize, - compression_context: IpcWriteContext, /// Coalesces small batches into target_batch_size before serialization. /// Lazily initialized on first write to capture the schema. coalescer: Option, @@ -68,7 +70,6 @@ impl, W: Write> BufBatchWriter { shuffle_block_writer, writer, buffer_max_size, - compression_context: IpcWriteContext::default(), coalescer: None, batch_size, #[cfg(debug_assertions)] @@ -112,6 +113,7 @@ impl, W: Write> BufBatchWriter { &mut self, batch: &RecordBatch, scratch: &mut Vec, + codec_context: &mut ShuffleCodecContext, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result { @@ -142,7 +144,8 @@ impl, W: Write> BufBatchWriter { let mut bytes_written = 0; for batch in &completed { - bytes_written += self.write_batch_to_buffer(batch, scratch, encode_time, write_time)?; + bytes_written += + self.write_batch_to_buffer(batch, scratch, codec_context, encode_time, write_time)?; } Ok(bytes_written) } @@ -152,6 +155,7 @@ impl, W: Write> BufBatchWriter { &mut self, batch: &RecordBatch, scratch: &mut Vec, + codec_context: &mut ShuffleCodecContext, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result { @@ -160,7 +164,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(); @@ -178,6 +182,7 @@ impl, W: Write> BufBatchWriter { pub(crate) fn flush( &mut self, scratch: &mut Vec, + codec_context: &mut ShuffleCodecContext, encode_time: &Time, write_time: &Time, ) -> datafusion::common::Result<()> { @@ -191,7 +196,7 @@ impl, W: Write> BufBatchWriter { } } for batch in &remaining { - self.write_batch_to_buffer(batch, scratch, encode_time, write_time)?; + self.write_batch_to_buffer(batch, scratch, codec_context, encode_time, write_time)?; } // Flush the scratch buffer to the underlying writer @@ -243,9 +248,14 @@ mod tests { .unwrap(); let mut output = Vec::new(); let time = Time::default(); + let mut codec_context = ShuffleCodecContext::default(); let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 << 20, 8192); - writer.write(&batch, scratch, &time, &time).unwrap(); - writer.flush(scratch, &time, &time).unwrap(); + writer + .write(&batch, scratch, &mut codec_context, &time, &time) + .unwrap(); + writer + .flush(scratch, &mut codec_context, &time, &time) + .unwrap(); output } @@ -293,7 +303,8 @@ mod tests { let time = Time::default(); let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 << 20, 8192); let mut dirty = vec![0xAB, 0xCD]; - let _ = writer.write(&batch, &mut dirty, &time, &time); + let mut codec_context = ShuffleCodecContext::default(); + let _ = writer.write(&batch, &mut dirty, &mut codec_context, &time, &time); } /// Swapping in a different scratch mid-writer would silently abandon any bytes still @@ -308,10 +319,13 @@ mod tests { let mut output = Vec::new(); let time = Time::default(); let mut writer = BufBatchWriter::new(block_writer, &mut output, 1 << 20, 8192); + let mut codec_context = ShuffleCodecContext::default(); let mut first = Vec::new(); - writer.write(&batch, &mut first, &time, &time).unwrap(); + writer + .write(&batch, &mut first, &mut codec_context, &time, &time) + .unwrap(); let mut second = Vec::new(); - let _ = writer.write(&batch, &mut second, &time, &time); + let _ = writer.write(&batch, &mut second, &mut codec_context, &time, &time); } /// A block that crosses `buffer_max_size` grows the scratch past the cap; `flush` @@ -328,15 +342,20 @@ mod tests { ShuffleBlockWriter::try_new(batch.schema().as_ref(), CompressionCodec::None).unwrap(); let mut output = Vec::new(); let time = Time::default(); + let mut codec_context = ShuffleCodecContext::default(); let mut scratch = Vec::new(); let mut writer = BufBatchWriter::new(block_writer, &mut output, buffer_max_size, batch_size); - writer.write(&batch, &mut scratch, &time, &time).unwrap(); + writer + .write(&batch, &mut scratch, &mut codec_context, &time, &time) + .unwrap(); assert!( scratch.capacity() > buffer_max_size, "oversized block must have grown the scratch past the cap" ); - writer.flush(&mut scratch, &time, &time).unwrap(); + writer + .flush(&mut scratch, &mut codec_context, &time, &time) + .unwrap(); assert!(scratch.is_empty()); assert!( scratch.capacity() <= buffer_max_size, @@ -354,13 +373,17 @@ mod tests { let mut output = Vec::new(); let mut scratch = Vec::new(); let mut writer = BufBatchWriter::new(block_writer, &mut output, large_cap, batch_size); - writer.write(&batch, &mut scratch, &time, &time).unwrap(); + writer + .write(&batch, &mut scratch, &mut codec_context, &time, &time) + .unwrap(); let cap_after_write = scratch.capacity(); assert!( cap_after_write > 0 && cap_after_write <= large_cap, "write must have serialized the batch into the scratch" ); - writer.flush(&mut scratch, &time, &time).unwrap(); + writer + .flush(&mut scratch, &mut codec_context, &time, &time) + .unwrap(); assert!(scratch.is_empty()); assert_eq!( scratch.capacity(), diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index e22e339f949..ae68cd1b28e 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; @@ -75,6 +76,10 @@ 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`]). 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. /// Has `num_output_partitions + 1` elements. @@ -136,6 +141,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, @@ -144,6 +150,16 @@ impl LocalPartitionWriter { }) } + #[cfg(test)] + pub(crate) fn holds_zstd_cctx(&self) -> bool { + 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 { @@ -179,7 +195,13 @@ impl PartitionWriter for LocalPartitionWriter { // `finish_all`. for batch in iter.by_ref() { let batch = batch?; - writer.write(&batch, scratch, &metrics.encode_time, &metrics.write_time)?; + writer.write( + &batch, + scratch, + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; } } DataOutput::Multi { @@ -191,7 +213,13 @@ 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, recycled_buffer)?; + spill_writers[pid].write( + iter, + &mut self.codec_context, + runtime, + metrics, + recycled_buffer, + )?; } } @@ -225,7 +253,13 @@ impl PartitionWriter for LocalPartitionWriter { // flushed once in `finish_all`. for batch in iter.by_ref() { let batch = batch?; - writer.write(&batch, scratch, &metrics.encode_time, &metrics.write_time)?; + writer.write( + &batch, + scratch, + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; } } DataOutput::Multi { @@ -259,18 +293,21 @@ impl PartitionWriter for LocalPartitionWriter { write_buffer_size, batch_size, ); + let codec_context = &mut self.codec_context; let result: datafusion::common::Result<()> = (|| { for batch in iter.by_ref() { let batch = batch?; buf_batch_writer.write( &batch, recycled_buffer, + codec_context, &metrics.encode_time, &metrics.write_time, )?; } buf_batch_writer.flush( recycled_buffer, + codec_context, &metrics.encode_time, &metrics.write_time, ) @@ -291,7 +328,12 @@ 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, scratch } => { - writer.flush(scratch, &metrics.encode_time, &metrics.write_time)?; + writer.flush( + scratch, + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; writer.writer_stream_position()? } DataOutput::Multi { output_writer, .. } => { @@ -323,8 +365,17 @@ 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(); + } } #[cfg(test)] diff --git a/native/shuffle/src/writers/local/spill.rs b/native/shuffle/src/writers/local/spill.rs index 77fe009046d..13a57384c22 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; @@ -51,12 +52,16 @@ 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. /// `recycled_buffer` is a scratch byte buffer shared by the sequential per-partition /// spill writes; it is left drained on return so one buffer's capacity serves every /// partition instead of each write regrowing its own. pub(crate) fn write>>( &mut self, iter: &mut I, + codec_context: &mut ShuffleCodecContext, runtime: &RuntimeEnv, metrics: &ShufflePartitionerMetrics, recycled_buffer: &mut Vec, @@ -74,6 +79,7 @@ impl SpillWriter { buf_batch_writer.write( &batch?, recycled_buffer, + codec_context, &metrics.encode_time, &metrics.write_time, )?; @@ -82,12 +88,14 @@ impl SpillWriter { buf_batch_writer.write( &batch, recycled_buffer, + codec_context, &metrics.encode_time, &metrics.write_time, )?; } buf_batch_writer.flush( recycled_buffer, + codec_context, &metrics.encode_time, &metrics.write_time, )?; @@ -267,6 +275,7 @@ mod tests { let mut spill = spill_writer(&batch, 10); let runtime = RuntimeEnv::default(); let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut codec_context = ShuffleCodecContext::default(); let mut recycled = Vec::new(); let mut iter = vec![ Ok(batch), @@ -275,7 +284,13 @@ mod tests { .into_iter(); assert!(spill - .write(&mut iter, &runtime, &metrics, &mut recycled) + .write( + &mut iter, + &mut codec_context, + &runtime, + &metrics, + &mut recycled + ) .is_err()); assert!( recycled.is_empty(), @@ -301,11 +316,18 @@ mod tests { let mut spill = spill_writer(&batch, 10); let runtime = pathless_backend::runtime(); let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut codec_context = ShuffleCodecContext::default(); let mut recycled = Vec::new(); let mut iter = vec![Ok(batch)].into_iter(); spill - .write(&mut iter, &runtime, &metrics, &mut recycled) + .write( + &mut iter, + &mut codec_context, + &runtime, + &metrics, + &mut recycled, + ) .unwrap(); assert!(spill.has_spill_file()); diff --git a/native/shuffle/src/writers/partition_writer.rs b/native/shuffle/src/writers/partition_writer.rs index 9b415f897da..0fcec17c00d 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 { /// [`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/rss/mod.rs b/native/shuffle/src/writers/rss/mod.rs index f4de8b850be..164680d4a62 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::IpcWriteContext; use arrow::record_batch::RecordBatch; use datafusion::common::{DataFusionError, Result}; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, Time}; @@ -216,7 +215,7 @@ mod tests { } .unwrap(); let time = Time::default(); - let write = |buffer: &mut Vec, context: &mut IpcWriteContext| { + let write = |buffer: &mut Vec, context: &mut ShuffleCodecContext| { buffer.clear(); let mut out = Cursor::new(buffer); if rss { @@ -230,13 +229,13 @@ mod tests { let mut results = Vec::new(); let mut outputs = Vec::new(); for reuse in [false, true] { - let mut context = IpcWriteContext::default(); + let mut context = ShuffleCodecContext::default(); let mut buffer = Vec::new(); write(&mut buffer, &mut context); let (counts, _) = allocations::measure(|| { for _ in 0..BLOCKS { if !reuse { - context = IpcWriteContext::default(); + context = ShuffleCodecContext::default(); } write(&mut buffer, &mut context); } @@ -500,14 +499,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 = IpcWriteContext::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 b77eb47bb9d..eb5774cc5fc 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::IpcWriteContext; use arrow_select::dictionary::garbage_collect_any_dictionary; use datafusion::common::{DataFusionError, Result}; use datafusion_comet_jni_bridge::errors::CometError; @@ -50,7 +50,10 @@ pub(crate) struct RssPartitionWriter { pusher: Arc, num_partitions: usize, max_frame_size: usize, - compression_context: IpcWriteContext, + /// 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, @@ -91,7 +94,7 @@ impl RssPartitionWriter { pusher, num_partitions, max_frame_size, - compression_context: IpcWriteContext::default(), + codec_context: ShuffleCodecContext::default(), next_partition_to_finish: 0, finished: false, failed: false, @@ -287,7 +290,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; @@ -1231,7 +1234,7 @@ mod buffer_tests { .write_rss_batch( &batch, &mut output, - &mut IpcWriteContext::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 46dd7086517..ed17973c889 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 IpcWriteContext, + 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 IpcWriteContext, + 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 IpcWriteContext, + codec_context: &mut ShuffleCodecContext, ipc_time: &Time, bounded_rss_codec: bool, ) -> Result { @@ -269,9 +273,52 @@ 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. + 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?; + + // 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 +329,345 @@ 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"); + } + + /// 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); + } + + /// 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>, + 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. + #[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() {