pepijnve commented on code in PR #19287:
URL: https://github.com/apache/datafusion/pull/19287#discussion_r2630322095
##########
datafusion/physical-plan/src/aggregates/row_hash.rs:
##########
@@ -1060,18 +1084,26 @@ impl GroupedHashAggregateStream {
Ok(Some(batch))
}
+ /// Determines if `spill_state_if_oom` can free up memory by spilling
state to disk
+ fn can_spill_on_oom(&self) -> bool {
+ match self.oom_mode {
+ // For spill mode, only spill if we are not already reading back
spilled state
+ OutOfMemoryMode::Spill => {
+ !self.group_values.is_empty() &&
!self.spill_state.is_stream_merging
+ }
+ // For emit early mode, never spill
+ OutOfMemoryMode::EmitEarly => false,
+ }
+ }
+
/// Optimistically, [`Self::group_aggregate_batch`] allows to exceed the
memory target slightly
/// (~ 1 [`RecordBatch`]) for simplicity. In such cases, spill the data to
disk and clear the
/// memory. Currently only [`GroupOrdering::None`] is supported for
spilling.
- fn spill_previous_if_necessary(&mut self, batch: &RecordBatch) ->
Result<()> {
- // TODO: support group_ordering for spilling
- if !self.group_values.is_empty()
+ fn spill_state_if_oom(&mut self, batch: &RecordBatch) -> Result<()> {
+ if self.can_spill_on_oom()
&& batch.num_rows() > 0
- && matches!(self.group_ordering, GroupOrdering::None)
- && !self.spill_state.is_stream_merging
&& self.update_memory_reservation().is_err()
Review Comment:
My second attempt eliminates all these somewhat odd conditional methods. The
flow is much more linear now and the interdependencies between the methods have
been removed.
--
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]
