perf: reuse zstd compression contexts across shuffle blocks

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<Mutex<Option<InputBatch>>>,
+    /// 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<Mutex<ShuffleDecodeContext>>,
     /// Cache of plan properties.
     cache: Arc<PlanProperties>,
     /// 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<InputBatch, CometError> {
@@ -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<RecordBatch> {
     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,16 +634,21 @@ mod tests {
             .write_batch(
                 &dict_batch,
                 &mut buf,
-                &mut CompressionContext::default(),
+                &mut ShuffleCodecContext::default(),
                 &ipc_time,
             )
             .unwrap();
         let bytes = buf.into_inner();
         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,

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

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<Schema> {
+    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<u8> {
+    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<u8> {
+    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<u8>, 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<u8>, RecordBatch)> {
+    small_frames(CompressionCodec::Zstd(3))
+}
+
+fn scenario_large_frame() -> Vec<(Vec<u8>, RecordBatch)> {
+    let batch = make_batch(0, LARGE_ROWS);
+    vec![(make_frame(&batch, CompressionCodec::Zstd(3)), batch)]
+}
+
+fn scenario_over_cap_recovery() -> Vec<(Vec<u8>, 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<u8>, 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<u8>, 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);