From 255571fcf90bd9456019ef7167b94f8e479943a5 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 28 Jul 2026 23:03:14 +0300 Subject: [PATCH 1/3] chore(ordered-partial-aggregate): move `OrderedPartialAggregateStream` to generators for readability --- .../physical-plan/src/aggregates/mod.rs | 2 +- .../src/aggregates/ordered_partial_stream.rs | 349 +++++++----------- 2 files changed, 132 insertions(+), 219 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 787cd4f03ff6c..382828e038114 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -698,7 +698,7 @@ impl From for SendableRecordBatchStream { StreamType::PartialReduceHash(stream) => Box::pin(stream), StreamType::FinalHash(stream) => Box::pin(stream), StreamType::SingleHash(stream) => Box::pin(stream), - StreamType::OrderedPartialAggregate(stream) => Box::pin(stream), + StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), StreamType::GroupedHash(stream) => Box::pin(stream), StreamType::GroupedPriorityQueue(stream) => Box::pin(stream), diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 73d15a8278692..634fa79b46c52 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -24,8 +24,8 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result}; -use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use futures::stream::{Stream, StreamExt}; use super::AggregateExec; @@ -33,7 +33,7 @@ use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; use crate::aggregates::AggregateMode; use crate::aggregates::order::GroupOrdering; use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput, SpillMetrics}; -use crate::stream::EmptyRecordBatchStream; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; /// Partial aggregate stream for `InputOrderMode::Sorted` and @@ -118,7 +118,7 @@ pub(crate) struct OrderedPartialAggregateStream { reservation: MemoryReservation, baseline_metrics: BaselineMetrics, reduction_factor: metrics::RatioMetrics, - state: Option, + table: Option>, } /// See comments at `poll_next()` for details. @@ -178,7 +178,83 @@ impl OrderedPartialAggregateStream { reservation, baseline_metrics, reduction_factor, - state: Some(OrderedPartialAggregateState::ReadingInput { table }), + table: Some(table) + }) + } + + pub(crate) fn into_stream(self) -> SendableRecordBatchStream { + let schema_clone = Arc::clone(&self.schema); + + let cloned_metrics = self.baseline_metrics.clone(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema_clone, + self.create_stream(), + )); + + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) + } + + /// Entry point for the ordered partial aggregate state machine. + /// + /// See comments in [`OrderedPartialAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling ordered input and aggregating batches + /// into the ordered partial aggregate table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one input batch. If the ordering proves some groups are + /// complete, yield one partial-state batch immediately, then continue + /// reading input. Otherwise continue directly with the next input batch. + /// -> DrainingFinal + /// Input was exhausted. Mark the table input as done so every remaining + /// group is safe to emit. + /// + /// DrainingFinal + /// -> DrainingFinal + /// One remaining partial-state batch was yielded; repeat to continue + /// draining the table. + /// -> Done + /// All remaining groups were emitted. + /// + /// Done + /// -> (end) + /// ``` + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|mut emitter| async move { + let mut table = self + .table + .take() + .expect("OrderedPartialAggregateStream state should not be None"); + + self.handle_reading_input(&mut table, &mut emitter) + .await?; + + // Input has exhausted, move to the final draining stage. + self.close_input(); + table.input_done(); + + let last_batch = self + .handle_draining_final(&mut table, &mut emitter) + .await?; + + // Clear memory before emitting last batch so we don't have to wait for next poll to clear + { + // Clear memory + drop(table); + let _ = self.reservation.try_resize(0); + } + + if let Some(last_batch) = last_batch { + emitter.emit(last_batch).await; + } + + Ok(()) }) } @@ -193,110 +269,45 @@ impl OrderedPartialAggregateStream { /// See comments at `poll_next()` for details. /// /// Returns the next operator state with control flow decision. - fn handle_reading_input( + async fn handle_reading_input( &mut self, - cx: &mut Context<'_>, - original_state: OrderedPartialAggregateState, - ) -> OrderedPartialAggregateStateTransition { - let OrderedPartialAggregateState::ReadingInput { mut table } = original_state - else { - unreachable!("expected reading input state") - }; + table: &mut OrderedAggregateTable, + emitter: &mut TryEmitter, + ) -> Result<()> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break(( - Poll::Pending, - OrderedPartialAggregateState::ReadingInput { table }, - )), - Poll::Ready(Some(Ok(batch))) => { - let input_rows = batch.num_rows(); - self.reduction_factor.add_total(input_rows); - - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = table.aggregate_batch(&batch); - timer.done(); - - if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedPartialAggregateState::ReadingInput { table }, - )); - } + while let Some(batch) = self.input.next().await.transpose()? { + let input_rows = batch.num_rows(); + self.reduction_factor.add_total(input_rows); - // Check memory reservation. See function comments for details. - match self.resize_or_take_state_batch(&mut table) { - Ok(Some(batch)) => { - self.reduction_factor.add_part(batch.num_rows()); - return ControlFlow::Break(( - Poll::Ready(Some(Ok( - batch.record_output(&self.baseline_metrics) - ))), - OrderedPartialAggregateState::ReadingInput { table }, - )); - } - Ok(None) => {} - Err(e) => { - self.close_input(); - self.reservation.free(); - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedPartialAggregateState::Done, - )); - } - } + let timer = elapsed_compute.timer(); + + table.aggregate_batch(&batch)?; - let timer = elapsed_compute.timer(); - let result = table.next_output_batch(); - timer.done(); - - match result { - // There is some previous group results can be emitted: emit - // them, and next continuing aggreagting input (loop in the - // current state) - Ok(Some(batch)) => { - self.reduction_factor.add_part(batch.num_rows()); - if let Err(e) = self.reservation.try_resize(table.memory_size()) { - self.close_input(); - self.reservation.free(); - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedPartialAggregateState::Done, - )); - } - let next_state = - OrderedPartialAggregateState::ReadingInput { table }; - - ControlFlow::Break(( - Poll::Ready(Some(Ok( - batch.record_output(&self.baseline_metrics) - ))), - next_state, - )) - } - // Can't do early emit, continue aggregating. - Ok(None) => ControlFlow::Continue( - OrderedPartialAggregateState::ReadingInput { table }, - ), - Err(e) => ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedPartialAggregateState::ReadingInput { table }, - )), + // Check memory reservation. See function comments for details. + match self.resize_or_take_state_batch(table)? { + Some(batch) => { + self.reduction_factor.add_part(batch.num_rows()); + drop(timer); + emitter.emit(batch).await; + continue; } + None => {} } - Poll::Ready(Some(Err(e))) => ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedPartialAggregateState::ReadingInput { table }, - )), - // Input has exhausted, move to the final draining stage. - Poll::Ready(None) => { - self.close_input(); - table.input_done(); - ControlFlow::Continue(OrderedPartialAggregateState::DrainingFinal { - table, - }) - } + + let Some(batch) = table.next_output_batch()? else { + // Can't do early emit, continue aggregating. + continue; + }; + + self.reduction_factor.add_part(batch.num_rows()); + self.reservation.try_resize(table.memory_size())?; + + drop(timer); + emitter.emit(batch).await; } + + Ok(()) } /// Update the memory reservation, and: @@ -342,50 +353,30 @@ impl OrderedPartialAggregateStream { /// /// See comments at `poll_next()` for details. /// - /// Returns the next operator state with control flow decision. - fn handle_draining_final( + /// Returns the last batch to emit so we can clear all the state before emitting so we dont need to hold while waiting for next poll to reset the state + async fn handle_draining_final( &mut self, - original_state: OrderedPartialAggregateState, - ) -> OrderedPartialAggregateStateTransition { - let OrderedPartialAggregateState::DrainingFinal { table } = original_state else { - unreachable!("expected draining final state") - }; - - let mut table = table; + table: &mut OrderedAggregateTable, + emitter: &mut TryEmitter, + ) -> Result> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = table.next_output_batch(); - timer.done(); - - match result { - Ok(Some(batch)) => { - self.reduction_factor.add_part(batch.num_rows()); - let next_state = if table.is_empty() { - OrderedPartialAggregateState::Done - } else { - OrderedPartialAggregateState::DrainingFinal { table } - }; - if let Err(e) = self.resize_reservation_for_state(&next_state) { - return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)); - } + let mut timer = elapsed_compute.timer(); + while let Some(batch) = table.next_output_batch()? { + self.reduction_factor.add_part(batch.num_rows()); - ControlFlow::Break(( - Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), - next_state, - )) - } - Err(e) => ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedPartialAggregateState::DrainingFinal { table }, - )), - Ok(None) => { - let next_state = OrderedPartialAggregateState::Done; - if let Err(e) = self.resize_reservation_for_state(&next_state) { - return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)); - } - ControlFlow::Continue(next_state) + if table.is_empty() { + return Ok(Some(batch)); } + + self.reservation.try_resize(table.memory_size())?; + + timer.done(); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); } + + // was empty + Ok(None) } fn resize_reservation_for_state( @@ -402,81 +393,3 @@ impl OrderedPartialAggregateStream { self.reservation.try_resize(new_size) } } - -impl Stream for OrderedPartialAggregateStream { - type Item = Result; - - /// Entry point for the ordered partial aggregate state machine. - /// - /// See comments in [`OrderedPartialAggregateStream`] for high-level ideas. - /// - /// State transition graph: - /// - /// ```text - /// (start) - /// -> ReadingInput - /// The stream starts by polling ordered input and aggregating batches - /// into the ordered partial aggregate table. - /// - /// ReadingInput - /// -> ReadingInput - /// Aggregate one input batch. If the ordering proves some groups are - /// complete, yield one partial-state batch immediately, then continue - /// reading input. Otherwise continue directly with the next input batch. - /// -> DrainingFinal - /// Input was exhausted. Mark the table input as done so every remaining - /// group is safe to emit. - /// - /// DrainingFinal - /// -> DrainingFinal - /// One remaining partial-state batch was yielded; repeat to continue - /// draining the table. - /// -> Done - /// All remaining groups were emitted. - /// - /// Done - /// -> (end) - /// ``` - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - loop { - let cur_state = self - .state - .take() - .expect("OrderedPartialAggregateStream state should not be None"); - - let next_state = match cur_state { - state @ OrderedPartialAggregateState::ReadingInput { .. } => { - self.handle_reading_input(cx, state) - } - state @ OrderedPartialAggregateState::DrainingFinal { .. } => { - self.handle_draining_final(state) - } - state @ OrderedPartialAggregateState::Done => { - let _ = self.reservation.try_resize(0); - self.state = Some(state); - return Poll::Ready(None); - } - }; - - match next_state { - ControlFlow::Continue(next_state) => { - self.state = Some(next_state); - continue; - } - ControlFlow::Break((poll, next_state)) => { - self.state = Some(next_state); - return poll; - } - } - } - } -} - -impl RecordBatchStream for OrderedPartialAggregateStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} From f0b3a3e569e38cd706bb1f8d63a895f5a7932061 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 28 Jul 2026 23:10:49 +0300 Subject: [PATCH 2/3] update comments lint and fmt --- .../physical-plan/src/aggregates/mod.rs | 5 +- .../src/aggregates/ordered_partial_stream.rs | 65 ++++--------------- 2 files changed, 15 insertions(+), 55 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 382828e038114..d8a85eed1b34c 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4465,9 +4465,8 @@ mod tests { .with_session_config(session_config), ); - let mut stream: SendableRecordBatchStream = Box::pin( - OrderedPartialAggregateStream::new(&aggregate, &task_ctx, 0)?, - ); + let mut stream: SendableRecordBatchStream = + OrderedPartialAggregateStream::new(&aggregate, &task_ctx, 0)?.into_stream(); while let Some(result) = stream.next().await { if let Err(e) = result { diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 634fa79b46c52..606ae0c3e66c0 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -17,9 +17,7 @@ //! Partial aggregate stream for ordered group input. -use std::ops::ControlFlow; use std::sync::Arc; -use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; @@ -32,9 +30,9 @@ use super::AggregateExec; use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; use crate::aggregates::AggregateMode; use crate::aggregates::order::GroupOrdering; -use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput, SpillMetrics}; +use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; -use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; +use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// Partial aggregate stream for `InputOrderMode::Sorted` and /// `InputOrderMode::PartiallySorted`. @@ -121,23 +119,6 @@ pub(crate) struct OrderedPartialAggregateStream { table: Option>, } -/// See comments at `poll_next()` for details. -enum OrderedPartialAggregateState { - ReadingInput { - table: OrderedAggregateTable, - }, - DrainingFinal { - table: OrderedAggregateTable, - }, - Done, -} - -type OrderedPartialAggregatePoll = Poll>>; -type OrderedPartialAggregateStateTransition = ControlFlow< - (OrderedPartialAggregatePoll, OrderedPartialAggregateState), - OrderedPartialAggregateState, ->; - impl OrderedPartialAggregateStream { pub fn new( agg: &AggregateExec, @@ -178,7 +159,7 @@ impl OrderedPartialAggregateStream { reservation, baseline_metrics, reduction_factor, - table: Some(table) + table: Some(table), }) } @@ -232,16 +213,13 @@ impl OrderedPartialAggregateStream { .take() .expect("OrderedPartialAggregateStream state should not be None"); - self.handle_reading_input(&mut table, &mut emitter) - .await?; + self.handle_reading_input(&mut table, &mut emitter).await?; // Input has exhausted, move to the final draining stage. self.close_input(); table.input_done(); - let last_batch = self - .handle_draining_final(&mut table, &mut emitter) - .await?; + let last_batch = self.handle_draining_final(&mut table, &mut emitter).await?; // Clear memory before emitting last batch so we don't have to wait for next poll to clear { @@ -267,8 +245,6 @@ impl OrderedPartialAggregateStream { /// if the ordering proves any group is ready. /// /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. async fn handle_reading_input( &mut self, table: &mut OrderedAggregateTable, @@ -285,14 +261,11 @@ impl OrderedPartialAggregateStream { table.aggregate_batch(&batch)?; // Check memory reservation. See function comments for details. - match self.resize_or_take_state_batch(table)? { - Some(batch) => { - self.reduction_factor.add_part(batch.num_rows()); - drop(timer); - emitter.emit(batch).await; - continue; - } - None => {} + if let Some(batch) = self.resize_or_take_state_batch(table)? { + self.reduction_factor.add_part(batch.num_rows()); + drop(timer); + emitter.emit(batch).await; + continue; } let Some(batch) = table.next_output_batch()? else { @@ -351,9 +324,11 @@ impl OrderedPartialAggregateStream { /// `table.input_done()` has already made every remaining group safe to emit, /// so this state keeps draining until the table is empty. /// + /// Returns the last batch to emit so we can free all the state and memory before emitting, + /// and we won't need to hold while waiting for the next poll. + /// /// See comments at `poll_next()` for details. /// - /// Returns the last batch to emit so we can clear all the state before emitting so we dont need to hold while waiting for next poll to reset the state async fn handle_draining_final( &mut self, table: &mut OrderedAggregateTable, @@ -378,18 +353,4 @@ impl OrderedPartialAggregateStream { // was empty Ok(None) } - - fn resize_reservation_for_state( - &mut self, - state: &OrderedPartialAggregateState, - ) -> Result<()> { - let new_size = match state { - OrderedPartialAggregateState::ReadingInput { table } - | OrderedPartialAggregateState::DrainingFinal { table } => { - table.memory_size() - } - OrderedPartialAggregateState::Done => 0, - }; - self.reservation.try_resize(new_size) - } } From 5c3ef5c77f43d542385fa7100b18bdddda3ff85a Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 28 Jul 2026 23:18:05 +0300 Subject: [PATCH 3/3] update comments --- .../physical-plan/src/aggregates/ordered_partial_stream.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 606ae0c3e66c0..5db0fbe67f2fd 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -244,7 +244,7 @@ impl OrderedPartialAggregateStream { /// Consumes one ordered input batch, then immediately emits completed groups /// if the ordering proves any group is ready. /// - /// See comments at `poll_next()` for details. + /// See comments at [`Self::create_stream`] for details. async fn handle_reading_input( &mut self, table: &mut OrderedAggregateTable, @@ -327,7 +327,7 @@ impl OrderedPartialAggregateStream { /// Returns the last batch to emit so we can free all the state and memory before emitting, /// and we won't need to hold while waiting for the next poll. /// - /// See comments at `poll_next()` for details. + /// See comments at [`Self::create_stream`] for details. /// async fn handle_draining_final( &mut self,