perf: cache the IPC schema across shuffle blocks on the reader

native/Cargo.toml

@@ -39,6 +39,7 @@ rust-version = "1.94.0"
 
 [workspace.dependencies]
 arrow = { version = "59.2.0", features = ["prettyprint", "ffi", "chrono-tz"] }
+arrow-data = { version = "59.2.0" }
 arrow-select = { version = "59.2.0" }
 async-trait = { version = "0.1" }
 bytes = { version = "1.11.1" }

native/core/src/execution/jni_api.rs

@@ -101,7 +101,7 @@ use tokio::sync::mpsc;
 use crate::execution::memory_pools::{create_memory_pool, parse_memory_pool_config};
 use crate::execution::operators::{ScanExec, ShuffleScanExec};
 use crate::execution::shuffle::{
-    decode_remote_shuffle_batch, read_ipc_compressed, CompressionCodec,
+    decode_remote_shuffle_batch_with, CompressionCodec, ShuffleBlockDecoder,
 };
 use crate::execution::spark_plan::SparkPlan;
 
@@ -1317,11 +1317,31 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock(
 ) -> jlong {
     try_unwrap_or_throw(&e, |env| {
         with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || {
-            decode_shuffle_block(env, byte_buffer, length, array_addrs, schema_addrs, None)
+            // This entry point carries no native handle, so the schema cache lives in a
+            // thread-local decoder. Spark runs one task per thread at a time and the decoder
+            // compares the schema message bytes before reusing a cached schema, so a decoder
+            // last used by a different task on this thread simply misses and re-parses.
+            LOCAL_SHUFFLE_BLOCK_DECODER.with_borrow_mut(|decoder| {
+                decode_shuffle_block(
+                    env,
+                    decoder,
+                    byte_buffer,
+                    length,
+                    array_addrs,
+                    schema_addrs,
+                    None,
+                )
+            })
         })
     })
 }
 
+thread_local! {
+    /// Schema-caching decoder for the handle-less local `decodeShuffleBlock` entry point.
+    static LOCAL_SHUFFLE_BLOCK_DECODER: std::cell::RefCell<ShuffleBlockDecoder> =
+        std::cell::RefCell::new(ShuffleBlockDecoder::new());
+}
+
 #[no_mangle]
 /// Parse the expected schema once for a remote shuffle iterator.
 ///
@@ -1338,14 +1358,19 @@ pub extern "system" fn Java_org_apache_comet_Native_createRemoteShuffleDecoder(
         })?;
         let decoder = RemoteShuffleDecoder {
             expected_types: schema.fields.iter().map(to_arrow_datatype).collect(),
+            decoder: ShuffleBlockDecoder::new(),
         };
         Ok(Box::into_raw(Box::new(decoder)) as jlong)
     })
 }
 
-/// Immutable decoding state owned by one JVM remote shuffle iterator, not shared across tasks.
+/// Decoding state owned by one JVM remote shuffle iterator, not shared across tasks. The JVM
+/// side serializes decode calls on a handle, which is what the schema cache relies on.
 struct RemoteShuffleDecoder {
     expected_types: Vec<ArrowDataType>,
+    /// Lives as long as the JVM-side reader holds the handle, so the IPC schema message is
+    /// parsed once per distinct schema encoding rather than once per block.
+    decoder: ShuffleBlockDecoder,
 }
 
 #[no_mangle]
