jayzhan211 commented on code in PR #25877:
URL: https://github.com/apache/datafusion/pull/25877#discussion_r4184132499
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -316,20 +316,44 @@ impl PartialHashAggregateStream {
break;
}
HandleInputResult::OOM => {
- let materialized_group_states =
hash_table.take_state_batch()?.ok_or_else(|| {
- internal_datafusion_err!(
+ let materialized_group_states =
+ hash_table.take_state_batches()?;
+ if materialized_group_states.is_empty() {
+ return Err(internal_datafusion_err!(
"Partial hash aggregate ran out of memory with
no aggregated groups"
- )
- })?;
+ ));
+ }
self.early_emit_count.add(1);
timer.done();
- self.emit_on_memory_pressure(
- materialized_group_states,
- &mut emitter,
- hash_table.memory_size(),
- )
- .await?;
+
+ // Blocked storage returns one batch per block, moved
out of
+ // the table without copying. Emit them in turn,
keeping the
+ // batches not emitted yet in the reservation.
+ let pending_memory: usize = materialized_group_states
+ .iter()
+ // Don't include the first batch since we will
emit it right away and release its memory
+ // if it needs slicing then a we will try to hold
on that reservation while slicing
+ .skip(1)
+ .map(|b| b.memory_size)
+ .sum();
+
+ // Make sure we can hold on the hash tables and all
the batches that need to be emitted (except the first one)
+ // if we can't hold it than we can't do anything about
it.
+ self.reservation.try_resize(
Review Comment:
Agreed that the held blocks must be counted, and my earlier suggestion (emit
without reserving) didn't do that. But the error doesn't count them either: at
this point they were never reserved. Logged in the failing tests:
`aggregate_grouping_sets_with_yielding_with_spill` has 0 B reserved while the
blocks hold 872 B (pool 500 B), and `..._does_not_fit` has 3.5 KB reserved for
4 × 12.5 KB blocks, so `?` frees nothing and fails a query main finishes.
`MemoryReservation::grow` is infallible, so you can count them honestly, over
the limit, until they drain:
```diff
- self.reservation.try_resize(
- hash_table.memory_size()
- // The batches that need to be held while
emitting each one
- + pending_memory,
- )?;
+ let held = hash_table.memory_size() +
pending_memory;
+ match self.reservation.try_resize(held) {
+ Ok(()) => {}
+ // The blocks already exist, so failing would
not free
+ // them: count them anyway, over the pool
limit, until
+ // they are emitted
+ Err(DataFusionError::ResourcesExhausted(_)) => {
+ self.reservation.grow(held -
self.reservation.size())
+ }
+ Err(e) => return Err(e),
+ }
```
With this both `aggregate_grouping_sets_*_with_spill` pass. The two
`test_partial_hash_stream_*` failures assume one flat state batch (as their own
comment anticipates): a `Utf8` key + `sum` in
`partial_stream_under_memory_limit` keeps the slicing coverage, plus an Int32 +
`count` variant asserting no error at 32 KB.
--
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]
