Support Apache Celeborn remote shuffle service

[pingzh]

What is the problem the feature request solves?

Remote shuffle services such as Apache Celeborn are widely used in cloud-native environments. Comet does not currently support Celeborn as a remote shuffle service.

I plan to work on this.

Describe the potential solution

The proposed implementation consists of approximately eight small PRs:

1. Add an RSS partition writer and JNI callback. Introduce a backend-neutral native partition writer and task-owned JVM callback for sending encoded Arrow IPC shuffle frames.

2. Add partition-writer destinations to shuffle plans. Extend the native shuffle plan to represent local versus remote partition writers while maintaining compatibility with existing shuffle plans.

3. Add destination-aware native shuffle execution. Allow native shuffle execution to write either to existing local shuffle files or to a remote partition writer without changing existing local behavior.

4. Bind task-owned JVM pushers to native shuffle planning. Connect each Spark map task’s JVM-side pusher to native shuffle execution through a task-scoped JNI handle.

5. Add a Celeborn shuffle manager and partition pusher. Introduce the Celeborn-specific adapter and a composite shuffle manager that delegates ordinary Spark shuffles to the existing Celeborn manager.

6. Complete the Celeborn map-side push lifecycle. Add bounded frame handling, backpressure, asynchronous push completion, cancellation, retries, speculative execution, commit coordination, and support for nested Arrow batches.

7. Add a native-aware Celeborn shuffle reader. Read remote Celeborn shuffle partitions through either Spark-compatible columnar reads or direct Comet native shuffle scans, with appropriate metrics and failure handling.

8. Enable native-only Celeborn shuffle planning. Enable the remote path only for explicitly configured, supported native shuffle plans; preserve fallback for unsupported plans, JVM shuffle, and Spark I/O encryption; add user-facing configuration and documentation.

flowchart TD
    Manager["Comet + Celeborn shuffle manager"]
    Fallback["Existing Spark / Celeborn shuffle"]

    Manager -->|"Supported native exchange"| Native
    Manager -->|"Ordinary or unsupported exchange"| Fallback

    subgraph Map["Map-side executor"]
        direction LR

        Native["Comet native operators"]
        Encode["Native partitioning + Arrow IPC"]
        Destination{"Partition writer"}
        Local["Existing local shuffle writer"]
        RSS["Backend-neutral RSS writer"]
        Callback["Task-owned JNI callback"]
        Client["Celeborn client adapter"]

        Native --> Encode --> Destination
        Destination -->|"Local"| Local
        Destination -->|"Remote"| RSS
        RSS --> Callback --> Client
    end

    Callback -.->|"Shared interface"| Other["Uniffle / other RSS adapters"]
    Client --> Celeborn[("Apache Celeborn")]

    subgraph Reduce["Reduce-side executor"]
        direction LR

        Reader["Celeborn-backed shuffle reader"]
        NativeScan["Comet native shuffle scan"]
        SparkRead["Spark-compatible columnar read"]

        Reader -->|"Raw shuffle frames"| NativeScan
        Reader -->|"Decoded batches"| SparkRead
    end

    Celeborn --> Reader

Loading

Additional context

Follow-up to #1241.

[wForget]

FYI: #4884 has basically implemented the uniffle remote shuffle service. I’ll complete it as soon as possible.

[pingzh]

thanks @wForget 🙏