feat: Support Apache Uniffle remote shuffle service for Comet native shuffle

[wForget]

What is the problem the feature request solves?

Support Apache Uniffle remote shuffle service for Comet native shuffle

Describe the potential solution

Note

The following solution proposal and UML diagrams were generated by AI based on the implementation in #4884 and should be reviewed for technical accuracy.

Introduce a pluggable partition writer and shuffle block reader abstraction, with Apache Uniffle as the first remote shuffle service implementation.

Shuffle writer

The native shuffle writer should support two output backends:

• LocalPartitionWriter, which continues writing local data and index files.
• RssPartitionWriter, which pushes encoded partition data to a remote shuffle service.

CometUniffleShuffleWriter creates a native RssPartitionPusher handle and passes it to the native shuffle plan. The native planner uses the handle to construct an RssPartitionWriter.

RssPartitionWriter maintains one buffered writer per output partition. Each writer encodes Comet shuffle blocks and writes them through an RssPartitionPusher.

The native RssPartitionPusher implements std::io::Write. Its write() implementation calls the JVM ShufflePartitionPusher.pushPartitionData(partitionId, bytes, length) interface through JNI.

CometUniffleShuffleWriter implements this JVM interface by forwarding the data to Uniffle's buffer manager. Uniffle is then responsible for creating, sending, retrying, and committing the remote shuffle blocks.

classDiagram
    class PartitionWriter {
        <<interface>>
        +write(partitionId, batches, metrics)
        +finishPartition(partitionId, batches, metrics)
        +finishAll(metrics)
    }

    class RssPartitionWriter {
        -partitionWriters
        +write(partitionId, batches, metrics)
        +finishPartition(partitionId, batches, metrics)
        +finishAll(metrics)
    }

    class RssPartitionPusher {
        -partitionId
        -jvmObject
        +cloneWithPartitionId(partitionId)
        +write(bytes)
        +pushPartitionData(bytes)
    }

    class ShufflePartitionPusher {
        <<interface>>
        +pushPartitionData(partitionId, bytes, length)
    }

    class CometUniffleShuffleWriter {
        -nativePusherHandle
        +writeImpl(inputs)
        +pushPartitionData(partitionId, bytes, length)
        +stop(success)
    }

    class UniffleBufferManager {
        +addPartitionData(partitionId, bytes, length)
        +clear(ratio)
    }

    class RssShuffleWriter {
        +processShuffleBlockInfos(blocks)
        +checkDataIfAnyFailure()
        +sendCommit()
    }

    PartitionWriter <|.. RssPartitionWriter
    RssPartitionWriter *-- RssPartitionPusher
    RssPartitionPusher ..> ShufflePartitionPusher : JNI callback
    ShufflePartitionPusher <|.. CometUniffleShuffleWriter
    RssShuffleWriter <|-- CometUniffleShuffleWriter
    CometUniffleShuffleWriter --> UniffleBufferManager

Loading

The write path is:

sequenceDiagram
    participant Spark as Spark task
    participant Writer as CometUniffleShuffleWriter
    participant Native as Native ShuffleWriterExec
    participant RPW as RssPartitionWriter
    participant Pusher as RssPartitionPusher
    participant JVM as ShufflePartitionPusher
    participant Uniffle as Uniffle client

    Spark->>Writer: writeImpl(inputs)
    Writer->>Native: nativeWrite(inputs, pusherHandle)
    Native->>RPW: create one buffered writer per partition
    Native->>RPW: write(partitionId, RecordBatch)
    RPW->>RPW: encode and buffer Comet shuffle blocks
    RPW->>Pusher: write(encodedBytes)
    Pusher->>JVM: pushPartitionData(partitionId, bytes, length)
    JVM->>Uniffle: addPartitionData()
    JVM->>Uniffle: processShuffleBlockInfos()
    JVM->>Uniffle: checkDataIfAnyFailure()
    Writer->>Uniffle: flush remaining blocks and wait
    Writer->>Uniffle: commit

Loading

Shuffle reader

CometUniffleShuffleReader should implement two read paths backed by the same Uniffle block iterator.

read()

read() provides the regular Spark ShuffleReader API:

1. Create an Uniffle ShuffleReadClient for each requested reducer partition.
2. Fetch remote ShuffleBlock instances.
3. Reassemble a complete Comet shuffle block, even when it spans multiple Uniffle blocks.
4. Parse the 16-byte Comet header containing the compressed length and field count.
5. Copy the compressed body into a reusable direct ByteBuffer.
6. Call Native.decodeShuffleBlock to produce a ColumnarBatch.
7. Update Spark shuffle-read metrics and return an interruptible iterator.

readAsShuffleBlockIterator()

readAsShuffleBlockIterator() supports Comet native shuffle scan.

Instead of decoding the block into a JVM ColumnarBatch, it returns a CometShuffleBlockIterator. Native scan pulls compressed blocks through:

• hasNext()
• getBuffer()
• getCurrentBlockLength()
• close()

This avoids an unnecessary JVM decode and allows the compressed Comet shuffle block to be consumed directly by native execution.

sequenceDiagram
    participant Consumer as Spark or Comet native scan
    participant Reader as CometUniffleShuffleReader
    participant Iterator as CometUniffleShuffleBlockIterator
    participant Client as Uniffle ShuffleReadClient
    participant Decode as Native.decodeShuffleBlock

    alt Regular Spark read
        Consumer->>Reader: read()
        Reader->>Iterator: create iterator
        loop Requested reducer partitions
            Iterator->>Client: readShuffleBlockData()
            Client-->>Iterator: ShuffleBlock ByteBuffer
            Iterator->>Iterator: parse header and assemble compressed body
        end
        Reader->>Decode: decodeShuffleBlock(buffer, length, fieldCount)
        Decode-->>Reader: ColumnarBatch
        Reader-->>Consumer: Iterator of ColumnarBatch
    else Comet native scan
        Consumer->>Reader: readAsShuffleBlockIterator()
        Reader-->>Consumer: CometShuffleBlockIterator
        loop Pull compressed blocks
            Consumer->>Iterator: hasNext()
            Iterator->>Client: readShuffleBlockData()
            Iterator-->>Consumer: block length
            Consumer->>Iterator: getBuffer()
            Iterator-->>Consumer: direct ByteBuffer
        end
    end

Loading

Additional context

No response