feat: add RSS partition writer and task-owned JNI callback (1/n)

native/jni-bridge/src/lib.rs

@@ -195,13 +195,15 @@ mod comet_s3_credential_dispatcher;
 mod comet_task_memory_manager;
 mod comet_udf_bridge;
 mod shuffle_block_iterator;
+mod shuffle_partition_pusher;
 
 use arrow_array_stream::ArrowArrayStream;
 pub use comet_metric_node::*;
 pub use comet_s3_credential_dispatcher::CometS3CredentialDispatcher;
 pub use comet_task_memory_manager::*;
 use comet_udf_bridge::CometUdfBridge;
 use shuffle_block_iterator::CometShuffleBlockIterator;
+pub use shuffle_partition_pusher::{JavaShufflePartitionPusher, ShufflePartitionPusher};
 
 /// The JVM classes that are used in the JNI calls.
 #[allow(dead_code)] // we need to keep references to Java items to prevent GC

native/jni-bridge/src/shuffle_partition_pusher.rs

@@ -0,0 +1,147 @@
+// 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 crate::{check_exception, errors::CometError, JVMClasses};
+use datafusion::common::{DataFusionError, Result};
+use jni::objects::{Global, JMethodID, JObject, JValue};
+use jni::signature::{Primitive, ReturnType};
+use jni::Env;
+
+/// Receives a complete encoded shuffle block for one output partition.
+///
+/// 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 {
+    /// Sends one complete, length-prefixed Arrow IPC shuffle block.
+    fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()>;
+}
+
+/// Invokes a task-owned JVM shuffle callback from any native execution thread.
+///
+/// The global callback reference keeps both the Java object and its class alive
+/// for the lifetime of the cached method ID. No thread-local JNI environment
+/// or Spark task context is retained between invocations.
+pub struct JavaShufflePartitionPusher {
+    callback: Global<JObject<'static>>,
+    push_method: JMethodID,
+}
+
+impl JavaShufflePartitionPusher {
+    /// Captures the callback while running on an attached JVM thread.
+    pub fn try_new(env: &mut Env<'_>, callback: &JObject<'_>) -> Result<Self> {
+        if callback.is_null() {
+            return Err(DataFusionError::Execution(
+                "Remote shuffle callback must not be null".to_string(),
+            ));
+        }
+
+        let callback_class = env.get_object_class(callback).map_err(CometError::from)?;
+        let push_method = env
+            .get_method_id(
+                &callback_class,
+                jni::jni_str!("pushPartitionData"),
+                jni::jni_sig!("(I[BI)V"),
+            )
+            .map_err(CometError::from)?;
+        let callback = env.new_global_ref(callback).map_err(CometError::from)?;
+
+        Ok(Self {
+            callback,
+            push_method,
+        })
+    }
+
+    fn checked_payload_length(partition_id: i32, payload_length: usize) -> Result<i32> {
+        if partition_id < 0 {
+            return Err(DataFusionError::Execution(format!(
+                "Remote shuffle partition must be nonnegative, got {partition_id}"
+            )));
+        }
+
+        i32::try_from(payload_length).map_err(|_| {
+            DataFusionError::Execution(format!(
+                "Remote shuffle payload size {payload_length} exceeds the JVM array limit"
+            ))
+        })
+    }
+}
+
+impl ShufflePartitionPusher for JavaShufflePartitionPusher {
+    fn push_partition_data(&self, partition_id: i32, data: &[u8]) -> Result<()> {
+        let payload_length = Self::checked_payload_length(partition_id, data.len())?;
+
+        JVMClasses::with_env(|env| {
+            let payload = env.byte_array_from_slice(data).map_err(CometError::from)?;
+
+            // SAFETY: `push_method` was resolved against the callback object's
+            // class with this exact argument list and void return type. The
+            // global object reference keeps its defining class alive.
+            let result = unsafe {
+                env.call_method_unchecked(
+                    self.callback.as_obj(),
+                    self.push_method,
+                    ReturnType::Primitive(Primitive::Void),
+                    &[
+                        JValue::Int(partition_id).as_jni(),
+                        JValue::Object(&payload).as_jni(),
+                        JValue::Int(payload_length).as_jni(),
+                    ],
+                )
+            };
+
+            // Inspect the pending exception before consuming the JNI result so
+            // its original throwable survives the DataFusion error boundary.
+            if let Some(exception) = check_exception(env)? {
+                return Err(exception.into());
+            }
+
+            result.map_err(CometError::from)?;
+            Ok(())
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::JavaShufflePartitionPusher;
+
+    #[test]
+    fn accepts_payload_lengths_representable_by_jvm_arrays() {
+        assert_eq!(
+            JavaShufflePartitionPusher::checked_payload_length(0, 0).unwrap(),
+            0
+        );
+        assert_eq!(
+            JavaShufflePartitionPusher::checked_payload_length(7, i32::MAX as usize).unwrap(),
+            i32::MAX
+        );
+    }
+
+    #[test]
+    fn rejects_negative_partition_ids() {
+        let error = JavaShufflePartitionPusher::checked_payload_length(-1, 1).unwrap_err();
+        assert!(error.to_string().contains("partition must be nonnegative"));
+    }
+
+    #[test]
+    fn rejects_payload_lengths_larger_than_a_jvm_array() {
+        let oversized_length = i32::MAX as usize + 1;
+        let error =
+            JavaShufflePartitionPusher::checked_payload_length(0, oversized_length).unwrap_err();
+        assert!(error.to_string().contains("exceeds the JVM array limit"));
+    }
+}

native/shuffle/src/lib.rs

@@ -29,4 +29,4 @@ pub use comet_partitioning::CometPartitioning;
 pub use ipc::read_ipc_compressed;
 pub use schema_align::SchemaAlignExec;
 pub use shuffle_writer::ShuffleWriterExec;
-pub use writers::{CompressionCodec, ShuffleBlockWriter};
+pub use writers::{CompressionCodec, RssPartitionWriter, ShuffleBlockWriter};

native/shuffle/src/writers/mod.rs

@@ -19,10 +19,12 @@ mod buf_batch_writer;
 mod checksum;
 mod local;
 mod partition_writer;
+mod rss;
 mod shuffle_block_writer;
 
 pub(crate) use buf_batch_writer::BufBatchWriter;
 pub(crate) use checksum::Checksum;
 pub(crate) use local::local_partition_writer::LocalPartitionWriter;
 pub(crate) use partition_writer::PartitionWriter;
+pub use rss::rss_partition_writer::RssPartitionWriter;
 pub use shuffle_block_writer::{CompressionCodec, ShuffleBlockWriter};