andygrove opened a new pull request, #5387: URL: https://github.com/apache/datafusion-comet/pull/5387
## Which issue does this PR close? Part of #5383, following the measurements in https://github.com/apache/datafusion-comet/issues/5383#issuecomment-5319082686. Independent of #5384 (they touch different crates and can merge in either order). ## Rationale for this change `buffer_partitioned_batch_may_spill` reserved exactly what each batch newly pinned. For an 8192-row `i32` column that is 32 KB, and against Comet's unified memory pool every `try_grow` is a JNI round-trip into Spark's memory manager plus contention on its executor-wide lock. Measuring a 201-partition shuffle of 10M rows with the counters from #5384, `ShuffleRepartitioner[0]` was the only consumer touching the pool, and it made **256 acquisitions per task to reserve 24 MB — 247 of them for exactly 32 KB** — followed by a single release of the whole amount. Nothing about that pattern needs a round-trip per batch. ## What changes are included in this PR? The repartitioner now tracks the exact resident figure itself in `buffered_bytes` and grows `reservation` in 1 MB steps on top of it, so a run of small batches costs one acquisition instead of one each. Two details keep the spill behaviour where it was: - `max_buffer_bytes` is compared against `buffered_bytes` rather than `reservation.size()`. The reservation is now rounded up, so comparing against it would trip the limit early and spill sooner than before. - if the pool denies a rounded-up request, it is retried at the exact deficit before spilling. A batch that fit before still fits; the extra call happens only on the path that was about to spill anyway. `used()` now reports `buffered_bytes`, which is the same number it reported before this change (the exact resident bytes), so the spill log is unaffected. The cost is up to one step of reservation held beyond what a writer is using, per active writer — 1 MB per concurrent shuffle-write task. I picked 1 MB rather than something larger to keep that bounded; it already removes 95% of the calls, and the step is capped at `max_buffer_bytes / 8` so a writer with a small limit cannot reserve a multiple of its own budget. ## How are these changes tested? New unit test `small_batches_share_one_reservation_step`, over a `MemoryPool` that counts the `try_grow` calls reaching it: 40 distinct batches of 100 rows produce **1** acquisition, where growing by the exact amount produces **40** (verified by setting the step to 1 byte). The batches are built separately rather than cloned, because `count_new_buffers` dedups by buffer address and re-inserting one batch would charge nothing after the first insert. Existing coverage: `cargo test -p datafusion-comet-shuffle --lib` (32 tests) passes, including `max_buffer_bytes_triggers_spill_without_memory_pressure`, `max_buffer_bytes_none_leaves_spilling_to_memory_pressure` and `max_buffer_bytes_preserves_output`, which are the tests guarding the limit semantics this touches. `cargo test -p datafusion-comet --lib` (163) and `CometNativeShuffleSuite` (27 tests) also pass. `cargo clippy --all-targets` and `cargo fmt --check` are clean. ## Measured effect Same benchmark case as the issue comment (`SQL Single INT Shuffle(201 Partition)`, 10M rows, native shuffle, release build, 4 GB off-heap, default `fair_unified`, `local[5]`), with the #5384 counters enabled: | Per task | Before | After | |---|---|---| | Acquire calls | 256 | **13** | | Release calls | 1 | 1 | | Time inside the JNI call | 327 µs | **50 µs** | | Bytes reserved | 25,165,824 | 25,198,592 | End-to-end wall time is **not** a useful signal here and I do not want to claim one: the saving is ~0.3 ms out of a ~91 ms task, and between the two runs every case moved by 5-10% including the pure-Spark baseline that this change cannot affect (620 -> 659 ms). The case for the change rests on the call count and the JNI time, both measured directly, plus the reduced traffic on Spark's shared memory-manager lock at higher concurrency than a local benchmark produces. -- 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]
