perf: reuse zstd compression contexts across shuffle blocks

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();

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

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<CCtx<'static>>,
+    /// 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))
+}

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};

native/shuffle/src/partitioners/multi_partition.rs

@@ -535,6 +535,7 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
                     )
                 })
             };
+            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

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())