andygrove commented on code in PR #7096:
URL: https://github.com/apache/datafusion-comet/pull/7096#discussion_r4357393616
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -153,22 +219,32 @@ impl MemoryPool for CometFairMemoryPool {
) -> Result<(), DataFusionError> {
if additional > 0 {
let mut state = self.state.lock();
- let num = state.consumers.len();
- let limit = self
- .pool_size
- .checked_div(num)
- .expect("overflow in checked_div");
- let consumer_used = *state.consumer_used(reservation);
- if limit < consumer_used.saturating_add(additional) {
- return resources_err!(
- "Failed to acquire {additional} bytes where this consumer
already holds {consumer_used} bytes and the fair limit is {limit} bytes, {num}
registered ({} bytes overcommitted)",
- self.spark.overcommit()
- );
- }
- // The shares alone do not bound the pool's total, because a
consumer keeps what it
- // reserved before another consumer registered.
+ let can_spill = reservation.consumer().can_spill();
let used = state.used;
- if self.pool_size < used.saturating_add(additional) {
+ if can_spill {
Review Comment:
Agreed, the exception went one step too far. In 6303b8e39 it applies only
where `try_grow` picks its check (`limited_to_its_share`). `register`,
`unregister`, `grow`, `shrink` and the `unspillable` bookkeeping go back to
`can_spill()`, so the aggregate counts in `num_spill`, its bytes are not
`unspillable`, and a sort below it stays at half of the task's share. I added
your sort-first case as
`a_sort_below_an_ordered_final_aggregate_leaves_it_a_share`: with `pool(&fake,
100)` the sort is refused at 51, then the sort's 50 and the aggregate's 50 both
succeed. It fails with the bookkeeping at d79cb838b.
`an_ordered_final_aggregate_is_limited_as_a_consumer_that_cannot_spill` passes
either way, as you said. `CometAggregateSuite` passes on 4.1.
##########
native/core/src/execution/memory_pools/mod.rs:
##########
@@ -145,6 +179,71 @@ mod tests {
}
}
+ #[test]
+ fn a_task_shares_the_pool_with_the_other_tasks_running_native_plans() {
+ let tasks = RunningTasks::default();
+ let fake = spark_memory::fake::FakeSpark::with(usize::MAX);
+ let fair = MemoryPoolConfig::new(MemoryPoolType::FairUnified, 100);
+ let greedy = MemoryPoolConfig::new(MemoryPoolType::GreedyUnified, 0);
+ let pool = create_pool(&fair, -3004, &tasks, || fake.memory());
+ let other_task = create_pool(&greedy, -3005, &tasks, || fake.memory());
+ let sort = MemoryConsumer::new("ExternalSorter[0]")
+ .with_can_spill(true)
+ .register(&pool);
+
+ // A task with the other pool type counts too.
+ sort.try_grow(50).unwrap();
+ assert!(sort.try_grow(1).is_err());
+
+ // Once it finishes, the task has the pool to itself.
+ drop(other_task);
+ sort.try_grow(50).unwrap();
+ assert!(sort.try_grow(1).is_err());
+ }
+
+ /// The tasks of an executor share one count, which `create_memory_pool`
takes from
+ /// [`RunningTasks::of_executor`]. A pool created with a count of its own
would see only its
+ /// own task and give it the whole pool. No other test counts tasks in the
executor's count.
+ #[test]
+ fn the_tasks_of_an_executor_share_one_count() {
Review Comment:
Done in 6303b8e39. `the_tasks_of_an_executor_share_one_count` now creates
the other task with the `greedy_unified` config, so it covers both the
executor-wide count and the mixed pool types, and
`a_task_shares_the_pool_with_the_other_tasks_running_native_plans` is gone.
--
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]
