[DNM] feat: Support apache uniffle remote shuffle service

.github/workflows/pr_build_linux.yml

@@ -582,3 +582,79 @@ jobs:
           SPARK_TPCDS_JOIN_CONF: |
             spark.sql.autoBroadcastJoinThreshold=-1
             spark.sql.join.forceApplyShuffledHashJoin=true
+
+  # Uniffle integration test - verifies Comet shuffle against a local Uniffle cluster
+  uniffle-integration-test:
+    needs: build-native
+    name: Uniffle Integration Test
+    runs-on: ubuntu-24.04
+    container:
+      image: amd64/rust
+    env:
+      JAVA_TOOL_OPTIONS: --add-exports=java.base/sun.nio.ch=ALL-UNNAMED --add-exports=java.base/sun.util.calendar=ALL-UNNAMED --add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.lang=ALL-UNNAMED
+      UNIFFLE_VERSION: 0.10.0
+      HADOOP_VERSION: 2.10.2
+    steps:
+      - uses: actions/checkout@v7
+
+      - name: Setup Rust & Java toolchain
+        uses: ./.github/actions/setup-builder
+        with:
+          rust-version: ${{ env.RUST_VERSION }}
+          jdk-version: 17
+
+      - name: Download native library
+        uses: actions/download-artifact@v8
+        with:
+          name: native-lib-linux
+          path: native/target/release/
+
+      - name: Cache Maven dependencies
+        uses: actions/cache@v6
+        with:
+          path: |
+            ~/.m2/repository
+            /root/.m2/repository
+          key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }}
+          restore-keys: |
+            ${{ runner.os }}-java-maven-
+
+      - name: Install Uniffle
+        run: |
+          wget "https://www.apache.org/dyn/closer.lua/uniffle/${UNIFFLE_VERSION}/apache-uniffle-${UNIFFLE_VERSION}-bin.tar.gz?action=download" \
+            -O "/opt/apache-uniffle-${UNIFFLE_VERSION}-bin.tar.gz"
+          wget "https://www.apache.org/dyn/closer.lua/hadoop/common/hadoop-${HADOOP_VERSION}/hadoop-${HADOOP_VERSION}.tar.gz?action=download" \
+            -O "/opt/hadoop-${HADOOP_VERSION}.tar.gz"
+          mkdir /opt/uniffle
+          tar xzf "/opt/apache-uniffle-${UNIFFLE_VERSION}-bin.tar.gz" \
+            -C /opt/uniffle --strip-components=1
+          tar xzf "/opt/hadoop-${HADOOP_VERSION}.tar.gz" -C /opt/
+          mkdir /opt/uniffle/shuffle_data
+          printf 'XMX_SIZE=16g\nHADOOP_HOME=/opt/hadoop-%s\n' "${HADOOP_VERSION}" \
+            > /opt/uniffle/conf/rss-env.sh
+          printf 'rss.coordinator.shuffle.nodes.max 1\nrss.rpc.server.port 19999\n' \
+            > /opt/uniffle/conf/coordinator.conf
+          printf '%s\n' \
+            'rss.server.app.expired.withoutHeartbeat 7200000' \
+            'rss.server.heartbeat.delay 3000' \
+            'rss.rpc.server.port 19997' \
+            'rss.rpc.server.type GRPC_NETTY' \
+            'rss.jetty.http.port 19996' \
+            'rss.server.netty.port 19995' \
+            'rss.storage.basePath /opt/uniffle/shuffle_data' \
+            'rss.storage.type MEMORY_LOCALFILE' \
+            'rss.coordinator.quorum localhost:19999' \
+            'rss.server.flush.thread.alive 10' \
+            'rss.server.single.buffer.flush.threshold 64m' \
+            > /opt/uniffle/conf/server.conf
+          cd /opt/uniffle
+          bash ./bin/start-coordinator.sh
+          bash ./bin/start-shuffle-server.sh
+
+      - name: Build project
+        run: |
+          ./mvnw -B -Prelease install -DskipTests -Pspark-3.5,uniffle
+
+      - name: Run Uniffle integration tests
+        run: |
+          export COMET_SHUFFLE_MANAGER=uniffle && ./mvnw test -Pspark-3.5,uniffle,uniffle-comet-test -pl spark

