perf: coalesce native shuffle batches before broadcast serialization

.github/workflows/pr_build_linux.yml

@@ -543,6 +543,7 @@ jobs:
               org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite
               org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite
               org.apache.spark.sql.comet.execution.shuffle.CometBlockStoreShuffleReaderSuite
+              org.apache.spark.sql.comet.execution.shuffle.NativeBatchDecoderIteratorSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite
               org.apache.comet.exec.DisableAQECometShuffleSuite

.github/workflows/pr_build_macos.yml

@@ -209,6 +209,7 @@ jobs:
               org.apache.spark.sql.comet.execution.shuffle.CometNativePositionalRoundRobinSuite
               org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite
               org.apache.spark.sql.comet.execution.shuffle.CometBlockStoreShuffleReaderSuite
+              org.apache.spark.sql.comet.execution.shuffle.NativeBatchDecoderIteratorSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite
               org.apache.comet.exec.DisableAQECometShuffleSuite

docs/source/contributor-guide/memory_management.md

@@ -120,6 +120,12 @@ the tasks of an executor that read it at the same time share one such copy. The
 usage log reports the allocator's total, and the part of it charged to the import allocator
 described below, so that the overhead can be sized for it.
 
+The opt-in broadcast producer batching path also retains native decoded shuffle arrays outside
+the execution pool. Each group is bounded by `min(spark.comet.batchSize, 65536)` rows, 8 MiB of
+retained/copy estimates, and 1,024 input batches, then concatenated once before Arrow export.
+Inputs and output coexist during that copy, and one decoded lookahead can be retained. The grouping
+limit therefore does not cap total reader RSS, and oversized input batches pass through unchanged.
+
 Sharing lowers the bytes that are resident, not the bytes that are reserved. Each task's native
 join reserves its whole build side, because DataFusion skips only the buffers that the same plan
 has already counted, so with N tasks sharing one copy the pool holds N times the bytes that are

docs/source/contributor-guide/native_shuffle.md

@@ -214,6 +214,14 @@ read time. See [Direct Read](#direct-read-shufflescan) below for how the choice
 3. The native `ShuffleScanExec` decodes each block with `read_ipc_compressed` and feeds the
    resulting `RecordBatch` straight into the plan. No Arrow FFI export or import happens.
 
+For a direct local native shuffle feeding a broadcast exchange, the opt-in
+`spark.comet.exec.broadcastExchange.producerCoalesce.enabled` mode uses one decoder across the task's fetched
+blocks. `DecodedBatchCoalescer` groups eligible small decoded batches before FFI export and
+`Utils.serializeBatches`. The read RDD is copied with the same dependency and AQE partition
+specifications, so other consumers of a reused shuffle keep their original reader. Grouping is
+bounded as described in [Memory Management](memory_management.md). Task completion releases
+pending native batches, exported Arrow batches, and the raw stream.
+
 **JVM decode**, used when direct read does not apply:
 
 1. `CometBlockStoreShuffleReader` fetches shuffle blocks via `ShuffleBlockFetcherIterator`.

docs/source/user-guide/latest/tuning/shuffle.md

@@ -53,6 +53,20 @@ scalar types. Hash partitioning keys must be scalar types unless
 and later) map keys. That setting is disabled by default until the performance of the nested hashing paths has been
 measured. Columns that are not partitioning keys may contain complex types like maps, structs, and arrays.
 
+### Batching shuffle inputs to broadcast
+
+The experimental `spark.comet.exec.broadcastExchange.producerCoalesce.enabled` setting combines small
+local native shuffle batches before broadcast serialization. It is disabled by default. It
+can reduce Arrow export, serialization, and driver assembly work when a broadcast build reads
+many small shuffle blocks.
+
+The broadcast must read directly from a local Comet native shuffle, including an AQE shuffle
+read, and have only INT, BIGINT, BOOLEAN, STRING, or BINARY columns. Grouping preserves the
+selected shuffle partitions and row order. Large batches pass through without a merge copy.
+
+Reader buffers are outside Spark's execution memory pool. The contributor guide describes the
+[grouping limits and memory lifetime](../../../contributor-guide/memory_management.md).
+
 ### Columnar (JVM) Shuffle
 
 Comet Columnar shuffle is JVM-based and supports `HashPartitioning`, `RoundRobinPartitioning`, `RangePartitioning`, and

native/core/src/execution/jni_api.rs

