Weijun-H commented on code in PR #24852:
URL: https://github.com/apache/datafusion/pull/24852#discussion_r3967349927
##########
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:
Yes, that is the case I meant. Emitting the whole batch when the second
reservation fails sounds reasonable, provided larger output batches are
supported. We should still account for the remaining hash table and add a test
that exercises this fallback.
--
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]