@@ -1384,24 +1409,29 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlockWit
 ) -> jlong {
     try_unwrap_or_throw(&e, |env| {
         with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || {
-            let decoder = unsafe { (decoder_handle as *const RemoteShuffleDecoder).as_ref() }
+            // SAFETY: the handle was returned by `createRemoteShuffleDecoder`, has not been
+            // released, and the JVM side does not decode concurrently on one handle, so this
+            // is the only live reference for the duration of the call.
+            let remote = unsafe { (decoder_handle as *mut RemoteShuffleDecoder).as_mut() }
                 .ok_or_else(|| {
                     CometError::Internal("Remote shuffle decoder is not initialized".to_owned())
                 })?;
             decode_shuffle_block(
                 env,
+                &mut remote.decoder,
                 byte_buffer,
                 length,
                 array_addrs,
                 schema_addrs,
-                Some(&decoder.expected_types),
+                Some(&remote.expected_types),
             )
         })
     })
 }
 
 fn decode_shuffle_block(
     env: &mut Env,
+    decoder: &mut ShuffleBlockDecoder,
     byte_buffer: JByteBuffer,
     length: jint,
     array_addrs: JLongArray,
@@ -1414,9 +1444,9 @@ fn decode_shuffle_block(
     let batch = if let Some(expected_types) = expected_types {
         // Reject incompatible logical types, then decode dictionaries before JVM import. The
         // JVM importer supports fewer dictionary key/value layouts than the shuffle writer.
-        decode_remote_shuffle_batch(slice, expected_types)?
+        decode_remote_shuffle_batch_with(decoder, slice, expected_types)?
     } else {
-        read_ipc_compressed(slice)?
+        decoder.decode(slice)?
     };
     prepare_output(env, array_addrs, schema_addrs, batch, false)
 }

native/core/src/execution/operators/shuffle_scan.rs

@@ -20,7 +20,7 @@ use crate::{
     execution::{
         operators::ExecutionError,
         planner::TEST_EXEC_CONTEXT_ID,
-        shuffle::{decode_remote_shuffle_batch, read_ipc_compressed},
+        shuffle::{decode_remote_shuffle_batch_with, ShuffleBlockDecoder},
     },
     jvm_bridge::{jni_call, JVMClasses},
 };
@@ -73,6 +73,10 @@ pub struct ShuffleScanExec {
     decode_time: Time,
     /// Remote inputs require Arrow array and logical schema validation; queried once at construction.
     requires_validation: bool,
+    /// Block decoder held for the life of the scan so the IPC schema message, which every block
+    /// from the same writer repeats verbatim, is parsed once rather than per block. Behind a
+    /// mutex only because this exec is `Clone`; it is used from `get_next_batch` on one thread.
+    decoder: Arc<Mutex<ShuffleBlockDecoder>>,
 }
 
 impl ShuffleScanExec {
@@ -115,6 +119,7 @@ impl ShuffleScanExec {
             schema,
             decode_time,
             requires_validation,
+            decoder: Arc::new(Mutex::new(ShuffleBlockDecoder::new())),
         })
     }
 