@@ -62,6 +62,7 @@ use datafusion::{
 use datafusion_comet_common::decode_string_arrays;
 use datafusion_comet_proto::spark_expression::agg_expr::ExprStruct as AggExprStruct;
 use datafusion_comet_proto::spark_operator::{AggregateMode, Operator, ShuffleScan};
+use datafusion_comet_shuffle::decoded_batch_coalescer::DecodedBatchCoalescer;
 use datafusion_comet_spark_expr::url_funcs::{CometParseUrl, CometTryParseUrl};
 use datafusion_spark::function::array::array_contains::SparkArrayContains;
 use datafusion_spark::function::array::repeat::SparkArrayRepeat;
@@ -2165,6 +2166,137 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock(
     })
 }
 
+/// Decode a local frame into owned arrays. The direct buffer can be reused after this call.
+fn append_shuffle_batch(
+    coalescer: &mut DecodedBatchCoalescer,
+    block: &[u8],
+    num_columns: usize,
+) -> CometResult<()> {
+    let batch = read_ipc_compressed(block)?;
+    if batch.num_columns() != num_columns {
+        return Err(CometError::Internal(
+            "Shuffle frame column count mismatch".into(),
+        ));
+    }
+    coalescer.push(batch)?;
+    Ok(())
+}
+
+/// The reader owns this handle and serializes native calls with cleanup.
+/// # Safety
+/// A nonzero handle must refer to a live coalescer exclusively borrowed for this call.
+unsafe fn shuffle_batch_coalescer<'a>(handle: jlong) -> CometResult<&'a mut DecodedBatchCoalescer> {
+    unsafe { (handle as *mut DecodedBatchCoalescer).as_mut() }
+        .ok_or_else(|| CometError::Internal("Closed shuffle batch coalescer".into()))
+}
+
+fn shuffle_batch_ready_columns(coalescer: &DecodedBatchCoalescer) -> jint {
+    coalescer
+        .peek_ready()
+        .map_or(-1, |batch| batch.num_columns() as jint)
+}
+
+#[no_mangle]
+pub extern "system" fn Java_org_apache_comet_Native_createShuffleBatchCoalescer(
+    e: EnvUnowned,
+    _class: JClass,
+    max_rows: jint,
+) -> jlong {
+    try_unwrap_or_throw(&e, |_| {
+        let coalescer = DecodedBatchCoalescer::try_new(max_rows.max(0) as usize)?;
+        Ok(Box::into_raw(Box::new(coalescer)) as jlong)
+    })
+}
+
+#[no_mangle]
+/// Decode and retain one local frame, returning ready columns or -1 for more input.
+/// # Safety
+/// The handle must be live and exclusively owned by this reader. On error, close the reader.
+pub unsafe extern "system" fn Java_org_apache_comet_Native_appendShuffleBatch(
+    e: EnvUnowned,
+    _class: JClass,
+    handle: jlong,
+    byte_buffer: JByteBuffer,
+    length: jint,
+    num_columns: jint,
+    tracing_enabled: jboolean,
+) -> jint {
+    try_unwrap_or_throw(&e, |env| {
+        with_trace("appendShuffleBatch", tracing_enabled != JNI_FALSE, || {
+            if num_columns < 0 {
+                return Err(CometError::Internal("Invalid shuffle frame bounds".into()));
+            }
+            let block = unsafe { shuffle_frame_bytes(env, &byte_buffer, length)? };
+            let coalescer = unsafe { shuffle_batch_coalescer(handle)? };
+            append_shuffle_batch(coalescer, block, num_columns as usize)?;
+            Ok(shuffle_batch_ready_columns(coalescer))
+        })
+    })
+}
+
+#[no_mangle]
+/// Report ready output, flushing the pending group at EOF.
+/// # Safety
+/// The handle must be live and exclusively owned by this reader.
+pub unsafe extern "system" fn Java_org_apache_comet_Native_peekShuffleBatch(
+    e: EnvUnowned,
+    _class: JClass,
+    handle: jlong,
+    end_of_input: jboolean,
+    tracing_enabled: jboolean,
+) -> jint {
+    try_unwrap_or_throw(&e, |_| {
+        let coalescer = unsafe { shuffle_batch_coalescer(handle)? };
+        if end_of_input != JNI_FALSE {
+            with_trace(
+                "coalesceShuffleBatches",
+                tracing_enabled != JNI_FALSE,
+                || coalescer.finish(),
+            )?;
+        }
+        Ok(shuffle_batch_ready_columns(coalescer))
+    })
+}
+
+#[no_mangle]
+/// Transfer one ready output using the ordinary export path, including offset repair.
+/// # Safety
+/// The handle and Arrow FFI output addresses must be live and exclusively owned by this reader.
+/// Both address arrays must have the column count reported by `peekShuffleBatch`.
+pub unsafe extern "system" fn Java_org_apache_comet_Native_takeShuffleBatch(
+    e: EnvUnowned,
+    _class: JClass,
+    handle: jlong,
+    array_addrs: JLongArray,
+    schema_addrs: JLongArray,
+    tracing_enabled: jboolean,
+) -> jlong {
+    try_unwrap_or_throw(&e, |env| {
+        with_trace("exportShuffleBatch", tracing_enabled != JNI_FALSE, || {
+            let batch = unsafe { shuffle_batch_coalescer(handle)? }
+                .take_ready()
+                .ok_or_else(|| CometError::Internal("No decoded shuffle batch is ready".into()))?;
+            prepare_output(env, array_addrs, schema_addrs, batch, false)
+        })
+    })
+}
+
+#[no_mangle]
+/// # Safety
+/// A nonzero handle must be live, must not be in use, and must be released exactly once.
+pub unsafe extern "system" fn Java_org_apache_comet_Native_closeShuffleBatchCoalescer(
+    e: EnvUnowned,
+    _class: JClass,
+    handle: jlong,
+) {
+    try_unwrap_or_throw(&e, |_| {
+        if handle != 0 {
+            drop(unsafe { Box::from_raw(handle as *mut DecodedBatchCoalescer) });
+        }
+        Ok(())
+    })
+}
+
 #[no_mangle]
 /// Parse the expected schema once for a remote shuffle iterator.
 ///
