andygrove opened a new issue, #7076: URL: https://github.com/apache/datafusion-comet/issues/7076
### What is the problem the feature request solves? The native Celeborn writer reserves about 7x the uncompressed batch size for a frame and keeps the whole reservation through the network round trip, long after the native and JVM copies are gone. On an executor running several tasks this limits how many frames can be in flight. In `RssPartitionWriter::push_batch_within_limit` (`native/shuffle/src/writers/rss/rss_partition_writer.rs`) the reservation for a frame is `frame_bound x 3` (native output buffer, its JNI `byte[]`, and Celeborn's copied request) plus an encoding scratch of `4 x original_size` plus metadata, compaction and codec workspace. When the push returns, the native `Vec` is dropped and the JNI array is unreachable, but `CelebornShufflePartitionPusher.releasableBytes` returns the full `reservation.bytes` only once both `nativeReleased` and `transportComplete` are true. With the 512 MiB default, about 9 frames of 8 MiB batches can be in flight per executor. ### Describe the potential solution On native release, shrink the reservation to what Celeborn's request still holds (the frame length plus the 16-byte batch header) instead of keeping the full amount until transport completion. Charge the encode scratch to the DataFusion memory pool rather than to the push admission budget, so it is freed when encoding ends. Test with the existing `CelebornShufflePartitionPusherSuite` ownership tests, plus a multi-task case that measures peak admitted bytes and frames in flight. ### Additional context Finding C1 in #5905. Related: #6654 (copies in native Celeborn shuffle writes), #6196 and #6336 (charging the local writer's buffers to the memory pool). -- 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]