native/core/src/execution/jni_api.rs

@@ -1218,6 +1218,7 @@ pub extern "system" fn Java_org_apache_comet_Native_getRustThreadId(
 
 use crate::execution::columnar_to_row::ColumnarToRowContext;
 use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema};
+use datafusion_comet_shuffle::RssPartitionPusher;
 use datafusion_spark::function::math::bin::SparkBin;
 use datafusion_spark::function::string::soundex::SparkSoundex;
 
@@ -1351,3 +1352,46 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose(
         Ok(())
     })
 }
+
+/// Create a native RSS partition pusher backed by the supplied Java pusher.
+///
+/// The returned handle owns the native pusher and must eventually be passed to
+/// `releaseRssPartitionPusher`.
+///
+/// # Safety
+/// `pusher` must be a valid JNI reference for the duration of this call.
+#[no_mangle]
+pub unsafe extern "system" fn Java_org_apache_comet_Native_createRssPartitionPusher(
+    e: EnvUnowned,
+    _class: JClass,
+    pusher: JObject,
+) -> jlong {
+    try_unwrap_or_throw(&e, |env| {
+        let rss_partition_pusher = Box::new(
+            RssPartitionPusher::try_new(Arc::new(jni_new_global_ref!(env, pusher)?)).unwrap(),
+        );
+
+        Ok(Box::into_raw(rss_partition_pusher) as i64)
+    })
+}
+
+/// Release a native RSS partition pusher created by `createRssPartitionPusher`.
+///
+/// # Safety
+/// `handle` must be a non-zero valid handle returned by `createRssPartitionPusher`. It must not be
+/// released more than once or used after this function returns.
+#[no_mangle]
+pub unsafe extern "system" fn Java_org_apache_comet_Native_releaseRssPartitionPusher(
+    e: EnvUnowned,
+    _class: JClass,
+    handle: jlong,
+) {
+    try_unwrap_or_throw(&e, |_env| {
+        debug_assert!(handle != 0, "releaseRssPartitionPusher: handle is null");
+        if handle != 0 {
+            let _pusher: Box<RssPartitionPusher> = Box::from_raw(handle as *mut RssPartitionPusher);
+            // _pusher is dropped here, freeing the resources
+        }
+        Ok(())
+    })
+}

native/core/src/execution/planner.rs

@@ -114,6 +114,9 @@ use datafusion::physical_plan::limit::GlobalLimitExec;
 use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec};
 use datafusion_comet_proto::spark_expression::ListLiteral;
 use datafusion_comet_proto::spark_operator::SparkFilePartition;
+use datafusion_comet_proto::spark_partitioning::partition_writer::PartitionWriterStruct::{
+    LocalPartitionWriter, RssPartitionWriter,
+};
 use datafusion_comet_proto::{
     spark_expression::{
         self, agg_expr::ExprStruct as AggExprStruct, expr::ExprStruct, literal::Value, AggExpr,
@@ -126,6 +129,7 @@ use datafusion_comet_proto::{
     },
     spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning},
 };
