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/06 08:11:47 UTC
(datafusion) branch gh-readonly-queue/main/pr-24017-8642154323c3800948f495af9ba2ee1539399105 created (now 5bf6aef8e3)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24017-8642154323c3800948f495af9ba2ee1539399105
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 5bf6aef8e3 chore: convert `PartialHashAggregateStream` to async generators and cleanup (#24017)
This branch includes the following new commits:
new 5bf6aef8e3 chore: convert `PartialHashAggregateStream` to async generators and cleanup (#24017)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: chore: convert `PartialHashAggregateStream` to async generators and cleanup (#24017)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24017-8642154323c3800948f495af9ba2ee1539399105
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 5bf6aef8e37c91d882c0f75b772201d62733c88f
Author: Raz Luvaton <16...@users.noreply.github.com>
AuthorDate: Sun Sep 6 08:11:24 2026 +0000
chore: convert `PartialHashAggregateStream` to async generators and cleanup (#24017)
## Which issue does this PR close?
Related to:
- #23974
## Rationale for this change
Cleanup the code and remove state
## What changes are included in this PR?
changed `PartialHashAggregateStream` to async generator and remove
unneeded code due to that
## Are these changes tested?
existing tests
## Are there any user-facing changes?
nope
---
.../physical-plan/src/aggregates/hash_stream.rs | 696 ++++++---------------
datafusion/physical-plan/src/aggregates/mod.rs | 2 +-
2 files changed, 201 insertions(+), 497 deletions(-)
diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs
index b88be67353..7ebf32c3ed 100644
--- a/datafusion/physical-plan/src/aggregates/hash_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs
@@ -26,9 +26,7 @@
//! See issue for details: <https://github.com/apache/datafusion/issues/22710>
use std::mem::size_of;
-use std::ops::ControlFlow;
use std::sync::Arc;
-use std::task::{Context, Poll};
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
@@ -57,7 +55,7 @@ use crate::sorts::IncrementalSortIterator;
use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
use crate::spill::spill_manager::SpillManager;
use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter};
-use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics};
+use crate::{InputOrderMode, SendableRecordBatchStream, metrics};
/// Hash aggregation is implemented in two stages: partial and final. This
/// stream implements the partial stage.
@@ -179,47 +177,10 @@ pub(crate) struct PartialHashAggregateStream {
/// be empty. See struct comments for details.
group_values_soft_limit: Option<usize>,
- /// Tracks the high-level stream lifecycle. The hash table owns the lower-level
- /// state for emitting output batches.
- state: Option<PartialHashAggregateState>,
+ /// The hash table owns the lower-level state for emitting output batches.
+ hash_table: Option<AggregateHashTable<PartialMarker>>,
}
-/// States for partial hash aggregation processing.
-enum PartialHashAggregateState {
- ReadingInput {
- hash_table: AggregateHashTable<PartialMarker>,
- },
- /// A fully materialized partial-state batch being emitted incrementally.
- EmittingOnMemoryPressure {
- hash_table: AggregateHashTable<PartialMarker>,
- // After each incremental emitting step, the `remaining_groups` will be updated
- // with batch slicing.
- remaining_groups: RecordBatch,
- },
- ProducingOutput {
- hash_table: AggregateHashTable<PartialMarker>,
- /// If `None`, partial skip was never triggered and this state will
- /// finish in `Done`. If `Some`, partial skip has triggered and the
- /// stream will move to `SkippingAggregation` after these accumulated
- /// groups are emitted.
- skip_hash_table: Option<AggregateHashTable<PartialSkipMarker>>,
- },
- SkippingAggregation {
- hash_table: AggregateHashTable<PartialSkipMarker>,
- },
- 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 PartialHashAggregatePoll = Poll<Option<Result<RecordBatch>>>;
-type PartialHashAggregateStateTransition = ControlFlow<
- (PartialHashAggregatePoll, PartialHashAggregateState),
- PartialHashAggregateState,
->;
-
/// Spill configuration and accumulated runs for final hash aggregation.
///
/// Each spill event drains all currently buffered groups, sorts their intermediate
@@ -405,6 +366,15 @@ impl FinalSpillContext {
}
}
+#[derive(PartialEq)]
+enum HandleInputResult {
+ ProcessNext,
+ ReachedLimit,
+ #[expect(clippy::upper_case_acronyms)]
+ OOM,
+ SwitchToSkipAggregation,
+}
+
impl PartialHashAggregateStream {
pub fn new(
agg: &AggregateExec,
@@ -467,26 +437,85 @@ impl PartialHashAggregateStream {
reduction_factor,
skip_aggregation_probe,
group_values_soft_limit: agg.limit_options().map(|config| config.limit()),
- state: Some(PartialHashAggregateState::ReadingInput { hash_table }),
+ hash_table: Some(hash_table),
})
}
- fn close_input(&mut self) {
- let input_schema = self.input.schema();
- self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
- }
+ pub(crate) fn into_stream(self) -> SendableRecordBatchStream {
+ let schema = Arc::clone(&self.schema);
- fn break_with_err(error: DataFusionError) -> PartialHashAggregateStateTransition {
- ControlFlow::Break((
- Poll::Ready(Some(Err(error))),
- PartialHashAggregateState::Error,
- ))
+ Box::pin(RecordBatchStreamAdapter::new(schema, self.create_stream()))
}
- fn break_with_internal_err(
- message: impl std::fmt::Display,
- ) -> PartialHashAggregateStateTransition {
- Self::break_with_err(internal_datafusion_err!("{message}"))
+ /// Entry point for the partial hash aggregate state machine.
+ ///
+ /// See comments in [`PartialHashAggregateStream`] for high-level ideas.
+ fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> {
+ async_try_stream(|mut emitter| async move {
+ let mut hash_table = self
+ .hash_table
+ .take()
+ .expect("hash_table should not be None");
+
+ debug_assert!(hash_table.is_building());
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+
+ let mut last_state = HandleInputResult::ProcessNext;
+ while let Some(batch) = self.input.next().await.transpose()? {
+ let timer = elapsed_compute.timer();
+ last_state = self.handle_input_batch(&batch, &mut hash_table)?;
+
+ match last_state {
+ HandleInputResult::ProcessNext => {}
+ HandleInputResult::ReachedLimit
+ | HandleInputResult::SwitchToSkipAggregation => {
+ break;
+ }
+ HandleInputResult::OOM => {
+ let state_batch_result = hash_table.take_state_batch();
+
+ // Emitting clears the aggregate table and releases its
+ // accumulated memory. Update the reservation accordingly.
+ self.reservation.try_resize(hash_table.memory_size())?;
+
+ let materialized_group_states = state_batch_result?.ok_or_else(|| {
+ internal_datafusion_err!(
+ "Partial hash aggregate ran out of memory with no aggregated groups"
+ )
+ })?;
+
+ timer.done();
+ self.emit_on_memory_pressure(
+ materialized_group_states,
+ &mut emitter,
+ )
+ .await?;
+ }
+ }
+ }
+
+ let timer = elapsed_compute.timer();
+
+ let skip_hash_table =
+ if last_state == HandleInputResult::SwitchToSkipAggregation {
+ Some(hash_table.partial_skip_table()?)
+ } else {
+ self.close_input();
+
+ None
+ };
+ hash_table.start_output()?;
+
+ timer.done();
+
+ self.produce_output(hash_table, &mut emitter).await?;
+
+ if let Some(hash_table) = skip_hash_table {
+ self.skip_rest_of_aggregation(hash_table, emitter).await?;
+ }
+
+ Ok(())
+ })
}
/// See comments in [`Self::group_values_soft_limit`] for details.
@@ -513,477 +542,150 @@ impl PartialHashAggregateStream {
.is_some_and(|probe| probe.should_skip())
}
- fn start_output(
- &mut self,
- hash_table: &mut AggregateHashTable<PartialMarker>,
- close_input: bool,
- ) -> Result<()> {
- if close_input {
- let input_schema = self.input.schema();
- self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
- }
- hash_table.start_output()
+ fn close_input(&mut self) {
+ let input_schema = self.input.schema();
+ self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
}
- /// Handle ReadingInput state - aggregate input batches into the hash table.
- ///
- /// See comments at `poll_next()` for details.
- ///
- /// Returns the next operator state with control flow decision.
- fn handle_reading_input(
+ /// Aggregate input batch into the hash table
+ fn handle_input_batch(
&mut self,
- cx: &mut Context<'_>,
- original_state: PartialHashAggregateState,
- ) -> PartialHashAggregateStateTransition {
- let PartialHashAggregateState::ReadingInput { mut hash_table } = original_state
- else {
- return Self::break_with_internal_err(
- "Partial hash aggregate stream expected ReadingInput state",
- );
- };
- debug_assert!(hash_table.is_building());
-
- match self.input.poll_next_unpin(cx) {
- Poll::Pending => ControlFlow::Break((
- Poll::Pending,
- PartialHashAggregateState::ReadingInput { hash_table },
- )),
- Poll::Ready(Some(Ok(batch))) => {
- // ----------------------------------
- // Step 1: Aggregate the input batch
- // ----------------------------------
- let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
- let timer = elapsed_compute.timer();
- let input_rows = batch.num_rows();
- self.reduction_factor.add_total(input_rows);
- let result = hash_table.aggregate_batch(&batch);
- timer.done();
-
- if let Err(e) = result {
- return Self::break_with_err(e);
- }
+ batch: &RecordBatch,
+ hash_table: &mut AggregateHashTable<PartialMarker>,
+ ) -> Result<HandleInputResult> {
+ // ----------------------------------
+ // Step 1: Aggregate the input batch
+ // ----------------------------------
+ let input_rows = batch.num_rows();
+ self.reduction_factor.add_total(input_rows);
+ hash_table.aggregate_batch(batch)?;
+
+ // --------------------------------
+ // Step 2: Soft limit optimization
+ // --------------------------------
+ if self.hit_soft_group_limit(hash_table) {
+ return Ok(HandleInputResult::ReachedLimit);
+ }
- // --------------------------------
- // Step 2: Soft limit optimization
- // --------------------------------
- if self.hit_soft_group_limit(&hash_table) {
- let timer = elapsed_compute.timer();
- let result = self.start_output(&mut hash_table, true);
- timer.done();
+ // ----------------------------------------------
+ // Step 3: Skip partial aggregation optimization
+ // ----------------------------------------------
+ self.update_skip_aggregation_probe(input_rows, hash_table.building_group_count());
- if let Err(e) = result {
- return Self::break_with_err(e);
- }
+ // True branch: a decision has been made to skip partial aggregation.
+ if self.should_skip_aggregation() {
+ return Ok(HandleInputResult::SwitchToSkipAggregation);
+ }
- return ControlFlow::Continue(
- PartialHashAggregateState::ProducingOutput {
- hash_table,
- skip_hash_table: None,
- },
- );
- }
+ // -------------------------------------------------
+ // Step 4: Larger-than-memory execution (early emit)
+ // -------------------------------------------------
+ let resize_result = self.reservation.try_resize(hash_table.memory_size());
+ match resize_result {
+ Ok(()) => Ok(HandleInputResult::ProcessNext),
+ Err(DataFusionError::ResourcesExhausted(_)) => Ok(HandleInputResult::OOM),
+ Err(e) => Err(e),
+ }
+ }
- // ----------------------------------------------
- // Step 3: Skip partial aggregation optimization
- // ----------------------------------------------
- self.update_skip_aggregation_probe(
- input_rows,
- hash_table.building_group_count(),
- );
-
- // True branch: a decision has been made to skip partial aggregation.
- if self.should_skip_aggregation() {
- let timer = elapsed_compute.timer();
- let result = match hash_table.partial_skip_table() {
- Ok(skip_hash_table) => self
- .start_output(&mut hash_table, false)
- .map(|()| skip_hash_table),
- Err(e) => Err(e),
- };
- timer.done();
-
- match result {
- Ok(skip_hash_table) => {
- // Move to `ProducingOutput` first. Its `skip_hash_table`
- // field moves the stream to skip-partial aggregation after
- // the accumulated batches have been output.
- return ControlFlow::Continue(
- PartialHashAggregateState::ProducingOutput {
- hash_table,
- skip_hash_table: Some(skip_hash_table),
- },
- );
- }
- Err(e) => return Self::break_with_err(e),
- }
- }
+ /// emit a materialized partial-state on memory pressure
+ /// batch in `batch_size`(from configuration) slices
+ async fn emit_on_memory_pressure(
+ &mut self,
+ // After each incremental emitting step, the `remaining_groups` will be updated
+ // with batch slicing.
+ mut remaining_groups: RecordBatch,
+ emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ while remaining_groups.num_rows() > self.batch_size {
+ // More batch to output, continue in the current state.
+ let output = remaining_groups.slice(0, self.batch_size);
- // -------------------------------------------------
- // Step 4: Larger-than-memory execution (early emit)
- // -------------------------------------------------
- let timer = elapsed_compute.timer();
- let resize_result = self.reservation.try_resize(hash_table.memory_size());
- timer.done();
- match resize_result {
- Ok(()) => {}
- Err(DataFusionError::ResourcesExhausted(_)) => {
- let elapsed_compute =
- self.baseline_metrics.elapsed_compute().clone();
- // Stops on drop
- let _timer = elapsed_compute.timer();
- let state_batch_result = hash_table.take_state_batch();
+ remaining_groups = remaining_groups.slice(
+ self.batch_size,
+ remaining_groups.num_rows() - self.batch_size,
+ );
- // 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());
-
- if let Err(e) = resize_result {
- return Self::break_with_err(e);
- }
-
- let materialized_group_states = match state_batch_result {
- Ok(Some(batch)) => batch,
- Ok(None) => {
- return Self::break_with_err(internal_datafusion_err!(
- "Partial hash aggregate ran out of memory with no aggregated groups"
- ));
- }
- Err(e) => return Self::break_with_err(e),
- };
-
- return ControlFlow::Continue(
- PartialHashAggregateState::EmittingOnMemoryPressure {
- hash_table,
- remaining_groups: materialized_group_states,
- },
- );
- }
- Err(e) => return Self::break_with_err(e),
- }
+ self.reduction_factor.add_part(output.num_rows());
+ debug_assert!(output.num_rows() > 0);
- ControlFlow::Continue(PartialHashAggregateState::ReadingInput {
- hash_table,
- })
- }
- Poll::Ready(Some(Err(e))) => Self::break_with_err(e),
- Poll::Ready(None) => {
- let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
- let timer = elapsed_compute.timer();
- let result = self.start_output(&mut hash_table, true);
- timer.done();
-
- match result {
- Ok(()) => ControlFlow::Continue(
- PartialHashAggregateState::ProducingOutput {
- hash_table,
- skip_hash_table: None,
- },
- ),
- Err(e) => Self::break_with_err(e),
- }
- }
+ emitter
+ .emit(output.record_output(&self.baseline_metrics))
+ .await;
}
- }
- /// Handle EmittingOnMemoryPressure state - emit a materialized partial-state
- /// batch in `batch_size`(from configuration) slices, then 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: PartialHashAggregateState,
- ) -> PartialHashAggregateStateTransition {
- let PartialHashAggregateState::EmittingOnMemoryPressure {
- hash_table,
- remaining_groups: batch,
- } = original_state
- else {
- return Self::break_with_internal_err(
- "Partial hash aggregate stream expected EmittingOnMemoryPressure state",
- );
- };
+ self.reduction_factor.add_part(remaining_groups.num_rows());
+ debug_assert!(remaining_groups.num_rows() > 0);
- let (output_batch, next_state) = if batch.num_rows() <= self.batch_size {
- // Last batch to output, go back to `ReadingInput`
- (
- batch,
- PartialHashAggregateState::ReadingInput { hash_table },
- )
- } else {
- // More batch 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,
- PartialHashAggregateState::EmittingOnMemoryPressure {
- hash_table,
- remaining_groups: remaining,
- },
- )
- };
+ emitter
+ .emit(remaining_groups.record_output(&self.baseline_metrics))
+ .await;
- self.reduction_factor.add_part(output_batch.num_rows());
- debug_assert!(output_batch.num_rows() > 0);
- ControlFlow::Break((
- Poll::Ready(Some(Ok(output_batch.record_output(&self.baseline_metrics)))),
- next_state,
- ))
+ Ok(())
}
- /// Handle ProducingOutput state - emit partial aggregate state batches.
- ///
- /// See comments at `poll_next()` for details.
- ///
- /// Returns the next operator state with control flow decision.
- fn handle_producing_output(
+ /// emit partial aggregate state batches.
+ async fn produce_output(
&mut self,
- original_state: PartialHashAggregateState,
- ) -> PartialHashAggregateStateTransition {
- let PartialHashAggregateState::ProducingOutput {
- mut hash_table,
- skip_hash_table,
- } = original_state
- else {
- return Self::break_with_internal_err(
- "Partial hash aggregate stream expected ProducingOutput state",
- );
- };
+ mut hash_table: AggregateHashTable<PartialMarker>,
+ emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
debug_assert!(!hash_table.is_building());
let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
- let timer = elapsed_compute.timer();
- let result = hash_table.next_output_batch();
- timer.done();
-
- match result {
- Ok(Some(batch)) => {
- let _ = self.reservation.try_resize(hash_table.memory_size());
- self.reduction_factor.add_part(batch.num_rows());
- debug_assert!(batch.num_rows() > 0);
- let next_state = if hash_table.is_done() {
- match skip_hash_table {
- Some(hash_table) => {
- PartialHashAggregateState::SkippingAggregation { hash_table }
- }
- None => PartialHashAggregateState::Done,
- }
- } else {
- PartialHashAggregateState::ProducingOutput {
- hash_table,
- skip_hash_table,
- }
- };
+ let mut timer = elapsed_compute.timer();
- ControlFlow::Break((
- Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))),
- next_state,
- ))
- }
- Ok(None) => {
- let _ = self.reservation.try_resize(0);
- // If the previous `Aggregating` stage decided to skip partial
- // aggregation, go to the `SkippingAggregation` stage; otherwise finish.
- let next_state = match skip_hash_table {
- Some(hash_table) => {
- PartialHashAggregateState::SkippingAggregation { hash_table }
- }
- None => PartialHashAggregateState::Done,
- };
- ControlFlow::Continue(next_state)
- }
- Err(e) => Self::break_with_err(e),
- }
- }
+ loop {
+ let Some(batch) = hash_table.next_output_batch()? else {
+ // Only reachable when the table held no groups at all: a
+ // non-empty table always reports its last batch together with
+ // the `Done` state, which the `try_resize` below already zeroes.
+ self.reservation.try_resize(0)?;
+ return Ok(());
+ };
- /// Handle SkippingAggregation state - convert raw input directly to partial states.
- ///
- /// See comments at `poll_next()` for details.
- ///
- /// Returns the next operator state with control flow decision.
- fn handle_skipping_aggregation(
- &mut self,
- cx: &mut Context<'_>,
- original_state: PartialHashAggregateState,
- ) -> PartialHashAggregateStateTransition {
- let PartialHashAggregateState::SkippingAggregation { mut hash_table } =
- original_state
- else {
- return Self::break_with_internal_err(
- "Partial hash aggregate stream expected SkippingAggregation state",
- );
- };
+ debug_assert!(batch.num_rows() > 0);
- match self.input.poll_next_unpin(cx) {
- Poll::Pending => ControlFlow::Break((
- Poll::Pending,
- PartialHashAggregateState::SkippingAggregation { hash_table },
- )),
- Poll::Ready(Some(Ok(batch))) => {
- if let Some(probe) = self.skip_aggregation_probe.as_mut() {
- probe.record_skipped(&batch);
- }
+ // The table hands over its groups as they are materialized and
+ // reports a size of 0 once it reaches `Done`, so this releases the
+ // reservation before the final batch goes downstream.
+ let _ = self.reservation.try_resize(hash_table.memory_size());
+ self.reduction_factor.add_part(batch.num_rows());
- let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
- let timer = elapsed_compute.timer();
- let result = hash_table.convert_batch_to_state(&batch);
- timer.done();
-
- match result {
- Ok(batch) => ControlFlow::Break((
- Poll::Ready(Some(
- Ok(batch.record_output(&self.baseline_metrics)),
- )),
- PartialHashAggregateState::SkippingAggregation { hash_table },
- )),
- Err(e) => Self::break_with_err(e),
- }
- }
- Poll::Ready(Some(Err(e))) => Self::break_with_err(e),
- Poll::Ready(None) => {
- let input_schema = self.input.schema();
- self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
- ControlFlow::Continue(PartialHashAggregateState::Done)
- }
+ timer.done();
+ emitter
+ .emit(batch.record_output(&self.baseline_metrics))
+ .await;
+ timer = elapsed_compute.timer();
}
}
-}
-impl Stream for PartialHashAggregateStream {
- type Item = Result<RecordBatch>;
+ /// convert raw input directly to partial states.
+ async fn skip_rest_of_aggregation(
+ &mut self,
+ mut hash_table: AggregateHashTable<PartialSkipMarker>,
+ mut emitter: TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
- /// Entry point for the partial hash aggregate state machine.
- ///
- /// See comments in [`PartialHashAggregateStream`] for high-level ideas.
- ///
- /// State transition graph:
- ///
- /// ```text
- /// (start)
- /// -> ReadingInput
- /// The stream starts by polling input and aggregating batches into the
- /// in-memory hash table.
- ///
- /// ReadingInput
- /// -> ReadingInput
- /// Aggregate one 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(skip=None)
- /// Input was exhausted, or the soft group limit was reached. Move to
- /// the next state to start outputting.
- /// -> ProducingOutput(skip=Some)
- /// Partial skip aggregation was triggered. First move to the
- /// `ProducingOutput` state to drain the accumulated state, then move to
- /// the `SkippingAggregation` state to convert input directly to partial
- /// state without aggregation.
- ///
- /// 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(skip=None)
- /// -> ProducingOutput(skip=None)
- /// One accumulated output batch was yielded, repeat to continue producing
- /// output incrementally.
- /// -> Done
- /// All accumulated output was emitted.
- ///
- /// ProducingOutput(skip=Some)
- /// -> ProducingOutput(skip=Some)
- /// One accumulated output batch was yielded, repeat to continue producing
- /// output incrementally.
- /// -> SkippingAggregation
- /// All accumulated output was emitted. Continue by converting raw
- /// input batches directly to partial aggregate state.
- ///
- /// SkippingAggregation
- /// -> SkippingAggregation
- /// One `convert_to_state` batch was yielded; repeat to continue
- /// processing.
- /// -> Done
- /// Input was exhausted.
- ///
- /// Any active state
- /// -> Error
- /// An error drops state-owned resources before it is returned.
- ///
- /// Error
- /// -> (end)
- ///
- /// Done
- /// -> (end)
- /// ```
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- loop {
- let cur_state = self
- .state
- .take()
- .expect("PartialHashAggregateStream state should not be None");
+ while let Some(batch) = self.input.next().await.transpose()? {
+ if let Some(probe) = self.skip_aggregation_probe.as_mut() {
+ probe.record_skipped(&batch);
+ }
- let next_state = match cur_state {
- state @ PartialHashAggregateState::ReadingInput { .. } => {
- self.handle_reading_input(cx, state)
- }
- state @ PartialHashAggregateState::EmittingOnMemoryPressure { .. } => {
- self.handle_emitting_on_memory_pressure(state)
- }
- state @ PartialHashAggregateState::ProducingOutput { .. } => {
- self.handle_producing_output(state)
- }
- state @ PartialHashAggregateState::SkippingAggregation { .. } => {
- self.handle_skipping_aggregation(cx, state)
- }
- state @ PartialHashAggregateState::Error => {
- self.close_input();
- self.reservation.free();
- self.state = Some(state);
- return Poll::Ready(None);
- }
- state @ PartialHashAggregateState::Done => {
- let _ = self.reservation.try_resize(0);
- self.state = Some(state);
- return Poll::Ready(None);
- }
+ let result = {
+ let _timer = elapsed_compute.timer();
+ hash_table.convert_batch_to_state(&batch)?
};
- match next_state {
- ControlFlow::Continue(next_state) => {
- self.state = Some(next_state);
- }
- ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => {
- debug_assert!(matches!(next_state, PartialHashAggregateState::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(PartialHashAggregateState::Error);
- return Poll::Ready(Some(Err(e)));
- }
- ControlFlow::Break((poll, next_state)) => {
- self.state = Some(next_state);
- return poll;
- }
- }
+ emitter
+ .emit(result.record_output(&self.baseline_metrics))
+ .await;
}
- }
-}
-impl RecordBatchStream for PartialHashAggregateStream {
- fn schema(&self) -> SchemaRef {
- Arc::clone(&self.schema)
+ self.close_input();
+
+ Ok(())
}
}
@@ -1379,7 +1081,8 @@ mod tests {
// Execute and collect results
let mut stream =
- PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?;
+ PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?
+ .into_stream();
let mut results = Vec::new();
while let Some(result) = stream.next().await {
@@ -1523,7 +1226,8 @@ mod tests {
// Execute and collect results
let mut stream =
- PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?;
+ PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?
+ .into_stream();
let mut results = Vec::new();
while let Some(result) = stream.next().await {
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 7c87b12826..69fb09a8b1 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -713,7 +713,7 @@ impl From<StreamType> for SendableRecordBatchStream {
fn from(stream: StreamType) -> Self {
match stream {
StreamType::AggregateStream(stream) => Box::pin(stream),
- StreamType::PartialHash(stream) => Box::pin(stream),
+ StreamType::PartialHash(stream) => stream.into_stream(),
StreamType::PartialReduceHash(stream) => Box::pin(stream),
StreamType::FinalHash(stream) => stream.into_stream(),
StreamType::SingleHash(stream) => Box::pin(stream),
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org