@@ -2251,6 +2383,22 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlockWit
     })
 }
 
+/// Borrows a codec header and IPC payload from a direct buffer.
+/// # Safety
+/// The buffer's backing memory must remain valid and unmodified for the borrow.
+unsafe fn shuffle_frame_bytes<'a>(
+    env: &mut Env,
+    byte_buffer: &'a JByteBuffer<'_>,
+    length: jint,
+) -> CometResult<&'a [u8]> {
+    let capacity = env.get_direct_buffer_capacity(byte_buffer)?;
+    if length < 4 || length as usize > capacity {
+        return Err(CometError::Internal("Invalid shuffle frame bounds".into()));
+    }
+    let pointer = env.get_direct_buffer_address(byte_buffer)?;
+    Ok(unsafe { std::slice::from_raw_parts(pointer, length as usize) })
+}
+
 /// Converts the JNI arguments of the `decodeShuffleBlock` entry points for
 /// [`decode_shuffle_block`].
 fn decode_shuffle_block_jni(
@@ -2261,9 +2409,7 @@ fn decode_shuffle_block_jni(
     schema_addrs: JLongArray,
     expected_types: Option<&[ArrowDataType]>,
 ) -> CometResult<jlong> {
-    let raw_pointer = env.get_direct_buffer_address(&byte_buffer)?;
-    let length = length as usize;
-    let block: &[u8] = unsafe { std::slice::from_raw_parts(raw_pointer, length) };
+    let block = unsafe { shuffle_frame_bytes(env, &byte_buffer, length)? };
     let array_addrs = unsafe { array_addrs.get_elements(env, ReleaseMode::NoCopyBack)? };
     let schema_addrs = unsafe { schema_addrs.get_elements(env, ReleaseMode::NoCopyBack)? };
     unsafe { decode_shuffle_block(block, &array_addrs, &schema_addrs, expected_types) }
@@ -3398,6 +3544,36 @@ mod tests {
         Ok((rows, arrow::array::make_array(data)))
     }
 
+    #[test]
+    fn append_shuffle_batches_owns_decoded_input_and_preserves_values() {
+        let input = int_batch();
+        let mut block = shuffle_block(&input);
+        let mut coalescer = DecodedBatchCoalescer::try_new(16).unwrap();
+        append_shuffle_batch(&mut coalescer, &block, 1).unwrap();
+        assert_eq!(shuffle_batch_ready_columns(&coalescer), -1);
+        // The JVM reuses its direct input buffer before the pending group is exported.
+        block.fill(0);
+        append_shuffle_batch(&mut coalescer, &shuffle_block(&input), 1).unwrap();
+        coalescer.finish().unwrap();
+        assert_eq!(shuffle_batch_ready_columns(&coalescer), 1);
+        let output = coalescer.take_ready().unwrap();
+        assert_eq!(output.num_rows(), 6);
+        assert_eq!(output.slice(0, 3), input);
+        assert_eq!(output.slice(3, 3), input);
+    }
+
+    #[test]
+    fn append_shuffle_batch_rejects_frame_count_mismatch() {
+        let mut coalescer = DecodedBatchCoalescer::try_new(16).unwrap();
+        let error =
+            append_shuffle_batch(&mut coalescer, &shuffle_block(&int_batch()), 2).unwrap_err();
+        assert!(error
+            .to_string()
+            .contains("Shuffle frame column count mismatch"));
+        coalescer.finish().unwrap();
+        assert!(coalescer.take_ready().is_none());
+    }
+
     #[test]
     fn decode_shuffle_block_exports_columns() {
         let batch = int_batch();