+use datafusion_comet_shuffle::ShufflePartitionWriter;
 use datafusion_comet_spark_expr::{
     jvm_udf::JvmScalarUdfExpr, ApproxPercentile, ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow,
     Correlation, Covariance, CreateNamedStruct, DecimalRescaleCheckOverflow, GetArrayStructFields,
@@ -1630,12 +1634,27 @@ impl PhysicalPlanner {
                 }?;
 
                 let write_buffer_size = writer.write_buffer_size as usize;
+                let partition_writer_struct = writer
+                    .partition_writer
+                    .as_ref()
+                    .and_then(|pw| pw.partition_writer_struct.clone())
+                    .ok_or_else(|| {
+                        GeneralError("ShuffleWriter is missing partition_writer".to_string())
+                    })?;
+                let shuffle_partition_writer = match partition_writer_struct {
+                    LocalPartitionWriter(local) => ShufflePartitionWriter::Local {
+                        output_data_file: local.output_data_file,
+                        output_index_file: local.output_index_file,
+                    },
+                    RssPartitionWriter(rss) => ShufflePartitionWriter::Rss {
+                        rss_partition_pusher_handle: rss.rss_partition_pusher,
+                    },
+                };
                 let shuffle_writer = Arc::new(ShuffleWriterExec::try_new(
                     writer_input,
                     partitioning,
                     codec,
-                    writer.output_data_file.clone(),
-                    writer.output_index_file.clone(),
+                    shuffle_partition_writer,
                     writer.tracing_enabled,
                     write_buffer_size,
                 )?);

native/jni-bridge/src/lib.rs

@@ -195,7 +195,9 @@ mod comet_s3_credential_dispatcher;
 mod comet_task_memory_manager;
 mod comet_udf_bridge;
 mod shuffle_block_iterator;
+mod shuffle_partition_pusher;
 
+use crate::shuffle_partition_pusher::ShufflePartitionPusher;
 use arrow_array_stream::ArrowArrayStream;
 pub use comet_metric_node::*;
 pub use comet_s3_credential_dispatcher::CometS3CredentialDispatcher;
@@ -238,6 +240,8 @@ pub struct JVMClasses<'a> {
     pub comet_udf_bridge: Option<CometUdfBridge<'a>>,
     /// JNI handles for the CometS3CredentialDispatcher SPI and the CometS3Credentials POJO.
     pub comet_s3_credential_dispatcher: CometS3CredentialDispatcher<'a>,
+    /// TODO: add comment
+    pub shuffle_partition_pusher: ShufflePartitionPusher<'a>,
 }
 
 unsafe impl Send for JVMClasses<'_> {}
@@ -316,6 +320,7 @@ impl JVMClasses<'_> {
                     bridge
                 },
                 comet_s3_credential_dispatcher: CometS3CredentialDispatcher::new(env).unwrap(),
+                shuffle_partition_pusher: ShufflePartitionPusher::new(env).unwrap(),
             }
         });
     }

native/jni-bridge/src/shuffle_partition_pusher.rs

