unikdahal commented on issue #6654: URL: https://github.com/apache/datafusion-comet/issues/6654#issuecomment-6022806757
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 `maxInFlightBytes` until 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. | Path | Frame copies alive | Payload copies after encoding | |---|---|---| | Current | 3 | 2 | | New API with a JVM array | 2 | 1 | | New API from native memory | 1 | 0 | ### Benchmarks **Celeborn push path** (one real worker, 256 MiB pushed per run, median of 7 rounds): | Batch size | Metric | Current `pushData(byte[])` | Zero-copy | Change | |---|---|---|---|---| | 64 KiB | Memory allocated | 277 MiB | 24 MiB | −91% | | 64 KiB | CPU | 770 ms | 470 ms | −39% | | 1 MiB | Memory allocated | 259 MiB | 3.1 MiB | −98.8% | | 1 MiB | CPU | 430 ms | 270 ms | −37% | | 1 MiB | Wall time | 559 ms | 506 ms | −9.5% | | 4 MiB | Memory allocated | 258 MiB | 2.0 MiB | −99.2% | | 4 MiB | CPU | 460 ms | 250 ms | −46% | | 4 MiB | Wall time | 292 ms | 254 ms | −13% | 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): | Metric | JVM array copy | Zero-copy | Change | |---|---|---|---| | GC count | 89 | 69 | −22% | | GC time | 458 ms | 391 ms | −15% | | Java heap allocated per run | 44.7 GiB | 41.3 GiB | −3.3 GiB | | Executor CPU | 71.4 s | 68.9 s | −3.5% | | Shuffle write time (shuffle-heavy queries) | 5.19 s | 4.79 s | −8% | | Total wall time | 87.4 s | 86.7 s | within noise | 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. -- 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]
