You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/02 01:18:36 UTC
(datafusion) branch main updated: refactor(hash-aggr): Let partial-reduce mode aggregation early emit when OOM (#24785)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 1608ee70fa refactor(hash-aggr): Let partial-reduce mode aggregation early emit when OOM (#24785)
1608ee70fa is described below
commit 1608ee70faa06063a38f01e8c57c1a6879708cd6
Author: Yongting You <20...@gmail.com>
AuthorDate: Wed Sep 2 01:08:05 2026 +0000
refactor(hash-aggr): Let partial-reduce mode aggregation early emit when OOM (#24785)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
Part of https://github.com/apache/datafusion/issues/22710
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
Please explain the problem you are trying to solve in terms of the
user-visible
behavior, rather than the implementation.
For example, "The code in `foo.rs` doesn't handle nulls" is a symptom of
the
implementation. "COUNT(DISTINCT) returns wrong results when the column
contains
nulls" is the user-visible problem.
-->
See https://github.com/apache/datafusion/pull/24486 for the rationale.
That discussion covers the issue and three potential solutions: error,
early-emit, and spill.
This PR implements the early-emit behavior agreed on there.
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here, but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
Key changes:
In datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs
- top comment explains the high-level ideas
- To understand the implementation, start from `poll_next()` and follow
along
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
UTs (note partial-reduce aggregation can't be planned from SQL, so we
can't do sqllogictests here)
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
No
---
datafusion/physical-plan/src/aggregates/mod.rs | 185 ++++++++++++++--
.../src/aggregates/partial_reduce_stream.rs | 235 +++++++++++++++++----
2 files changed, 359 insertions(+), 61 deletions(-)
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index f00f1a160e..6da7ee1018 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -187,7 +187,6 @@ use datafusion_common::{
assert_eq_or_internal_err, internal_err, not_impl_err,
};
use datafusion_execution::TaskContext;
-use datafusion_execution::memory_pool::MemoryLimit;
use datafusion_expr::{Accumulator, Aggregate};
use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
use datafusion_physical_expr::equivalence::ProjectionMapping;
@@ -1308,12 +1307,7 @@ impl AggregateExec {
&& self.group_by.is_single()
}
- fn should_use_partial_reduce_hash_stream(&self, context: &TaskContext) -> bool {
- // TODO: implement memory-limited path and remove this limitation
- if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) {
- return false;
- }
-
+ fn should_use_partial_reduce_hash_stream(&self, _context: &TaskContext) -> bool {
self.mode == AggregateMode::PartialReduce
&& self.limit_options.is_none()
&& self.input_order_mode == InputOrderMode::Linear
@@ -3241,7 +3235,7 @@ mod tests {
use arrow::array::{
BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array,
- Int64Array, StructArray, UInt32Array, UInt64Array,
+ Int64Array, NullArray, StructArray, UInt32Array, UInt64Array,
};
use arrow::compute::{SortOptions, concat_batches};
use arrow::datatypes::Int32Type;
@@ -4586,6 +4580,15 @@ mod tests {
}
fn partial_reduce_test_aggregate() -> Result<AggregateExec> {
+ partial_reduce_test_aggregate_with_batches(1)
+ }
+
+ /// Partial-reduce aggregate over `num_input_batches` identical input batches
+ /// of partial states, each reducing to groups `1, 2, 3` with sums
+ /// `50, 20, 30`.
+ fn partial_reduce_test_aggregate_with_batches(
+ num_input_batches: usize,
+ ) -> Result<AggregateExec> {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::UInt32, false),
Field::new("b", DataType::Float64, false),
@@ -4618,7 +4621,7 @@ mod tests {
],
)?;
let partial_reduce_input = TestMemoryExec::try_new_exec(
- &[vec![partial_state_batch]],
+ &[vec![partial_state_batch; num_input_batches]],
Arc::clone(&partial_schema),
None,
)?;
@@ -4654,26 +4657,162 @@ mod tests {
Ok(())
}
- /// Spilling behavior is not implemented for partial-reduce stream yet, so fall
- /// back to the existing `GroupedHashAggregateStream`
+ /// Partial-reduce hash aggregation emits its accumulated partial states early
+ /// under memory pressure instead of failing, and the early-emitted states
+ /// still merge into the correct result.
#[tokio::test]
- async fn partial_reduce_aggregate_with_memory_limit_planning() -> Result<()> {
- let partial_reduce = partial_reduce_test_aggregate()?;
+ async fn partial_reduce_aggregate_with_memory_limit_emits_early() -> Result<()> {
+ let num_input_batches = 3;
+ let partial_reduce =
+ partial_reduce_test_aggregate_with_batches(num_input_batches)?;
let runtime = RuntimeEnvBuilder::new()
.with_memory_limit(1, 1.0)
.build_arc()?;
- let task_ctx =
- Arc::new(
- TaskContext::default()
- .with_session_config(SessionConfig::new().set_bool(
- "datafusion.execution.enable_migration_aggregate",
- true,
- ))
- .with_runtime(runtime),
- );
+ // A batch size smaller than the number of flushed groups also covers
+ // splitting one flush across several output batches.
+ let batch_size = 2;
+ let task_ctx = Arc::new(
+ TaskContext::default()
+ .with_session_config(migrated_hash_session_config(batch_size))
+ .with_runtime(runtime),
+ );
let stream = partial_reduce.execute_typed(0, &task_ctx)?;
- assert!(matches!(stream, StreamType::GroupedHash(_)));
+ assert!(matches!(stream, StreamType::PartialReduceHash(_)));
+ let stream: SendableRecordBatchStream = stream.into();
+ let output = collect(stream).await?;
+
+ // The table is flushed after every input batch, so each of the three
+ // groups is emitted once per input batch instead of being merged into a
+ // single row. Each flush is sliced into batches of 2 and 1 rows.
+ assert_eq!(output.len(), 2 * num_input_batches);
+ assert_snapshot!(batches_to_string(&output), @r"
+ +---+-------------+
+ | a | SUM(b)[sum] |
+ +---+-------------+
+ | 1 | 50.0 |
+ | 2 | 20.0 |
+ | 3 | 30.0 |
+ | 1 | 50.0 |
+ | 2 | 20.0 |
+ | 3 | 30.0 |
+ | 1 | 50.0 |
+ | 2 | 20.0 |
+ | 3 | 30.0 |
+ +---+-------------+
+ ");
+
+ Ok(())
+ }
+
+ /// Same shape as [`partial_reduce_test_aggregate_with_batches`], but with multiple
+ /// group keys.
+ fn partial_reduce_test_aggregate_rows_multi_group_keys(
+ num_input_batches: usize,
+ ) -> Result<AggregateExec> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::UInt32, false),
+ Field::new("n", DataType::Null, true),
+ Field::new("b", DataType::Float64, false),
+ ]));
+ let group_by = PhysicalGroupBy::new_single(vec![
+ (col("a", &schema)?, "a".to_string()),
+ (col("n", &schema)?, "n".to_string()),
+ ]);
+ let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new(
+ AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("SUM(b)")
+ .build()?,
+ )];
+
+ let empty_input =
+ TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?;
+ let partial = AggregateExec::try_new(
+ AggregateMode::Partial,
+ group_by.clone(),
+ aggregates.clone(),
+ vec![None],
+ empty_input,
+ Arc::clone(&schema),
+ )?;
+ let partial_schema = partial.schema();
+ let partial_state_batch = RecordBatch::try_new(
+ Arc::clone(&partial_schema),
+ vec![
+ Arc::new(UInt32Array::from(vec![1, 2, 1, 3])),
+ Arc::new(NullArray::new(4)),
+ Arc::new(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])),
+ ],
+ )?;
+ let partial_reduce_input = TestMemoryExec::try_new_exec(
+ &[vec![partial_state_batch; num_input_batches]],
+ Arc::clone(&partial_schema),
+ None,
+ )?;
+
+ AggregateExec::try_new(
+ AggregateMode::PartialReduce,
+ group_by,
+ aggregates,
+ vec![None],
+ partial_reduce_input,
+ partial_schema,
+ )
+ }
+
+ #[tokio::test]
+ async fn partial_reduce_aggregate_with_memory_limit_emits_early_multi_group_keys()
+ -> Result<()> {
+ let num_input_batches = 3;
+ let partial_reduce =
+ partial_reduce_test_aggregate_rows_multi_group_keys(num_input_batches)?;
+
+ // Pin the representation: this is exactly the condition
+ // `new_group_values` uses to pick `GroupValuesRows` over
+ // `GroupValuesColumn`. If a `Null` `GroupColumn` is ever added, this
+ // assertion fires and the test stops covering the row-encoded path.
+ let group_schema = partial_reduce
+ .group_by
+ .group_schema(&partial_reduce.schema())?;
+ assert!(
+ !group_values::multi_group_by::supported_schema(&group_schema),
+ "expected the Null group column to force the GroupValuesRows fallback"
+ );
+
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(1, 1.0)
+ .build_arc()?;
+ let batch_size = 2;
+ let task_ctx = Arc::new(
+ TaskContext::default()
+ .with_session_config(migrated_hash_session_config(batch_size))
+ .with_runtime(runtime),
+ );
+
+ let stream = partial_reduce.execute_typed(0, &task_ctx)?;
+ assert!(matches!(stream, StreamType::PartialReduceHash(_)));
+ let stream: SendableRecordBatchStream = stream.into();
+ let output = collect(stream).await?;
+
+ // Same flush cadence as the column-backed test: one flush per input
+ // batch, each sliced into batches of 2 and 1 rows.
+ assert_eq!(output.len(), 2 * num_input_batches);
+ assert_snapshot!(batches_to_string(&output), @r"
+ +---+---+-------------+
+ | a | n | SUM(b)[sum] |
+ +---+---+-------------+
+ | 1 | | 50.0 |
+ | 2 | | 20.0 |
+ | 3 | | 30.0 |
+ | 1 | | 50.0 |
+ | 2 | | 20.0 |
+ | 3 | | 30.0 |
+ | 1 | | 50.0 |
+ | 2 | | 20.0 |
+ | 3 | | 30.0 |
+ +---+---+-------------+
+ ");
Ok(())
}
diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs
index 9eedf86847..d8f1447bc9 100644
--- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs
@@ -28,7 +28,7 @@ use std::task::{Context, Poll};
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
-use datafusion_common::Result;
+use datafusion_common::{DataFusionError, Result};
use datafusion_execution::TaskContext;
use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
use futures::stream::{Stream, StreamExt};
@@ -68,6 +68,15 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream};
/// This stage is useful for tree-reduce plans. It consumes the same schema as
/// a final aggregate stage, but emits the same schema as a partial aggregate
/// stage.
+///
+/// # Memory Management
+///
+/// If the memory reservation cannot grow after aggregating an input batch, all
+/// accumulated partial states are emitted immediately, and the remaining input
+/// is aggregated with an empty table. This repeats until the input ends.
+///
+/// See [`crate::aggregates::AggregateMode::PartialReduce`] for why it's allowed
+/// to emit the same group multiple times.
pub(crate) struct PartialReduceHashAggregateStream {
/// Output schema: group columns followed by partial aggregate state columns.
schema: SchemaRef,
@@ -75,6 +84,9 @@ pub(crate) struct PartialReduceHashAggregateStream {
/// Input batches containing partial aggregate state rows.
input: SendableRecordBatchStream,
+ /// Target output batch size from configuration.
+ batch_size: usize,
+
/// Execution metrics shared with the aggregate plan node.
baseline_metrics: BaselineMetrics,
@@ -93,10 +105,22 @@ enum PartialReduceHashAggregateState {
ReadingInput {
hash_table: AggregateHashTable<PartialReduceMarker>,
},
+ /// A fully materialized partial-state batch being emitted incrementally
+ /// because the table ran out of memory while reading input.
+ EmittingOnMemoryPressure {
+ hash_table: AggregateHashTable<PartialReduceMarker>,
+ // After each incremental emitting step, `remaining_groups` is updated
+ // with batch slicing.
+ remaining_groups: RecordBatch,
+ },
ProducingOutput {
hash_table: AggregateHashTable<PartialReduceMarker>,
},
Done,
+ /// Sentinel state to use when returning error from any other states, because:
+ /// - It explicitly releases state-owned resources immediately
+ /// - More defensive against accidentally resuming execution after error
+ Error,
}
type PartialReduceHashAggregatePoll = Poll<Option<Result<RecordBatch>>>;
@@ -111,28 +135,34 @@ type PartialReduceHashAggregateStateTransition = ControlFlow<
impl PartialReduceHashAggregateState {
fn hash_table(&self) -> &AggregateHashTable<PartialReduceMarker> {
match self {
- Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => {
- hash_table
+ Self::ReadingInput { hash_table }
+ | Self::EmittingOnMemoryPressure { hash_table, .. }
+ | Self::ProducingOutput { hash_table } => hash_table,
+ Self::Done | Self::Error => {
+ unreachable!("Done and Error states do not hold a hash table")
}
- Self::Done => unreachable!("Done state does not hold a hash table"),
}
}
fn hash_table_mut(&mut self) -> &mut AggregateHashTable<PartialReduceMarker> {
match self {
- Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => {
- hash_table
+ Self::ReadingInput { hash_table }
+ | Self::EmittingOnMemoryPressure { hash_table, .. }
+ | Self::ProducingOutput { hash_table } => hash_table,
+ Self::Done | Self::Error => {
+ unreachable!("Done and Error states do not hold a hash table")
}
- Self::Done => unreachable!("Done state does not hold a hash table"),
}
}
fn into_hash_table(self) -> AggregateHashTable<PartialReduceMarker> {
match self {
- Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => {
- hash_table
+ Self::ReadingInput { hash_table }
+ | Self::EmittingOnMemoryPressure { hash_table, .. }
+ | Self::ProducingOutput { hash_table } => hash_table,
+ Self::Done | Self::Error => {
+ unreachable!("Done and Error states do not hold a hash table")
}
- Self::Done => unreachable!("Done state does not hold a hash table"),
}
}
@@ -178,19 +208,25 @@ impl PartialReduceHashAggregateStream {
Ok(Self {
schema,
input,
+ batch_size,
baseline_metrics,
reservation,
state: Some(PartialReduceHashAggregateState::ReadingInput { hash_table }),
})
}
- fn start_output(
- &mut self,
- hash_table: &mut AggregateHashTable<PartialReduceMarker>,
- ) -> Result<()> {
+ fn close_input(&mut self) {
let input_schema = self.input.schema();
self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
- hash_table.start_output()
+ }
+
+ fn break_with_err(
+ error: DataFusionError,
+ ) -> PartialReduceHashAggregateStateTransition {
+ ControlFlow::Break((
+ Poll::Ready(Some(Err(error))),
+ PartialReduceHashAggregateState::Error,
+ ))
}
/// Handle ReadingInput state - aggregate partial state batches into the hash table.
@@ -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`. The reservation will be reset after exiting the
+ /// `EmittingOnMemoryPressure` state.
+ ///
+ /// # 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`
+ (
+ batch,
+ PartialReduceHashAggregateState::ReadingInput { hash_table },
+ )
+ } else {
+ // More batches to output, continue in the current state.
+ let remaining =
+ batch.slice(self.batch_size, batch.num_rows() - self.batch_size);
+ let output = batch.slice(0, self.batch_size);
+ (
+ output,
+ PartialReduceHashAggregateState::EmittingOnMemoryPressure {
+ hash_table,
+ remaining_groups: remaining,
+ },
+ )
+ };
+
+ debug_assert!(output_batch.num_rows() > 0);
+ ControlFlow::Break((
+ Poll::Ready(Some(Ok(output_batch.record_output(&self.baseline_metrics)))),
+ next_state,
+ ))
+ }
+
/// Handle ProducingOutput state - emit merged partial aggregate state batches.
///
/// See comments at `poll_next()` for details.
@@ -281,6 +397,8 @@ impl PartialReduceHashAggregateStream {
match result {
Ok(Some(batch)) => {
+ // The output is already materialized, so a failed resize cannot
+ // be acted on: keep the reservation as is and finish the output.
let _ = self
.reservation
.try_resize(original_state.hash_table().memory_size());
@@ -300,7 +418,7 @@ impl PartialReduceHashAggregateStream {
let _ = self.reservation.try_resize(0);
ControlFlow::Continue(original_state.into_done())
}
- Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)),
+ Err(e) => Self::break_with_err(e),
}
}
}
@@ -325,10 +443,22 @@ impl Stream for PartialReduceHashAggregateStream {
/// Aggregate one partial-state input batch, update the inner aggregate
/// hash table, and continue with the next input batch.
///
+ /// -> EmittingOnMemoryPressure
+ /// The table cannot reserve enough memory. Materialize all accumulated
+ /// partial states and begin emitting them incrementally.
+ ///
/// -> ProducingOutput
/// Input was exhausted. Move to the next state to start outputting
/// merged partial aggregate states.
///
+ /// EmittingOnMemoryPressure
+ /// -> EmittingOnMemoryPressure
+ /// One batch-sized slice was yielded; repeat until all materialized
+ /// partial states are emitted.
+ ///
+ /// -> ReadingInput
+ /// The materialized states were emitted; continue with the empty table.
+ ///
/// ProducingOutput
/// -> ProducingOutput
/// One merged partial-state output batch was yielded; repeat to
@@ -337,6 +467,13 @@ impl Stream for PartialReduceHashAggregateStream {
/// -> Done
/// All merged partial-state output was emitted.
///
+ /// Any active state
+ /// -> Error
+ /// An error drops state-owned resources before it is returned.
+ ///
+ /// Error
+ /// -> (end)
+ ///
/// Done
/// -> (end)
/// ```
@@ -354,9 +491,18 @@ impl Stream for PartialReduceHashAggregateStream {
state @ PartialReduceHashAggregateState::ReadingInput { .. } => {
self.handle_reading_input(cx, state)
}
+ state @ PartialReduceHashAggregateState::EmittingOnMemoryPressure {
+ ..
+ } => self.handle_emitting_on_memory_pressure(state),
state @ PartialReduceHashAggregateState::ProducingOutput { .. } => {
self.handle_producing_output(state)
}
+ state @ PartialReduceHashAggregateState::Error => {
+ self.close_input();
+ self.reservation.free();
+ self.state = Some(state);
+ return Poll::Ready(None);
+ }
state @ PartialReduceHashAggregateState::Done => {
let _ = self.reservation.try_resize(0);
self.state = Some(state);
@@ -368,6 +514,19 @@ impl Stream for PartialReduceHashAggregateStream {
ControlFlow::Continue(next_state) => {
self.state = Some(next_state);
}
+ ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => {
+ debug_assert!(matches!(
+ next_state,
+ PartialReduceHashAggregateState::Error
+ ));
+
+ // The handler has already discarded its state-owned resources.
+ // Release the remaining stream-owned resources before returning.
+ self.close_input();
+ self.reservation.free();
+ self.state = Some(PartialReduceHashAggregateState::Error);
+ return Poll::Ready(Some(Err(e)));
+ }
ControlFlow::Break((poll, next_state)) => {
self.state = Some(next_state);
return poll;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org