feat: complete Celeborn map-side shuffle push lifecycle (6/n)

.github/workflows/pr_build_linux.yml

@@ -329,6 +329,7 @@ jobs:
               org.apache.comet.exec.CometNativeShuffleSuite
               org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
               org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite

.github/workflows/pr_build_macos.yml

@@ -145,6 +145,7 @@ jobs:
               org.apache.comet.exec.CometNativeShuffleSuite
               org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
               org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite

native/Cargo.toml

@@ -39,6 +39,7 @@ rust-version = "1.88"
 
 [workspace.dependencies]
 arrow = { version = "58.4.0", features = ["prettyprint", "ffi", "chrono-tz"] }
+arrow-select = { version = "58.4.0" }
 async-trait = { version = "0.1" }
 bytes = { version = "1.11.1" }
 parquet = { version = "58.4.0", default-features = false, features = ["experimental"] }

native/core/src/execution/planner.rs

@@ -47,7 +47,7 @@ use crate::execution::{
     serde::to_arrow_datatype,
     shuffle::{SchemaAlignExec, ShuffleWriterDestination, ShuffleWriterExec},
 };
-use crate::jvm_bridge::{jni_call, JVMClasses, JavaShufflePartitionPusher, ShufflePartitionPusher};
+use crate::jvm_bridge::{jni_call, JVMClasses, ShufflePartitionPusher};
 use arrow::compute::CastOptions;
 use arrow::datatypes::{DataType, Field, FieldRef, Schema, TimeUnit, DECIMAL128_MAX_PRECISION};
 use arrow::ffi_stream::FFI_ArrowArrayStream;
