Weijun-H commented on code in PR #24852:
URL: https://github.com/apache/datafusion/pull/24852#discussion_r3912934795
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -652,10 +655,19 @@ impl PartialHashAggregateStream {
let _timer = elapsed_compute.timer();
let state_batch_result = hash_table.take_state_batch();
+ // If we are holding on the memory due to slicing
account for that
+ let state_batch_size = match &state_batch_result {
+ Ok(Some(batch)) if batch.num_rows() >
self.batch_size => {
+ batch.get_array_memory_size()
+ }
+ _ => 0,
+ };
+
// Emitting clears the aggregate table and releases its
// accumulated memory. Update the reservation
accordingly.
- let resize_result =
-
self.reservation.try_resize(hash_table.memory_size());
+ let resize_result = self
+ .reservation
+ .try_resize(hash_table.memory_size() +
state_batch_size);
Review Comment:
This follows a failed resize, so charging the retained batch can make early
emission fail with OOM after the batch is already allocated. If it cannot be
retained, preserve progress (for example, yield it whole).
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -675,6 +687,7 @@ impl PartialHashAggregateStream {
PartialHashAggregateState::EmittingOnMemoryPressure {
hash_table,
remaining_groups: materialized_group_states,
+ batch_memory_size: state_batch_size,
Review Comment:
`MaterializedAggregateOutput` already implements this
materialize-once/slice-many lifecycle. Reuse it with a split child
`MemoryReservation`, and test that the reservation drops after the last slice
before polling to `Done`.
--
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]
