feat: support Apache Celeborn remote shuffle service

.github/workflows/pr_build_linux.yml

@@ -327,6 +327,10 @@ jobs:
               org.apache.comet.exec.CometShuffle4_0Suite
               org.apache.comet.exec.CometNativeColumnarToRowSuite
               org.apache.comet.exec.CometNativeShuffleSuite
+              org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite
               org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite

.github/workflows/pr_build_macos.yml

@@ -143,6 +143,10 @@ jobs:
               org.apache.comet.exec.CometShuffle4_0Suite
               org.apache.comet.exec.CometNativeColumnarToRowSuite
               org.apache.comet.exec.CometNativeShuffleSuite
+              org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite
               org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite

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

@@ -220,6 +220,32 @@ Comet provides a fully native shuffle implementation, which generally provides t
 supports `HashPartitioning`, `RangePartitioning` and `SinglePartitioning` but currently only supports primitive type
 partitioning keys. Columns that are not partitioning keys may contain complex types like maps, structs, and arrays.
 
+#### Apache Celeborn Remote Shuffle
+
+Comet native shuffle can write partition data to and read it from an existing Apache Celeborn remote shuffle
+service. Provide an Apache Celeborn 0.7.0 or newer Spark client on the driver and executor classpaths, then
+configure the composite shuffle manager and explicitly enable native shuffle:
+
+```
+spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager
+spark.celeborn.master.endpoints=celeborn-master:9097
+spark.comet.enabled=true
+spark.comet.exec.enabled=true
+spark.comet.shuffle.enabled=true
+spark.comet.shuffle.mode=native
+```
+
+The composite manager delegates ordinary Spark shuffles to Celeborn and uses native Comet shuffle only for
+supported exchanges whose child is already a Comet plan. Unsupported exchanges, the default `auto` shuffle mode,
+and JVM columnar shuffle remain on the existing Spark/Celeborn shuffle path. The Celeborn client is optional and
+is not bundled with Comet. Spark I/O encryption (`spark.io.encryption.enabled=true`) is not supported by the
+native remote shuffle path; encrypted applications continue to use the existing Spark/Celeborn shuffle path.
+
+Set `spark.comet.shuffle.celeborn.enabled=false` to disable Celeborn-backed partition pushers. The maximum
+complete partition frame size defaults to 64 MiB and can be configured with
+`spark.comet.shuffle.rss.maxFrameBytes`. The executor-wide in-flight push limit defaults to 256 MiB and can be
+configured with `spark.comet.shuffle.rss.maxInFlightBytes`.
+
 #### Columnar (JVM) Shuffle
 
 Comet Columnar shuffle is JVM-based and supports `HashPartitioning`, `RoundRobinPartitioning`, `RangePartitioning`, and

native/core/src/execution/jni_api.rs

@@ -21,8 +21,11 @@ use super::{serde, utils::SparkArrowConvert};
 use crate::{
     errors::{try_unwrap_or_throw, CometError, CometResult},
     execution::{
-        metrics::utils::update_comet_metric, planner::PhysicalPlanner, serde::to_arrow_datatype,
-        shuffle::spark_unsafe::row::process_sorted_row_partition, sort::RdxSort,
+        metrics::utils::update_comet_metric,
+        planner::{PhysicalPlanner, RegisteredShufflePusher},
+        serde::to_arrow_datatype,
+        shuffle::spark_unsafe::row::process_sorted_row_partition,
+        sort::RdxSort,
     },
     jvm_bridge::JVMClasses,
 };
@@ -41,7 +44,9 @@ use datafusion::{
     physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream},
     prelude::{SessionConfig, SessionContext},
 };
+use datafusion_comet_jni_bridge::shuffle_partition_pusher::JavaShufflePartitionPusher;
 use datafusion_comet_proto::spark_operator::Operator;
+use datafusion_comet_shuffle::PartitionPusher;
 use datafusion_comet_spark_expr::url_funcs::{CometParseUrl, CometTryParseUrl};
 use datafusion_spark::function::array::array_contains::SparkArrayContains;
 use datafusion_spark::function::array::repeat::SparkArrayRepeat;
@@ -120,6 +125,10 @@ use std::sync::OnceLock;
 #[cfg(feature = "jemalloc")]
 use tikv_jemalloc_ctl::{epoch, stats};
 
