wForget opened a new issue, #4913:
URL: https://github.com/apache/datafusion-comet/issues/4913
### 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.
```mermaid
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
```
The write path is:
```mermaid
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
```
## 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.
```mermaid
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
```
### Additional context
_No response_
--
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]
