pingzh opened a new pull request, #5467: URL: https://github.com/apache/datafusion-comet/pull/5467
## Which issue does this PR close? Closes #5352. ## Rationale for this change Remote shuffle services such as Apache Celeborn are widely used in cloud-native Spark deployments, but Comet’s native shuffle currently cannot integrate with them. This prevents users from combining Comet’s native execution with existing Celeborn infrastructure. This change enables native Comet shuffle data to be written directly to Celeborn and read back through a native-aware shuffle reader while preserving existing fallback behavior. ## What changes are included in this PR? - Add a native RSS partition writer that encodes complete Arrow IPC shuffle frames and pushes them through task-scoped JNI callbacks. - Introduce destination-aware native shuffle planning and execution. - Add a Celeborn partition pusher with bounded frame sizes, executor-wide backpressure, asynchronous completion tracking, cancellation, retries, and cleanup. - Implement `CometCelebornShuffleManager`, native shuffle writers, and a Celeborn-backed shuffle reader. - Coordinate shuffle generations, speculative attempts, stage retries, and Spark commit ownership. - Support unmodified Apache Celeborn 0.7.0+ clients, including their public mapper-completion and partition-reader APIs. - Enable Celeborn-backed native shuffle only for supported plans when native shuffle is explicitly configured. - Preserve existing Spark/Celeborn fallback paths for unsupported plans, disabled configurations, and Spark I/O encryption. - Add configuration options for Celeborn integration, maximum shuffle frame size, and executor-wide in-flight shuffle bytes. - Document configuration requirements and limitations. Example configuration: ```properties 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 ``` ## How are these changes tested? The full Maven reactor passes **136 tests** across five suites: - `CelebornShufflePartitionPusherSuite`: 64 tests. - `CometCelebornNativeShuffleWriterSuite`: 11 tests. - `CometCelebornShuffleManagerSuite`: 20 tests. - `CometCelebornShuffleReaderSuite`: 25 tests. - `CometSparkSessionExtensionsSuite`: 16 tests. Native Rust coverage includes **36 passing tests** for shuffle plan compatibility, JNI callbacks, RSS partition writing, shuffle destinations, and native planner integration. Coverage includes public Celeborn API compatibility, nested Arrow batches, backpressure, asynchronous push completion, cancellation, task failures, speculative execution, shuffle retries, reducer reads, metrics, unsupported-plan fallback, and encrypted shuffle fallback. The native library builds successfully, and Rust formatting, ScalaStyle, and Spotless checks all pass. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
