comphead commented on code in PR #7096:
URL: https://github.com/apache/datafusion-comet/pull/7096#discussion_r4354915418
##########
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:
This test and
`a_task_shares_the_pool_with_the_other_tasks_running_native_plans` above have
the same body: a 100-byte pool, a second task, the sort refused at 50, then 50
more once the other task drops. If `other_task` here were created with the
`greedy_unified` config, one test would cover both the executor-wide count and
the mixed pool types, and the one above could go.
##########
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:
Thanks for the quick turn. Looking at `d79cb838b`, I think my suggestion
went one step too far. Taking the aggregate out of `num_spill` also lifts the
cap on the sort below it. In a `CometSortAggregateExec` over a native sort, the
sort reads all of its input before the aggregate's first request. At
`6a1c4dffca` the sort was held to half of what the merge leaves, so the
aggregate kept the other half. Now the sort can hold `task_share` less its
merge reservation. When it fits in memory it keeps holding that while it emits
(`in_mem_sort_stream` in DataFusion 55.2). The aggregate's headroom is then
about the `sort_spill_reservation_bytes` that `sort()` frees before emitting,
plus whatever Spark grants past the share. That is the #17334 shape that
DataFusion's flag works around.
Could the exception apply only where `try_grow` picks its branch, while
`register`, `unregister`, `grow`, `shrink` and the `unspillable` bookkeeping
keep `can_spill()`? The aggregate would still count in `num_spill`, so the sort
leaves it a share, but it would not be refused at that share. It would also be
one call site instead of five, and your new test passes either way.
A test that grows the sort first would pin it. With `pool(&fake, 100)`, the
proposed version refuses `sort.try_grow(51)`, and `aggregate.try_grow(50)`
succeeds after `sort.try_grow(50)`. At the head, `sort.try_grow(90)` succeeds
and the `pool_size` check then refuses `aggregate.try_grow(11)`. I worked these
through `try_grow` by hand and did not run them.
--
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]
