unikdahal opened a new issue, #6654:
URL: https://github.com/apache/datafusion-comet/issues/6654
## Problem
The native Comet -> Celeborn shuffle write path currently appears to keep
**three frame-sized payload representations** alive and performs **two full
payload copies** after the shuffle frame has already been encoded.
Current flow:
1. **Comet native frame**
- Encoded frame lives in native memory.
2. **Native -> JVM copy**
- `shuffle_partition_pusher.rs` uses `byte_array_from_slice(...)`.
- This copies the full native frame into a JVM `byte[]`.
3. **JVM -> Celeborn request copy**
- `ShuffleClientImpl.pushOrMergeData(...)` allocates a new `byte[]`
containing the 16-byte Celeborn batch header + payload.
- The complete payload is copied again using `System.arraycopy(...)`.
4. **Celeborn -> Netty**
- `NettyManagedBuffer(Unpooled.wrappedBuffer(body))` appears to wrap the
existing buffer rather than performing another full payload copy.
So the current path is roughly:
**Native frame -> JNI byte[] -> Celeborn framed byte[] -> Netty**
That means approximately:
- **3 overlapping frame-sized representations**
- **2 full payload copies**
This also seems consistent with the current RSS reservation logic in Comet,
which reserves for the native output, JNI array, and Celeborn transport request
overlapping.
## Idea
Explore whether we can reduce these copies incrementally.
Potential progression:
- Current: **3 representations / 2 copies**
- Intermediate: **2 representations / 1 copy**
- Target: **~1 payload representation / 0 redundant copies after encoding**
Some directions worth exploring:
- Avoid Celeborn's `System.arraycopy(...)` by reserving the Celeborn header
space in the existing JVM buffer.
- Add/use a caller-owned `ByteBuffer` or `ByteBuf` push path in Celeborn.
- Explore passing a native-backed direct buffer across JNI.
- Use a small Celeborn header + native payload as a composite buffer instead
of materializing them into another array.
Celeborn's Flink client already has a `ByteBuf`-based path that writes the
batch header into an existing buffer, so there seems to be precedent for
avoiding the additional framing copy.
## Main challenge
The difficult part of a true zero-copy path is likely **ownership/lifetime**.
Celeborn pushes are asynchronous, and the payload may still be referenced by
Netty or Celeborn retry/revive logic after the original JNI call returns.
Any native-backed buffer therefore needs to remain alive until all
transport/retry owners are finished.
Comet's existing `CelebornTransportCallbackTracker` and push reservation
lifecycle may already provide useful machinery for this.
## Things to investigate
- Can we first eliminate only the Celeborn `System.arraycopy(...)`?
- What is the cleanest caller-owned buffer API for Celeborn?
- Can a direct `ByteBuffer` / `ByteBuf` safely reference the native Comet
frame?
- Can the existing callback tracker retain/release the native frame across
retries?
- How should CRC, encryption, and fallback paths behave?
- What is the actual impact on allocation, RSS, CPU, and shuffle throughput?
A small benchmark comparing the **current**, **one-copy**, and
**direct/zero-copy** paths should help determine whether pursuing the full
optimization is worthwhile.
## Relevant code
Comet:
- `native/jni-bridge/src/shuffle_partition_pusher.rs`
- `native/shuffle/src/writers/rss/rss_partition_writer.rs`
-
`spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java`
-
`spark/src/main/java/org/apache/comet/shuffle/CelebornTransportCallbackTracker.java`
Celeborn:
- `client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java`
-
`common/src/main/java/org/apache/celeborn/common/network/buffer/NettyManagedBuffer.java`
-
`client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/FlinkShuffleClientImpl.java`
This issue is mainly for **exploration and discussion** before deciding on a
concrete implementation.
--
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]
