jayzhan211 commented on code in PR #24785:
URL: https://github.com/apache/datafusion/pull/24785#discussion_r3903898957
##########
datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs:
##########
@@ -219,46 +255,126 @@ impl PartialReduceHashAggregateStream {
timer.done();
if let Err(e) = result {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- original_state,
- ));
+ return Self::break_with_err(e);
}
- if let Err(e) = self
- .reservation
- .try_resize(original_state.hash_table().memory_size())
- {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- original_state,
- ));
- }
-
- ControlFlow::Continue(original_state)
- }
- Poll::Ready(Some(Err(e))) => {
- ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state))
+ // Update the memory reservation. If OOM, do early emit.
+ self.resize_or_emit_early(original_state)
}
+ Poll::Ready(Some(Err(e))) => Self::break_with_err(e),
// Input ends, move to output state
Poll::Ready(None) => {
+ self.close_input();
let elapsed_compute =
self.baseline_metrics.elapsed_compute().clone();
let timer = elapsed_compute.timer();
- let result =
self.start_output(original_state.hash_table_mut());
+ let result = original_state.hash_table_mut().start_output();
timer.done();
match result {
Ok(()) => {
ControlFlow::Continue(original_state.into_producing_output())
}
- Err(e) => {
- ControlFlow::Break((Poll::Ready(Some(Err(e))),
original_state))
- }
+ Err(e) => Self::break_with_err(e),
}
}
}
}
+ /// Update the memory reservation. If the reservation succeeds, continue
reading
+ /// input. If OOM, clear the aggregated states in the hash table, and
early emit
+ /// them immediately.
+ ///
+ /// Returns the next state; the caller finishes the intended task based on
it.
+ ///
+ /// The reservation is left at its pre-emission size while the states are
being
+ /// emitted, because the cleared states are still held in memory as
+ /// `remaining_groups`. [`Self::handle_emitting_on_memory_pressure`]
updates the
+ /// reservation once the last slice has been emitted.
+ ///
+ /// # Implementation Note
+ /// All accumulated states are materialized at once, and then sliced into
+ /// `batch_size` output batches. Emit them incrementally after blocked
state
+ /// management is ready.
+ ///
+ /// Issue: <https://github.com/apache/datafusion/issues/7065>
+ fn resize_or_emit_early(
+ &mut self,
+ mut original_state: PartialReduceHashAggregateState,
+ ) -> PartialReduceHashAggregateStateTransition {
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+ let _timer = elapsed_compute.timer(); // Stop on drop
+ let resize_result = self
+ .reservation
+ .try_resize(original_state.hash_table().memory_size());
+
+ let oom = match resize_result {
+ Ok(()) => return ControlFlow::Continue(original_state),
+ Err(e @ DataFusionError::ResourcesExhausted(_)) => e,
+ Err(e) => return Self::break_with_err(e),
+ };
+
+ let state_batch_result =
original_state.hash_table_mut().take_state_batch();
+
+ match state_batch_result {
+ Ok(Some(remaining_groups)) => ControlFlow::Continue(
+ PartialReduceHashAggregateState::EmittingOnMemoryPressure {
+ hash_table: original_state.into_hash_table(),
+ remaining_groups,
+ },
+ ),
+ // No accumulated group to emit, so early emission cannot release
any
+ // memory: report the original error.
+ Ok(None) => Self::break_with_err(oom),
+ Err(e) => Self::break_with_err(e),
+ }
+ }
+
+ /// Handle EmittingOnMemoryPressure state - emit a materialized
partial-state
+ /// batch in `batch_size`(from configuration) slices. After all slices are
+ /// emitted, update the memory reservation and resume reading input.
+ ///
+ /// See comments at `poll_next()` for details.
+ ///
+ /// Returns the next operator state with control flow decision.
+ fn handle_emitting_on_memory_pressure(
+ &mut self,
+ original_state: PartialReduceHashAggregateState,
+ ) -> PartialReduceHashAggregateStateTransition {
+ let PartialReduceHashAggregateState::EmittingOnMemoryPressure {
+ hash_table,
+ remaining_groups: batch,
+ } = original_state
+ else {
+ unreachable!("expected the EmittingOnMemoryPressure state")
+ };
+
+ let (output_batch, next_state) = if batch.num_rows() <=
self.batch_size {
+ // Go back to `ReadingInput`
Review Comment:
```suggestion
// Last slice: the materialized states are now owned downstream,
so
// the reservation can drop back to the (empty) table's size.
let _ = self.reservation.resize(hash_table.memory_size());
// Go back to `ReadingInput`
```
--
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]
