diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs index e5014aa639e..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,10 +417,12 @@ impl RecordBatchStream for ShuffleScanStream { #[cfg(test)] mod tests { - use crate::execution::shuffle::{CompressionCodec, ShuffleBlockWriter}; + 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::ipc::writer::CompressionContext; use arrow::record_batch::RecordBatch; use datafusion::physical_plan::metrics::Time; use std::io::Cursor; @@ -416,7 +437,7 @@ mod tests { .write_batch( batch, &mut output, - &mut CompressionContext::default(), + &mut ShuffleCodecContext::default(), &Time::new(), ) .unwrap(); @@ -452,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!( @@ -464,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(), @@ -490,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); } @@ -535,7 +562,7 @@ mod tests { .write_batch( &batch, &mut buf, - &mut CompressionContext::default(), + &mut ShuffleCodecContext::default(), &ipc_time, ) .unwrap(); @@ -607,7 +634,7 @@ mod tests { .write_batch( &dict_batch, &mut buf, - &mut CompressionContext::default(), + &mut ShuffleCodecContext::default(), &ipc_time, ) .unwrap(); @@ -615,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 {:?}", @@ -689,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/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); 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..17ceddeadf9 --- /dev/null +++ b/native/shuffle/src/codec_context.rs @@ -0,0 +1,197 @@ +// 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}; + +/// 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. +/// +/// 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: 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 { + /// 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 => { + #[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 + } +} + +/// 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>, + /// 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 { + /// 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 => { + #[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)?; + 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() + } + + /// 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 { + 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..d98f170da1e 100644 --- a/native/shuffle/src/ipc.rs +++ b/native/shuffle/src/ipc.rs @@ -15,23 +15,67 @@ // 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 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()); +} + /// 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 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 (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 +} + +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()) })?; @@ -44,7 +88,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 +161,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 +318,150 @@ 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); + } + } + } + + /// 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, 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] + #[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. + 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/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 21d71dc1221..667b107233b 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(); @@ -577,6 +566,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 @@ -1160,6 +1217,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 +1230,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 +1250,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 +1347,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 +1359,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..870588c1132 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,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. @@ -121,6 +126,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, @@ -129,6 +135,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 { @@ -164,7 +180,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 +196,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 +230,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 +269,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 +294,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, .. } => { @@ -291,6 +330,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/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/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/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..066235bcbf2 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,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. 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() {