@@ -134,12 +139,14 @@ impl ShuffleScanExec {
 
         let mut current_batch = self.batch.try_lock().unwrap();
         if current_batch.is_none() {
+            let mut decoder = self.decoder.try_lock().unwrap();
             let next_batch = Self::get_next(
                 self.exec_context_id,
                 self.input_source.as_ref().unwrap().as_obj(),
                 &self.data_types,
                 &self.decode_time,
                 self.requires_validation,
+                &mut decoder,
             )?;
             *current_batch = Some(next_batch);
         }
@@ -156,6 +163,7 @@ impl ShuffleScanExec {
         data_types: &[DataType],
         decode_time: &Time,
         requires_validation: bool,
+        decoder: &mut ShuffleBlockDecoder,
     ) -> Result<InputBatch, CometError> {
         if exec_context_id == TEST_EXEC_CONTEXT_ID {
             return Ok(InputBatch::EOF);
@@ -191,7 +199,8 @@ 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(decoder, slice, data_types, requires_validation)
+            {
                 Ok(batch) => batch,
                 Err(failure) => {
                     // Remote inputs must invalidate the failed shuffle generation even when
@@ -230,16 +239,17 @@ impl ShuffleScanExec {
 }
 
 fn decode_shuffle_batch(
+    decoder: &mut ShuffleBlockDecoder,
     bytes: &[u8],
     expected_types: &[DataType],
     requires_validation: bool,
 ) -> DataFusionResult<RecordBatch> {
     if requires_validation {
         // Validate logical types before decoding dictionaries or normalizing nested fields.
         // Keep both validation and normalization failures inside get_next's recovery callback.
-        decode_remote_shuffle_batch(bytes, expected_types)
+        decode_remote_shuffle_batch_with(decoder, bytes, expected_types)
     } else {
-        check_column_count(read_ipc_compressed(bytes)?, expected_types.len())
+        check_column_count(decoder.decode(bytes)?, expected_types.len())
     }
 }
 
@@ -474,19 +484,29 @@ mod tests {
                 .values(),
             &[u32::MAX]
         );
-        let error = super::decode_shuffle_batch(&payload, &[DataType::Int32], true)
-            .unwrap_err()
-            .to_string();
+        let error = super::decode_shuffle_batch(
+            &mut crate::execution::shuffle::ShuffleBlockDecoder::new(),
+            &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)
-                .unwrap()
-                .column(0)
-                .data_type(),
+            super::decode_shuffle_batch(
+                &mut crate::execution::shuffle::ShuffleBlockDecoder::new(),
+                &payload,
+                &[DataType::Int32],
+                false
+            )
+            .unwrap()
+            .column(0)
+            .data_type(),
             &DataType::UInt32
         );
     }
@@ -500,7 +520,13 @@ 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 crate::execution::shuffle::ShuffleBlockDecoder::new(),
+            &payload,
+            &[],
+            true,
+        )
+        .unwrap();
         assert_eq!(decoded.num_columns(), 0);
         assert_eq!(decoded.num_rows(), 3);
     }
@@ -621,14 +647,24 @@ mod tests {
 
         // Local decoding preserves the wire encoding for get_next to unpack. Remote decoding
         // validates and unpacks first, so the same result can also be safely imported by the JVM.
-        let local =
-            super::decode_shuffle_batch(body, &[DataType::Int32, DataType::Utf8], false).unwrap();
+        let local = super::decode_shuffle_batch(
+            &mut crate::execution::shuffle::ShuffleBlockDecoder::new(),
+            body,
+            &[DataType::Int32, DataType::Utf8],
+            false,
+        )
+        .unwrap();
         assert!(matches!(
             local.column(1).data_type(),
             DataType::Dictionary(_, _)
         ));
-        let decoded =
-            super::decode_shuffle_batch(body, &[DataType::Int32, DataType::Utf8], true).unwrap();
+        let decoded = super::decode_shuffle_batch(
+            &mut crate::execution::shuffle::ShuffleBlockDecoder::new(),
+            body,
+            &[DataType::Int32, DataType::Utf8],
+            true,
+        )
+        .unwrap();
         assert_eq!(decoded.column(1).data_type(), &DataType::Utf8);
 
         // Create ShuffleScanExec with value types (Utf8, not Dictionary) — this is
@@ -697,8 +733,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 crate::execution::shuffle::ShuffleBlockDecoder::new(),
+            &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

@@ -30,6 +30,7 @@ publish = false
 
 [dependencies]
 arrow = { workspace = true }
+arrow-data = { workspace = true }
 arrow-select = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }

native/shuffle/benches/shuffle_reader.rs

@@ -24,7 +24,9 @@ use arrow::ipc::reader::StreamReader;
 use arrow::ipc::writer::IpcWriteContext;
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use datafusion::physical_plan::metrics::Time;
-use datafusion_comet_shuffle::{read_ipc_compressed, CompressionCodec, ShuffleBlockWriter};
+use datafusion_comet_shuffle::{
+    read_ipc_compressed, CompressionCodec, ShuffleBlockDecoder, ShuffleBlockWriter,
+};
 use std::hint::black_box;
 use std::io::Cursor;
 use std::sync::Arc;
@@ -93,13 +95,24 @@ fn criterion_benchmark(c: &mut Criterion) {
 
             let id = format!("{num_columns}col_{num_rows}row");
 
-            // full decode: schema parse plus record batch
+            // full decode with a throwaway decoder: schema parse plus record batch
             group.bench_with_input(
                 BenchmarkId::new("decode_block", &id),
                 &uncompressed,
                 |b, block| b.iter(|| black_box(read_ipc_compressed(black_box(block)).unwrap())),
             );
 
+            // full decode with a decoder held across blocks, the way a reader holds it: the
+            // schema message is byte-identical every time, so it is served from the cache
+            group.bench_with_input(
+                BenchmarkId::new("decode_block_cached_schema", &id),
+                &uncompressed,
+                |b, block| {
+                    let mut decoder = ShuffleBlockDecoder::new();
+                    b.iter(|| black_box(decoder.decode(black_box(block)).unwrap()))
+                },
+            );
+
             // schema parse alone: `try_new` stops before the record batch. Skips the codec tag.
             group.bench_with_input(
                 BenchmarkId::new("parse_schema_only", &id),