Repository navigation
perf: explore reducing copies in native Celeborn shuffle writes #6654
Description
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:
-
Comet native frame
- Encoded frame lives in native memory.
-
Native -> JVM copy
shuffle_partition_pusher.rsusesbyte_array_from_slice(...).- This copies the full native frame into a JVM
byte[].
-
JVM -> Celeborn request copy
ShuffleClientImpl.pushOrMergeData(...)allocates a newbyte[]containing the 16-byte Celeborn batch header + payload.- The complete payload is copied again using
System.arraycopy(...).
-
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
ByteBufferorByteBufpush 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/ByteBufsafely 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.rsnative/shuffle/src/writers/rss/rss_partition_writer.rsspark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.javaspark/src/main/java/org/apache/comet/shuffle/CelebornTransportCallbackTracker.java
Celeborn:
client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.javacommon/src/main/java/org/apache/celeborn/common/network/buffer/NettyManagedBuffer.javaclient-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.
Activity
unikdahal commentedon Oct 5, 2026
@pingzh would love to get your thoughts on this. Do you think reducing these copies / exploring a direct-buffer path is worth pursuing, or is there anything in the current retry/ownership design that makes this impractical?
pingzh commentedon Oct 5, 2026
@unikdahal thanks for opening this ticket. it is definitely worth pursuing. Internally, we did a quick win in our celeborn fork, pushOrMergeData to avoid one copy, but i did not put much thoughts in the Netty side yet.
unikdahal commentedon Oct 6, 2026
I spent some time prototyping this end to end, and the full zero-copy path works. The part that makes it practical is having Celeborn own the lifetime question instead of Comet.
Approach
On the Celeborn side, a small client-only API pushes a caller-owned
ByteBuffer(heap or direct) without copying it, and runs a release callback exactly once when the client is completely done with the buffer: after the batch is acknowledged, fails, or is dropped, and after every network write and retry of it has finished. Internally each batch is reference-counted, so retries and in-flight writes keep it alive. The wire format and servers stay unchanged.On the Comet side, the native writer hands each encoded frame to the JVM as a direct buffer over native memory and transfers ownership of it. Celeborn sends straight from native memory, and the release callback frees the frame. The frame stays counted against
maxInFlightBytesuntil it is released, so memory stays bounded. Since there is only one copy of the frame, the writer reserves one frame instead of three, and it trims each frame to its compressed size before handing it over.Benchmarks
Celeborn push path (one real worker, 256 MiB pushed per run, median of 7 rounds):
pushData(byte[])At 64 KiB, wall time is the same either way because the worker is the bottleneck.
Comet end to end (Spark 3.5, Celeborn 0.7.0 server, TPC-H SF1 plus two shuffle-heavy queries, median of 5 runs). Every run produced the same results as vanilla Spark.
Removing the copies (same API, JVM array vs. native memory):
For context, compared with what Celeborn users get today (Comet with the released client, which falls back to Celeborn's regular Spark shuffle), native shuffle through the new API cut total wall time from 132.9 s to 86.7 s (−35%) and executor CPU by 58%.
One caveat on memory: peak RSS goes up with native shuffle. That comes from Comet's native shuffle buffering into the off-heap pool, not from the push path. Native shuffle to local disk, with no Celeborn involved, peaks almost as high, and the copy and zero-copy variants peak at about the same level.
I'll try to get the API upstreamed in Celeborn first, since the Comet change depends on it, and then open the Comet PR.