[WIP]feat: Support apache uniffle remote shuffle service#4884
Draft
wForget wants to merge 6 commits into
Draft
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
wForget
changed the title
Member
Author
|
Most Comet Spark unit tests pass successfully. See: https://github.com/apache/datafusion-comet/actions/runs/29318829876/job/87049939110?pr=4884 The remaining failures are due to:
In a follow-up, I’ll enable only the TPC-H or TPC-DS integration tests for Uniffle. |
wForget
changed the title
There was a problem hiding this comment.
Pull request overview
This PR adds Apache Uniffle (RSS) as a remote shuffle backend for Comet native shuffle, while keeping the existing local-file native shuffle path. It introduces a pluggable “partition writer” abstraction end-to-end (Spark → protobuf → native planner → native shuffle writer) and adds CI coverage that runs Comet shuffle tests against a local Uniffle cluster.
Changes:
- Add a new
comet-unifflemodule withCometUniffleShuffleManager,CometUniffleShuffleWriter, andCometUniffleShuffleReaderto support remote shuffle reads and writes via Uniffle. - Refactor native shuffle writing into shared logic (
CometBaseNativeShuffleWriter) and extend the native shuffle writer to support both local output files and RSS push via a JNIShufflePartitionPusher. - Add Maven profiles and a dedicated GitHub Actions job to build with Uniffle and run TPC-H/TPC-DS Comet shuffle tests against a local Uniffle cluster.
Reviewed changes
Copilot reviewed 43 out of 43 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleWriter.scala | Implements an RSS-backed shuffle writer that pushes native-encoded blocks to Uniffle. |
| thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleReader.scala | Implements an RSS-backed shuffle reader plus native “encoded-block iterator” path. |
| thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleManager.scala | Wires Uniffle manager to Comet native shuffle dependency for read/write paths. |
| thirdparty/comet-uniffle/pom.xml | Adds the new Maven module and its dependencies. |
| spark/src/test/scala/org/apache/spark/sql/CometTPCHQuerySuite.scala | Switches test shuffle-manager config to be driven by CometTestShuffleManager. |
| spark/src/test/scala/org/apache/spark/sql/CometTPCDSQuerySuite.scala | Same as TPCH suite, enabling Uniffle via env-driven conf injection. |
| spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala | Adds CometTestShuffleManager to choose local vs Uniffle shuffle manager via env. |
| spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala | Forces local Comet shuffle manager configs for JVM columnar shuffle tests. |
| spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala | Minor adjustment to writer construction after CometNativeShuffleWriter refactor. |
| spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala | Adds explicit restriction preventing JVM columnar shuffle when Uniffle manager is used. |
| spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala | Uses CometNativeShuffleReader to obtain encoded shuffle blocks for native shuffle scan. |
| spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala | Refactors writer to reuse shared native write plan/exec logic. |
| spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleInputRDD.scala | Updates import to the new CometShuffleBlockIterator package. |
| spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala | Implements CometNativeShuffleReader for local shuffle blocks via a block iterator. |
| spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBaseNativeShuffleWriter.scala | New shared trait that builds shuffle writer plans and runs native execution. |
| spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala | Updates imports for the relocated CometShuffleBlockIterator. |
| spark/src/main/scala/org/apache/comet/Native.scala | Adds JNI entrypoints to create/release native RSS partition pushers. |
| spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala | Treats Uniffle shuffle manager as Comet-compatible for enabling shuffle execution. |
| spark/src/main/scala/org/apache/comet/CometExecIterator.scala | Updates imports for the relocated shuffle block iterator interface. |
| spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java | New JVM callback interface for native RSS shuffle writers. |
| spark/src/main/java/org/apache/comet/shuffle/CometShuffleBlockIterator.java | New interface for supplying encoded shuffle blocks to native via JNI. |
| spark/src/main/java/org/apache/comet/shuffle/CometNativeShuffleReader.java | New interface for readers that can provide encoded blocks for native shuffle scan. |
| spark/src/main/java/org/apache/comet/CometLocalShuffleBlockIterator.java | Renames/implements iterator as CometLocalShuffleBlockIterator using the new interface. |
| spark/pom.xml | Adds uniffle-comet-test profile and adjusts dependencies/classpath for Uniffle test runs. |
| spark-integration/pom.xml | Removes exclusions now handled via dependency optionality/shading behavior. |
| pom.xml | Adds uniffle profile/module and Uniffle dependency management plus duplicate-finder ignores. |
| native/shuffle/src/writers/rss/rss_partition_writer.rs | Adds RSS partition writer that flushes encoded blocks via a pusher per partition. |
| native/shuffle/src/writers/rss/rss_partition_pusher.rs | Adds JNI-backed pusher implementing Write to forward bytes to the JVM callback. |
| native/shuffle/src/writers/rss/mod.rs | New RSS writer module exports. |
| native/shuffle/src/writers/mod.rs | Wires RSS writer/pusher into the shuffle writer crate exports. |
| native/shuffle/src/shuffle_writer.rs | Introduces ShufflePartitionWriter enum and routes shuffle output to local or RSS backend. |
| native/shuffle/src/lib.rs | Exposes new shuffle writer/pusher types publicly from the crate. |
| native/shuffle/src/bin/shuffle_bench.rs | Updates benchmark harness to use the new ShufflePartitionWriter::Local form. |
| native/shuffle/Cargo.toml | Updates JNI dependency and adds paste (with cargo-machete ignore). |
| native/shuffle/benches/shuffle_writer.rs | Updates benches for the new ShufflePartitionWriter::Local API. |
| native/proto/src/proto/partitioning.proto | Adds PartitionWriter oneof with local and RSS writer configs. |
| native/proto/src/proto/operator.proto | Replaces shuffle writer output paths with a PartitionWriter field. |
| native/jni-bridge/src/shuffle_partition_pusher.rs | Adds cached JNI method metadata for ShufflePartitionPusher.pushPartitionData. |
| native/jni-bridge/src/shuffle_block_iterator.rs | Updates JNI class path for the moved CometShuffleBlockIterator. |
| native/jni-bridge/src/lib.rs | Registers ShufflePartitionPusher into the cached JVMClasses set. |
| native/core/src/execution/planner.rs | Decodes protobuf PartitionWriter into native ShufflePartitionWriter enum. |
| native/core/src/execution/jni_api.rs | Implements JNI createRssPartitionPusher and releaseRssPartitionPusher. |
| .github/workflows/pr_build_linux.yml | Adds a Uniffle integration-test job that spins up Uniffle and runs Comet shuffle suites. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+265
to
+273
| val oldLimit = current.limit() | ||
| try { | ||
| current.limit(current.position() + headerBuf.capacity()) | ||
| headerBuf.clear() | ||
| headerBuf.put(current) | ||
| headerBuf.flip() | ||
| } finally { | ||
| current.limit(oldLimit) | ||
| } |
Comment on lines
+278
to
+286
| // Subtract 8 because compressedLength includes the 8-byte field count we already read | ||
| val bytesToRead = compressedLength - 8 | ||
| if (bytesToRead > Integer.MAX_VALUE) { | ||
| throw new IllegalStateException( | ||
| "Native shuffle block size of " + bytesToRead + " exceeds maximum of " | ||
| + Integer.MAX_VALUE + ". Try reducing spark.comet.columnar.shuffle.batch.size.") | ||
| } | ||
| currentBlockLength = bytesToRead.toInt | ||
|
|
Comment on lines
+54
to
+62
| let length = buf.len() as i32; | ||
| JVMClasses::with_env(|env| { | ||
| let jbytes = env.byte_array_from_slice(buf).unwrap(); | ||
| let length: i32 = unsafe { | ||
| jni_call!(env, | ||
| shuffle_partition_pusher(self.jobject.as_ref()).push_partition_data(self.pid, &jbytes, length) -> i32)? | ||
| }; | ||
| Ok(length) | ||
| }) |
| currentBlockLength | ||
| } | ||
|
|
||
| private def nextShuffleBlock(): Option[ByteBuffer] = { |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #4913.
Rationale for this change
Comet native shuffle currently relies on local shuffle files, which prevents it from being used in deployments backed by a remote shuffle service.
This PR adds Apache Uniffle support while preserving the existing local shuffle path. Comet can push encoded native shuffle blocks to Uniffle and read them back for either regular Spark execution or native shuffle scans.
What changes are included in this PR?
LocalPartitionWriterwrites local data and index files as before.RssPartitionWriterpushes encoded partition data to a remote shuffle service.RssPartitionPusherthat forwards native shuffle data to the JVM.comet-unifflemodule containing:CometUniffleShuffleManagerCometUniffleShuffleWriterCometUniffleShuffleReaderNotes: The initial integration supports native shuffle mode only. JVM columnar shuffle with Uniffle is not supported by this PR.
How are these changes tested?
A dedicated GitHub Actions job:
uniffleprofile.CometUniffleShuffleManager.Most Comet Spark tests pass with Uniffle. The remaining failures are caused by the current native-only limitation and tests that assume deterministic shuffle ordering without an explicit sort.
Test run: https://github.com/apache/datafusion-comet/actions/runs/29318829876/job/87049939110?pr=4884