+#[cfg(test)]
+#[path = "rss_planner_jni_tests.rs"]
+mod rss_planner_jni_tests;
+
 static TOKIO_RUNTIME: Mutex<Option<Runtime>> = Mutex::new(None);
 
 #[cfg(feature = "jemalloc")]
@@ -394,6 +403,9 @@ struct ExecutionContext {
     pub class_loader: Option<Arc<Global<JObject<'static>>>>,
     /// Removes this context's tracing memory-pool entry on every exit path.
     memory_pool_registration: Option<ThreadMemoryPoolRegistration>,
+    /// Optional RSS callback registered by this task before the physical plan is built.
+    /// The callback's global JNI references are released with this execution context.
+    pub rss_pusher: Option<RegisteredShufflePusher>,
 }
 
 /// Accept serialized query plan and return the address of the native query plan.
@@ -583,13 +595,96 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
                 task_context,
                 class_loader,
                 memory_pool_registration,
+                rss_pusher: None,
             });
 
             Ok(Box::into_raw(exec_context) as i64)
         })
     })
 }
 
+/// Register an RSS callback on the native context that owns its Spark task attempt.
+///
+/// Registration must happen after `createPlan` and before the first `executePlan`. The handle is
+/// an opaque task-local identifier, not a native address or a process-wide registry key.
+///
+/// # Safety
+/// The execution context address must refer to a live native plan returned by `createPlan`.
+#[no_mangle]
+pub unsafe extern "system" fn Java_org_apache_comet_Native_registerRssPartitionPusher(
+    e: EnvUnowned,
+    _class: JClass,
+    exec_context: jlong,
+    handle: jlong,
+    object: JObject,
+    num_partitions: jint,
+    max_frame_bytes: jint,
+) {
+    try_unwrap_or_throw(&e, |env| {
+        if exec_context == 0 {
+            return Err(CometError::NullPointer(
+                "Comet execution context is null".into(),
+            ));
+        }
+
+        register_rss_partition_pusher(
+            env,
+            get_execution_context(exec_context),
+            handle,
+            &object,
+            num_partitions,
+            max_frame_bytes,
+        )
+    })
+}
+
+fn register_rss_partition_pusher(
+    env: &mut Env,
+    exec_context: &mut ExecutionContext,
+    handle: i64,
+    object: &JObject<'_>,
+    num_partitions: jint,
+    max_frame_bytes: jint,
+) -> CometResult<()> {
+    if handle <= 0 {
+        return Err(CometError::Config(
+            "RSS partition pusher handle must be positive".into(),
+        ));
+    }
+    if num_partitions <= 0 {
+        return Err(CometError::Config("Invalid RSS partition count".into()));
+    }
+    if max_frame_bytes <= 0 {
+        return Err(CometError::Config("Invalid RSS frame byte limit".into()));
+    }
+    if exec_context.root_op.is_some() {
+        return Err(CometError::Config(
+            "RSS partition pusher must be registered before execution".into(),
+        ));
+    }
+    if exec_context.rss_pusher.is_some() {
+        return Err(CometError::Config(
+            "RSS partition pusher is already registered for this task".into(),
+        ));
+    }
+
+    let num_partitions = num_partitions as usize;
+    let max_frame_bytes = max_frame_bytes as usize;
+    let pusher: Arc<dyn PartitionPusher> = Arc::new(JavaShufflePartitionPusher::try_new(
+        env,
+        object,
+        num_partitions,
+        max_frame_bytes,
+    )?);
+    exec_context.rss_pusher = Some(RegisteredShufflePusher {
+        handle,
+        num_partitions,
+        max_frame_bytes,
+        pusher,
+    });
+    Ok(())
+}
+
 /// Configure DataFusion session context.
 fn prepare_datafusion_session_context(
     batch_size: usize,
@@ -830,7 +925,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
                         .with_exec_id(exec_context_id)
                         .with_sql_text_pool(&exec_context.spark_plan)
                         .with_task_context(exec_context.task_context.clone())
-                        .with_class_loader(exec_context.class_loader.clone());
+                        .with_class_loader(exec_context.class_loader.clone())
+                        .with_rss_pusher(exec_context.rss_pusher.clone());
                 let (scans, shuffle_scans, root_op) = planner.create_plan(
                     &exec_context.spark_plan,
                     &mut exec_context.input_sources.clone(),