rluvaton commented on code in PR #24852:
URL: https://github.com/apache/datafusion/pull/24852#discussion_r3967235649
##########
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 will fail only if state batch have more rows than batch size (since
otherwise we don't count it due it being emitted right away) and the state
batch size takes as much or more memory as it did unmaterialized.
what I can do is that if it fail to reserve for that as well, we can just
emit the large batch to avoid failure
--
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]
