[PR] fix: CometFairMemoryPool compared pool-wide usage against a per-consumer limit [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: CometFairMemoryPool compared pool-wide usage against a per-consumer limit [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove commented on PR #5214:
URL: https://github.com/apache/datafusion-comet/pull/5214#issuecomment-5159170174

   Follow-up to my review above: rather than reason about how tight the old cap was, I measured it. I temporarily instrumented `CometFairMemoryPool::register` to print the consumer count per task pool and ran two Comet queries with `spark.memory.offHeap.enabled=true`, `spark.memory.offHeap.size=512m`, `fair_unified`, on `local[5]` (instrumentation reverted afterwards).
   
   Consumers registering per task pool:
   
   | Plan | Peak `num` | Consumers |
   |---|---|---|
   | broadcast join + group-by + shuffle | 3 | `HashJoinInput[0]`, `GroupedHashAggregateStream[0]`, `ShuffleRepartitioner[0]` |
   | final sort stage | 2 | `ExternalSorter[0]`, `ExternalSorterMerge[0]` |
   | 2x SMJ + window + `countDistinct` | 7 | two join inputs, the multi-phase distinct-agg chain, `ShuffleRepartitioner[0]` |
   
   So `num` is 1-3 for ordinary stages and reaches ~7 on a wide one. Worth noting a single sort costs two consumers (`sorts/sort.rs:283,288`).
   
   That lets us say where the old cap actually binds. The old code caps a task's *total* native usage at `offHeap.size / num`, while Spark's `ExecutionMemoryPool` independently caps each active task at roughly `pool / N_active_tasks`. Whichever is smaller wins:
   
   | Scenario | old Comet cap | Spark's grant | binding constraint |
   |---|---|---|---|
   | 8 cores packed, num=3 | 33% of pool | ~12.5% | Spark - old cap invisible |
   | 8 cores packed, num=7 | 14% | ~12.5% | Spark, marginally |
   | 16 cores packed, num=7 | 14% | ~6% | Spark - old cap invisible |
   | 1-2 active tasks, num=7 | 14% | up to ~100% | **Comet, ~7x tighter** |
   | 2 cores, num=7 | 14% | 50% | Comet, ~3.5x tighter |
   
   The conclusion I'd draw: the defect is real, but on a busy executor Spark's own per-task fair share is tighter than the buggy cap, so it is masked. It bites when consumers-per-task exceeds active-tasks-per-executor - stage tails, AQE-coalesced partitions, low partition counts, small executors - where Spark would hand a lone task most of the pool and this check clamps it to `1/num`. The expected symptom is unnecessary spilling on long-tail tasks rather than a uniform throughput loss, which I think is a more defensible framing for the PR than "the default config can only use 10% of its off-heap share".
   
   Caveat on my own numbers: the two probe queries were small enough that the check never actually fired (zero `fair limit` errors in the run), so this measures `num` and the resulting arithmetic, not a performance delta.
   
   If you want a reproducer that demonstrates the fix rather than arguing for it, the low-concurrency case is the one to target: `local[1]` with `spark.sql.shuffle.partitions=1`, a modest `offHeap.size`, and a wide join + distinct query should hit "fair limit is M bytes, 7 registered" on main and get further after the fix. Happy to build that out if it would help the review.
   


-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org

Re: [PR] fix: CometFairMemoryPool compared pool-wide usage against a per-consumer limit [datafusion-comet]

["andygrove (via GitHub)" <gi...@apache.org>]

andygrove closed pull request #5214: fix: CometFairMemoryPool compared pool-wide usage against a per-consumer limit
URL: https://github.com/apache/datafusion-comet/pull/5214


-- 
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: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org