@@ -0,0 +1,48 @@
+// 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 jni::{
+    errors::Result as JniResult,
+    objects::{JClass, JMethodID},
+    signature::{Primitive, ReturnType},
+    strings::JNIString,
+    Env,
+};
+
+pub struct ShufflePartitionPusher<'a> {
+    pub class: JClass<'a>,
+    pub method_push_partition_data: JMethodID,
+    pub method_push_partition_data_ret: ReturnType,
+}
+
+impl<'a> ShufflePartitionPusher<'a> {
+    pub const JVM_CLASS: &'static str = "org/apache/comet/shuffle/ShufflePartitionPusher";
+
+    pub fn new(env: &mut Env<'a>) -> JniResult<ShufflePartitionPusher<'a>> {
+        let class = env.find_class(JNIString::new(Self::JVM_CLASS))?;
+
+        Ok(ShufflePartitionPusher {
+            method_push_partition_data: env.get_method_id(
+                JNIString::new(Self::JVM_CLASS),
+                jni::jni_str!("pushPartitionData"),
+                jni::jni_sig!("(I[BI)I"),
+            )?,
+            method_push_partition_data_ret: ReturnType::Primitive(Primitive::Int),
+            class,
+        })
+    }
+}

native/proto/src/proto/operator.proto

@@ -324,8 +324,7 @@ enum CompressionCodec {
 
 message ShuffleWriter {
   spark.spark_partitioning.Partitioning partitioning = 1;
-  string output_data_file = 3;
-  string output_index_file = 4;
+  spark.spark_partitioning.PartitionWriter partition_writer = 2;
   CompressionCodec codec = 5;
   int32 compression_level = 6;
   bool tracing_enabled = 7;

native/proto/src/proto/partitioning.proto

@@ -58,3 +58,19 @@ message RoundRobinPartition {
   // Maximum number of columns to hash. 0 means no limit (hash all columns).
   int32 max_hash_columns = 2;
 }
+
+message PartitionWriter {
+  oneof partition_writer_struct {
+    LocalPartitionWriter local_partition_writer = 1;
+    RssPartitionWriter rss_partition_writer = 2;
+  }
+}
+
+message LocalPartitionWriter {
+  string output_data_file = 1;
+  string output_index_file = 2;
+}
+
+message RssPartitionWriter {
+  int64 rss_partition_pusher = 1;
+}

native/shuffle/Cargo.toml

@@ -41,7 +41,8 @@ datafusion-comet-jni-bridge = { workspace = true }
 datafusion-comet-spark-expr = { workspace = true }
 futures = { workspace = true }
 itertools = "0.15.0"
-jni = "0.21"
+jni = "0.22.4"
+paste = "1.0.14"
 log = "0.4"
 lz4_flex = { version = "0.13.1", default-features = false, features = ["frame"] }
 # parquet is only used by the shuffle_bench binary (shuffle-bench feature)
@@ -61,6 +62,9 @@ tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
 [features]
 shuffle-bench = ["clap", "parquet"]
 
+[package.metadata.cargo-machete]
+ignored = ["paste"]
+
 [lib]
 name = "datafusion_comet_shuffle"
 path = "src/lib.rs"

native/shuffle/benches/shuffle_writer.rs

@@ -30,7 +30,8 @@ use datafusion::{
     prelude::SessionContext,
 };
 use datafusion_comet_shuffle::{
-    CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShuffleWriterExec,
+    CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShufflePartitionWriter,
+    ShuffleWriterExec,
 };
 use itertools::Itertools;
 use std::io::Cursor;
@@ -149,8 +150,10 @@ fn create_shuffle_writer_exec(
         ))),
         partitioning,
         compression_codec,
-        "/tmp/data.out".to_string(),
-        "/tmp/index.out".to_string(),
+        ShufflePartitionWriter::Local {
+            output_data_file: "/tmp/data.out".to_string(),
+            output_index_file: "/tmp/index.out".to_string(),
+        },
         false,
         1024 * 1024,
     )

native/shuffle/src/bin/shuffle_bench.rs

@@ -45,7 +45,9 @@ use datafusion::physical_plan::common::collect;
 use datafusion::physical_plan::metrics::{MetricValue, MetricsSet};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
-use datafusion_comet_shuffle::{CometPartitioning, CompressionCodec, ShuffleWriterExec};
+use datafusion_comet_shuffle::{
+    CometPartitioning, CompressionCodec, ShufflePartitionWriter, ShuffleWriterExec,
+};
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use std::fs;
 use std::path::{Path, PathBuf};
@@ -476,8 +478,10 @@ async fn execute_shuffle_write(
         input,
         partitioning,
         codec,
-        data_file,
-        index_file,
+        ShufflePartitionWriter::Local {
+            output_data_file: data_file,
+            output_index_file: index_file,
+        },
         false,
         write_buffer_size,
     )

native/shuffle/src/lib.rs

@@ -28,5 +28,7 @@ pub(crate) mod writers;
 pub use comet_partitioning::CometPartitioning;
 pub use ipc::read_ipc_compressed;
 pub use schema_align::SchemaAlignExec;
+pub use shuffle_writer::ShufflePartitionWriter;
 pub use shuffle_writer::ShuffleWriterExec;
+pub use writers::RssPartitionPusher;
 pub use writers::{CompressionCodec, ShuffleBlockWriter};