andygrove opened a new pull request, #5214: URL: https://github.com/apache/datafusion-comet/pull/5214
## Which issue does this PR close? Part of https://github.com/apache/datafusion-comet/issues/5212 (position 1). ## Rationale for this change `CometFairMemoryPool::try_grow` computed the per-consumer fair share `pool_size / num` but compared it against the **pool-wide** `state.used`: ```rust let limit = self.pool_size.checked_div(num).expect("overflow in checked_div"); let used = state.used; // pool-wide total if limit < used + additional { return resources_err!(...); } ``` So a task could never use more than a *single* consumer's share in total, and usable memory shrank linearly as consumers registered — a plan with 10 memory consumers could use 10% of its off-heap share before erroring or spilling. This is the **default off-heap pool type** (`fair_unified`, `CometConf.scala:686`), so it affects the default configuration whenever `spark.memory.offHeap.enabled=true`. The comment justifying the comparison was incorrect: > We use `state.used` instead of `reservation.size()` because DataFusion 53+ calls `pool.try_grow()` before incrementing the reservation's atomic size, so `reservation.size()` would not include prior grows. `MemoryReservation::try_grow` (datafusion-execution 54.1.0, `memory_pool/mod.rs:471-475`) is: ```rust pub fn try_grow(&self, capacity: usize) -> Result<()> { self.registration.pool.try_grow(self, capacity)?; self.size.fetch_add(capacity, atomic::Ordering::Relaxed); Ok(()) } ``` `reservation.size()` therefore *does* include every prior grow. It excludes only the current `additional`, which is exactly why upstream `FairSpillPool` compares `reservation.size() + additional > available` (`memory_pool/pool.rs:249`). The corresponding comment on `shrink` **is** correct — `MemoryReservation::shrink` decrements the atomic before calling the pool — and the two cases look to have been conflated during the DataFusion 53 upgrade in #3629. ## What changes are included in this PR? - Compare the fair share against `reservation.size()` rather than the pool-wide `state.used`, matching upstream `FairSpillPool` semantics. - Extract the decision into a pure `check_fair_share` function so the policy is unit testable without a JVM, and document that `reserved` must be the requesting consumer's own usage. - Replace `expect("overflow in checked_div")` with a fallback to the whole pool. `num` cannot be zero while a reservation exists (a consumer registers before it can grow), but the previous code would have panicked with a message describing the wrong failure. - Include the consumer name and the pool-wide total in the error message, which should make future reports of this class of failure easier to read. Behaviour is otherwise unchanged: the pool still tracks `state.used` for `reserved()` and the `shrink` guard, and the hard bound is still whatever Spark's `TaskMemoryManager` actually grants (the `acquired < additional` path below the check). ## How are these changes tested? Six new unit tests in `fair_pool.rs` covering the fair-share policy: a sole consumer may use the whole pool, the share divides evenly, each consumer may independently reach its own share (the regression guard), a consumer may not exceed its share, zero registered consumers does not divide by zero, and an oversized request does not overflow. ``` running 6 tests test execution::memory_pools::fair_pool::tests::each_consumer_may_reach_its_own_share ... ok test execution::memory_pools::fair_pool::tests::consumer_may_not_exceed_its_own_share ... ok test execution::memory_pools::fair_pool::tests::oversized_request_does_not_overflow ... ok test execution::memory_pools::fair_pool::tests::share_is_divided_evenly_between_consumers ... ok test execution::memory_pools::fair_pool::tests::sole_consumer_may_use_whole_pool ... ok test execution::memory_pools::fair_pool::tests::no_registered_consumers_does_not_panic ... ok ``` Full crate suite passes (141 passed, 4 ignored), and `cargo clippy --all-targets` is clean. **Testing gap worth flagging:** these tests cover the policy function, not the `try_grow` call site — that `reservation.size()` rather than `state.used` is passed in is verified by inspection only. Testing the pool end to end needs a JNI harness, since `CometFairMemoryPool::new` requires a live `CometTaskMemoryManager` global ref. That is tracked as the test-coverage item in #5212; I did not want to grow this PR into building that harness. I have **not** measured the performance effect on a real workload. The change strictly widens what the pool admits, so I would expect fewer spurious spills and `ResourcesExhausted` errors on off-heap plans with several memory consumers, but that is reasoning rather than a measurement. ## Are there any user-facing changes? Yes, on the default off-heap configuration: plans that previously spilled or failed with "Failed to acquire N bytes … the fair limit is M bytes" will now be able to use their intended fair share of the pool. No configuration changes. ## Follow-ups not in this PR `register`/`unregister` count *every* consumer, whereas upstream `FairSpillPool` divides by `num_spill` (spillable consumers only) and gives unspillable consumers the remainder. Comet's limit is therefore still tighter than upstream's. Aligning that means tracking spillable and unspillable usage separately, which is a larger behavioural change than this fix, so I left it out — noted as a secondary divergence under position 1 in #5212. -- 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]
