andygrove opened a new pull request, #6205: URL: https://github.com/apache/datafusion-comet/pull/6205
## Which issue does this PR close? Closes #5961. ## Rationale for this change `fair_unified` is the default off-heap pool. Since the DataFusion 53 upgrade (#3629, shipped in 1.0.0), its `try_grow` has compared the pool's total reserved bytes against `pool_size / num_consumers`. That caps the whole task at one consumer's share rather than capping each consumer at its own. With three consumers registered the task can use a third of its budget, and every new registration tightens the ceiling on the consumers already running. The tuning guide documents the per-consumer limit. The comment that justified the change holds for `shrink` but not for `try_grow`. `MemoryReservation::try_grow` calls `pool.try_grow` before adding to its size in DataFusion 52.1, 53.0 and 55.1 alike, so `reservation.size()` inside the pool is the reservation's size before the request. @peterxcli worked this out in #5466. This PR takes the same approach and adds the pool-total check that @sunchao asked for in the review there. This is the small version of the fix, so it can go into 1.1.0 and back to `branch-1.0`. #5847 goes further and charges sibling reservations against their consumer's share. It can build on this. ## What changes are included in this PR? `CometFairMemoryPool::try_grow` now refuses a request, without calling Spark, if either of these holds: - The requesting reservation's size plus the request exceeds `pool_size / num_consumers`. This is the check from before DataFusion 53, and the one upstream `FairSpillPool` makes for spillable consumers. - The pool's total plus the request exceeds `pool_size`. The shares alone do not bound the total. A consumer keeps what it reserved before others registered, and sibling reservations from `new_empty()` or `split()` are each checked against the same share on their own. This never refuses a request that main accepts. A reservation's size is part of the pool's total, so `used + additional <= pool_size / num` implies both new conditions. It can therefore only reduce spilling and let a task reserve more, up to `pool_size` and whatever Spark grants the task. That is worth a line in the release notes, since deployments may have sized their off-heap memory against the accidental cap since 1.0.0. `shrink` is unchanged. DataFusion 53+ decrements the reservation's size before it calls `pool.shrink`, so `shrink` keeps working from the pool's total. `grow` is also unchanged, and ignores both limits because it records memory that already exists. The tuning guide, the memory management guide and the `review-comet-memory-pr` skill described the pool-wide comparison, and now describe both checks. The tuning guide also said `num_reservations`, but the divisor is the number of registered consumers. ## How are these changes tested? Two new unit tests in `fair_pool.rs` go through DataFusion's `MemoryReservation` and the pool's real admission path. The existing `FakeSpark` grants everything in these tests, so only the pool's own checks can refuse. - `each_consumer_is_limited_to_its_own_share` is the reproduction from the issue. With a 100-byte pool and two consumers, the first holds 40 bytes and the second can still grow by 20. Each then stops at 50. - `a_consumer_registered_late_is_limited_by_the_pool_total` has a consumer reserve 60 of 90 bytes on its own before a second one registers. The first is then held at its new 45-byte share. The pool total stops the second at 30 bytes, below its 45-byte share, until the first releases memory. Putting main's comparison back fails both tests, and deleting the pool-total check fails the second. ``` cargo test -p datafusion-comet --lib memory_pools # 17 passed cargo clippy -p datafusion-comet --all-targets -- -D warnings ``` The Comet JVM suites run this pool, because `CometTestBase` sets `spark.memory.offHeap.enabled=true` with 2 GiB, but their data rarely reaches the limit. An off-heap run comparing spilling against main is still to do, and I'll post the results here. -- 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]
