EmilyMatt commented on code in PR #25877:
URL: https://github.com/apache/datafusion/pull/25877#discussion_r4184594386
##########
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:
grow() should really not be used anywhere, you are suggesting adding a panic
here, which is much more severe than returning an error for the query - which
is also bad in my opinion, but preferable to exceeding the expected memory,
which can lead to SIGKILLs that are far harder to debug.
Erroring out if the memory is umavailable is the way forward in my opinion.
The next step is to ensure this situation never happens if we are able to
hold 2 batches in memory
--
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]