@@ -3994,7 +3994,7 @@ fn shuffle_writer_destination(
 
             Ok(ShuffleWriterDestination::Rss {
                 pusher: Arc::clone(pusher),
-                max_frame_size: JavaShufflePartitionPusher::MAX_PAYLOAD_SIZE,
+                max_frame_size: pusher.max_frame_size(),
             })
         }
         None => Err(GeneralError(
@@ -4860,6 +4860,24 @@ mod tests {
         }
     }
 
+    struct BoundedShufflePartitionPusher {
+        max_frame_size: usize,
+    }
+
+    impl ShufflePartitionPusher for BoundedShufflePartitionPusher {
+        fn push_partition_data(
+            &self,
+            _partition_id: i32,
+            _data: &[u8],
+        ) -> datafusion::common::Result<()> {
+            Ok(())
+        }
+
+        fn max_frame_size(&self) -> usize {
+            self.max_frame_size
+        }
+    }
+
     fn local_shuffle_partition_writer(
         output_data_file: &str,
         output_index_file: &str,
@@ -5055,6 +5073,24 @@ mod tests {
         }
     }
 
+    #[test]
+    fn shuffle_partition_writer_uses_its_task_callback_frame_limit() {
+        let writer = spark_operator::ShuffleWriter {
+            partition_writer: Some(rss_shuffle_partition_writer()),
+            ..Default::default()
+        };
+        let callback: Arc<dyn ShufflePartitionPusher> = Arc::new(BoundedShufflePartitionPusher {
+            max_frame_size: 4096,
+        });
+
+        match super::shuffle_writer_destination(&writer, Some(&callback)).unwrap() {
+            ShuffleWriterDestination::Rss { max_frame_size, .. } => {
+                assert_eq!(max_frame_size, 4096);
+            }
+            destination => panic!("expected an RSS shuffle destination, got {destination:?}"),
+        }
+    }
+
     #[test]
     fn shuffle_partition_writer_rejects_rss_with_legacy_data_path() {
         let writer = spark_operator::ShuffleWriter {

native/jni-bridge/src/shuffle_partition_pusher.rs

@@ -29,6 +29,28 @@ const MAX_JVM_ARRAY_LENGTH: i32 = i32::MAX - 8;
 /// Implementations must remain safe when invoked from native execution
 /// threads that do not inherit Spark's task-local JVM state.
 pub trait ShufflePartitionPusher: Send + Sync {
+    /// Reserves encoding scratch before a complete shuffle frame is materialized.
+    fn reserve_partition_data(&self, _reservation_bytes: usize) -> Result<()> {
+        Ok(())
+    }
+
+    /// Acknowledges that encoding buffers and JNI references have been released, even on success.
+    /// An asynchronous implementation must also wait for its transport buffers before releasing
+    /// the reservation. Reserve, push, and release form one synchronous invocation on a worker.
+    fn release_partition_data_reservation(&self) -> Result<()> {
+        Ok(())
+    }
+
+    /// Returns the largest complete frame accepted by this callback.
+    fn max_frame_size(&self) -> usize {
+        JavaShufflePartitionPusher::MAX_PAYLOAD_SIZE
+    }
+
+    /// Returns the largest encoding reservation this callback can admit before allocation.
+    fn max_reservation_size(&self) -> usize {
+        usize::MAX
+    }
+
     /// Sends one complete, length-prefixed Arrow IPC shuffle block.
     fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()>;
 }
@@ -41,6 +63,10 @@ pub trait ShufflePartitionPusher: Send + Sync {
 pub struct JavaShufflePartitionPusher {
     callback: Global<JObject<'static>>,
     push_method: JMethodID,
+    reserve_method: JMethodID,
+    release_method: JMethodID,
+    max_frame_size: usize,
+    max_reservation_size: usize,
 }
 
 impl JavaShufflePartitionPusher {
@@ -66,11 +92,85 @@ impl JavaShufflePartitionPusher {
                 jni::jni_sig!("(I[BI)V"),
             )
             .map_err(CometError::from)?;
+        let reserve_method = env
+            .get_method_id(
+                &callback_class,
+                jni::jni_str!("reservePartitionData"),
+                jni::jni_sig!("(I)V"),
+            )
+            .map_err(CometError::from)?;
+        let release_method = env
+            .get_method_id(
+                &callback_class,
+                jni::jni_str!("releasePartitionDataReservation"),
+                jni::jni_sig!("()V"),
+            )
+            .map_err(CometError::from)?;
+        let max_frame_method = env
+            .get_method_id(
+                &callback_class,
+                jni::jni_str!("maxFrameBytes"),
+                jni::jni_sig!("()I"),
+            )
+            .map_err(CometError::from)?;
+        // SAFETY: the cached ID was resolved on this callback class with a no-argument int ABI.
+        let configured_maximum = unsafe {
+            env.call_method_unchecked(
+                callback,
+                max_frame_method,
+                ReturnType::Primitive(Primitive::Int),
+                &[],
+            )
+        };
+        if let Some(exception) = check_exception(env)? {
+            return Err(exception.into());
+        }
+        let configured_maximum = configured_maximum
+            .map_err(CometError::from)?
+            .i()
+            .map_err(CometError::from)?;
+        if configured_maximum <= 0 || configured_maximum > MAX_JVM_ARRAY_LENGTH {
+            return Err(DataFusionError::Execution(format!(
+                "Remote shuffle maximum frame size {configured_maximum} is outside the JVM array limit of {MAX_JVM_ARRAY_LENGTH} bytes"
+            )));
+        }
+        let max_reservation_method = env
+            .get_method_id(
+                &callback_class,
+                jni::jni_str!("maxReservationBytes"),
+                jni::jni_sig!("()I"),
+            )
+            .map_err(CometError::from)?;
+        // SAFETY: the ID was resolved on this callback class with a no-argument int ABI.
+        let reservation_maximum = unsafe {
+            env.call_method_unchecked(
+                callback,
+                max_reservation_method,
+                ReturnType::Primitive(Primitive::Int),
+                &[],
+            )
+        };
+        if let Some(exception) = check_exception(env)? {
+            return Err(exception.into());
+        }
+        let reservation_maximum = reservation_maximum
+            .map_err(CometError::from)?
+            .i()
+            .map_err(CometError::from)?;
+        if reservation_maximum <= 0 {
+            return Err(DataFusionError::Execution(
+                "Remote shuffle maximum encoding reservation must be positive".to_string(),
+            ));
+        }
         let callback = env.new_global_ref(callback).map_err(CometError::from)?;
 
         Ok(Self {
             callback,
             push_method,
+            reserve_method,
+            release_method,
+            max_frame_size: configured_maximum as usize,
+            max_reservation_size: reservation_maximum as usize,
         })
     }
 
@@ -92,6 +192,62 @@ impl JavaShufflePartitionPusher {
 }
 
 impl ShufflePartitionPusher for JavaShufflePartitionPusher {
+    fn reserve_partition_data(&self, reservation_bytes: usize) -> Result<()> {
+        let bytes = i32::try_from(reservation_bytes).map_err(|_| {
+            DataFusionError::Execution(format!(
+                "Remote shuffle encoding reservation {reservation_bytes} exceeds the JVM integer limit"
+            ))
+        })?;
+        if bytes <= 0 {
+            return Err(DataFusionError::Execution(
+                "Remote shuffle encoding reservation must be positive".to_string(),
+            ));
+        }
+        JVMClasses::with_env(|env| {
+            // SAFETY: this default callback method has the cached `(I)V` signature.
+            let result = unsafe {
+                env.call_method_unchecked(
+                    self.callback.as_obj(),
+                    self.reserve_method,
+                    ReturnType::Primitive(Primitive::Void),
+                    &[JValue::Int(bytes).as_jni()],
+                )
+            };
+            if let Some(exception) = check_exception(env)? {
+                return Err(exception.into());
+            }
+            result.map_err(CometError::from)?;
+            Ok(())
+        })
+    }
+
+    fn release_partition_data_reservation(&self) -> Result<()> {
+        JVMClasses::with_env(|env| {
+            // SAFETY: this default callback method has the cached `()V` signature.
+            let result = unsafe {
+                env.call_method_unchecked(
+                    self.callback.as_obj(),
+                    self.release_method,
+                    ReturnType::Primitive(Primitive::Void),
+                    &[],
+                )
+            };
+            if let Some(exception) = check_exception(env)? {
+                return Err(exception.into());
+            }
+            result.map_err(CometError::from)?;
+            Ok(())
+        })
+    }
+
+    fn max_frame_size(&self) -> usize {
+        self.max_frame_size
+    }
+
+    fn max_reservation_size(&self) -> usize {
+        self.max_reservation_size
+    }
+
     fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()> {
         let payload_length = Self::checked_payload_length(partition_id, data.len())?;
 

native/shuffle/Cargo.toml

@@ -30,6 +30,7 @@ publish = false
 
 [dependencies]
 arrow = { workspace = true }
+arrow-select = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 clap = { version = "4", features = ["derive"], optional = true }
@@ -49,7 +50,7 @@ parquet = { workspace = true, optional = true }
 simd-adler32 = "0.3.9"
 snap = "1.1"
 tokio = { version = "1", features = ["rt-multi-thread"] }
-zstd = "0.13.3"
+zstd = { version = "0.13.3", features = ["experimental"] }
 
 [dev-dependencies]
 criterion = { version = "0.7", features = ["async", "async_tokio", "async_std"] }

native/shuffle/src/shuffle_writer.rs

@@ -267,12 +267,12 @@ async fn external_shuffle(
 ) -> Result<SendableRecordBatchStream> {
     let schema = input.schema();
 
-    let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone())?;
     let mut repartitioner = match destination {
         ShuffleWriterDestination::Local {
             output_data_file,
             output_index_file,
         } => {
+            let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone())?;
             let writer = LocalPartitionWriter::try_new(
                 output_data_file,
                 output_index_file,
@@ -298,6 +298,8 @@ async fn external_shuffle(
             pusher,
             max_frame_size,
         } => {
+            let shuffle_block_writer =
+                ShuffleBlockWriter::try_new_rss(Arc::clone(&schema), codec.clone())?;
             let writer = RssPartitionWriter::try_new(
                 shuffle_block_writer,
                 pusher,