From f442f34bd48e31b1932fc626e2b75a20f3e7e86e Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:19:39 +0300 Subject: [PATCH 1/6] chore(OrderedFinalAggregateStream): refactor to async generator --- .../physical-plan/src/aggregates/mod.rs | 2 +- .../src/aggregates/ordered_final_stream.rs | 521 +++++++----------- 2 files changed, 196 insertions(+), 327 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index db1eb951d6fbc..fa91fc7de4435 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::FinalHash(stream) => Box::pin(stream), StreamType::SingleHash(stream) => Box::pin(stream), StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), - StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), + StreamType::OrderedFinalAggregate(stream) => stream.into_stream(), StreamType::GroupedHash(stream) => Box::pin(stream), StreamType::GroupedPriorityQueue(stream) => Box::pin(stream), } diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 071c1e9011f41..2d1289b888f1f 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -17,14 +17,12 @@ //! Final aggregate stream for ordered partial-state input. -use std::ops::ControlFlow; use std::sync::Arc; -use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result, internal_err}; -use datafusion_execution::TaskContext; +use datafusion_execution::{async_try_stream, TaskContext, TryEmitter}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; @@ -39,7 +37,7 @@ use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; use crate::sorts::IncrementalSortIterator; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::SpillManager; -use crate::stream::EmptyRecordBatchStream; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// Final aggregate stream for `InputOrderMode::Sorted` and @@ -116,12 +114,6 @@ enum OrderedFinalAggregateState { Done, } -type OrderedFinalAggregatePoll = Poll>>; -type OrderedFinalAggregateStateTransition = ControlFlow< - (OrderedFinalAggregatePoll, OrderedFinalAggregateState), - OrderedFinalAggregateState, ->; - impl OrderedFinalSpillContext { fn new( agg: &AggregateExec, @@ -246,7 +238,7 @@ impl OrderedFinalSpillContext { group_by_metrics, None, )?; - Ok(Box::pin(replay)) + Ok(replay.into_stream()) } } @@ -356,18 +348,110 @@ impl OrderedFinalAggregateStream { }) } + 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 final aggregate state machine. + /// + /// See comments in [`OrderedFinalAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling ordered partial-state input and merging + /// those states into the ordered final aggregate table. + /// + /// ReadingInput + /// -> ReadingInput + /// Merge one input batch. If it fits in memory, optionally yield groups + /// proven complete by the input ordering, then read the next batch. + /// -> Spilling + /// The table cannot reserve enough memory. Move all current states into + /// one fully group-key-sorted spill run. + /// -> ProducingOutput + /// Input was exhausted without spilling. Mark every remaining group as + /// complete and produce its final result. + /// -> PreparingMergeInput + /// Input was exhausted after spilling. Spill the last in-memory run and + /// construct the ordered input used to merge all spill files. + /// + /// Spilling + /// -> ReadingInput + /// One sorted run was written; resume reading the original input. + /// + /// PreparingMergeInput + /// Spill the final in-memory run and build the input ordered replay stream. + /// -> MergingSpills + /// The final run was spilled and the ordered replay stream was built. + /// + /// MergingSpills + /// Aggregate the merged spill runs and emit final results. + /// -> MergingSpills + /// Forward one result batch from the fully ordered replay stream that + /// consumes the sort-preserving merge. + /// -> Done + /// The merged spill input was fully aggregated. + /// + /// ProducingOutput + /// -> ProducingOutput + /// One remaining final aggregate 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 { + loop { + let cur_state = self + .state + .take() + .expect("OrderedFinalAggregateStream state should not be None"); + + self.state = Some(match cur_state { + state @ OrderedFinalAggregateState::ReadingInput { .. } => { + self.handle_reading_input(&mut emitter, state).await? + } + state @ OrderedFinalAggregateState::Spilling { .. } => { + self.handle_spilling(state).await? + } + state @ OrderedFinalAggregateState::PreparingMergeInput { .. } => { + self.handle_preparing_merge_input(state)? + } + state @ OrderedFinalAggregateState::MergingSpills { .. } => { + self.handle_merging_spills(&mut emitter, state).await? + } + state @ OrderedFinalAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(&mut emitter, state).await? + } + state @ OrderedFinalAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Ok(()); + } + }); + } + }) + } + fn close_input(&mut self) { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); } - fn break_with_internal_err(message: &str) -> OrderedFinalAggregateStateTransition { - ControlFlow::Break(( - Poll::Ready(Some(internal_err!("{message}"))), - OrderedFinalAggregateState::Done, - )) - } - /// Reserve memory for the current aggregate table. fn reservation_size_for_table( table: &OrderedAggregateTable, @@ -388,94 +472,64 @@ impl OrderedFinalAggregateStream { /// 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<'_>, + emitter: &mut TryEmitter, original_state: OrderedFinalAggregateState, - ) -> OrderedFinalAggregateStateTransition { + ) -> Result { let OrderedFinalAggregateState::ReadingInput { mut table, spill_context, } = original_state else { - return Self::break_with_internal_err( - "Ordered final aggregate stream expected ReadingInput state", - ); + return internal_err!("Ordered final aggregate stream expected ReadingInput state"); }; - match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break(( - Poll::Pending, - OrderedFinalAggregateState::ReadingInput { - table, - spill_context, - }, - )), - Poll::Ready(Some(Ok(batch))) => { + match self.input.next().await.transpose()? { + Some(batch) => { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = table.aggregate_batch(&batch); + table.aggregate_batch(&batch)?; timer.done(); - if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::ReadingInput { - table, - spill_context, - }, - )); - } - // Check memory reservation, and potentially spill. let timer = elapsed_compute.timer(); let resize_result = - self.reservation - .try_resize(Self::reservation_size_for_table( - &table, - spill_context.as_deref(), - )); + self.reservation + .try_resize(Self::reservation_size_for_table( + &table, + spill_context.as_deref(), + )); timer.done(); match resize_result { Ok(()) => {} Err(e @ DataFusionError::ResourcesExhausted(_)) => { let Some(spill_context) = spill_context else { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::Done, - )); + return Err(e); }; if table.is_empty() { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::Done, - )); + return Err(e); } - return ControlFlow::Continue( - OrderedFinalAggregateState::Spilling { + return Ok(OrderedFinalAggregateState::Spilling { table, spill_context, - }, - ); + }); } Err(e) => { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::Done, - )); + return Err(e); } } let result = if spill_context - .as_ref() - .is_some_and(|spill_context| spill_context.has_spills()) + .as_ref() + .is_some_and(|spill_context| spill_context.has_spills()) { // Once one incomplete run is spilled, every remaining state // must participate in replay so no group is finalized twice. - Ok(None) + None } else { let timer = elapsed_compute.timer(); - let result = table.next_output_batch(); + let result = table.next_output_batch()?; timer.done(); result }; @@ -483,59 +537,35 @@ impl OrderedFinalAggregateStream { match result { // Some finalized groups can be emitted. Yield them, then // continue aggregating input in the current state. - Ok(Some(batch)) => { - if let Err(e) = - self.reservation - .try_resize(Self::reservation_size_for_table( - &table, - spill_context.as_deref(), - )) - { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::Done, - )); - } + Some(batch) => { + self.reservation + .try_resize(Self::reservation_size_for_table( + &table, + spill_context.as_deref(), + ))?; let next_state = OrderedFinalAggregateState::ReadingInput { table, spill_context, }; - ControlFlow::Break(( - Poll::Ready(Some(Ok( - batch.record_output(&self.baseline_metrics) - ))), - next_state, - )) + emitter.emit(batch).await; + + Ok(next_state) } // Can't do early emit, continue aggregating. - Ok(None) => { - ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput { + None => { + Ok(OrderedFinalAggregateState::ReadingInput { table, spill_context, }) } - Err(e) => ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::ReadingInput { - table, - spill_context, - }, - )), } } - Poll::Ready(Some(Err(e))) => ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::ReadingInput { - table, - spill_context, - }, - )), - Poll::Ready(None) => { + None => { self.close_input(); match spill_context { Some(spill_context) if spill_context.has_spills() => { - ControlFlow::Continue( + Ok( OrderedFinalAggregateState::PreparingMergeInput { table, spill_context, @@ -544,7 +574,7 @@ impl OrderedFinalAggregateStream { } _ => { table.input_done(); - ControlFlow::Continue( + Ok( OrderedFinalAggregateState::ProducingOutput { table }, ) } @@ -561,25 +591,22 @@ impl OrderedFinalAggregateStream { fn handle_spilling( &mut self, original_state: OrderedFinalAggregateState, - ) -> OrderedFinalAggregateStateTransition { + ) -> Result { let OrderedFinalAggregateState::Spilling { mut table, mut spill_context, } = original_state else { - return Self::break_with_internal_err( - "Ordered final aggregate stream expected Spilling state", + return internal_err!( + "Ordered final aggregate stream expected Spilling state" ); }; // Sanity check: it's impossible to OOM when the table is empty if table.is_empty() { - return ControlFlow::Break(( - Poll::Ready(Some(internal_err!( - "Ordered final aggregation entered Spilling with an empty table" - ))), - OrderedFinalAggregateState::Done, - )); + return internal_err!( + "Ordered final aggregation entered Spilling with an empty table" + ); } let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); @@ -589,23 +616,18 @@ impl OrderedFinalAggregateStream { // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. if let Err(e) = self.reservation.try_resize(table.memory_size()) { - result = - Err(e.context("Decreasing allocation after spilling should succeed")); + return Err(e.context("Decreasing allocation after spilling should succeed"))?; } timer.done(); - match result { - // Finished spilling the aggregate table, continue aggregating from input - Ok(()) => ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput { - table, - spill_context: Some(spill_context), - }), - Err(e) => ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::Done, - )), - } + result?; + + // Finished spilling the aggregate table, continue aggregating from input + Ok(OrderedFinalAggregateState::ReadingInput { + table, + spill_context: Some(spill_context), + }) } /// 1. Spills the last in-memory run. @@ -620,47 +642,35 @@ impl OrderedFinalAggregateStream { fn handle_preparing_merge_input( &mut self, original_state: OrderedFinalAggregateState, - ) -> OrderedFinalAggregateStateTransition { + ) -> Result { let OrderedFinalAggregateState::PreparingMergeInput { mut table, mut spill_context, } = original_state else { - return Self::break_with_internal_err( - "Ordered final aggregate stream expected PreparingMergeInput state", + return internal_err!( + "Ordered final aggregate stream expected PreparingMergeInput state" ); }; let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let replay = match spill_context.spill_table(&mut table) { - Ok(()) => { - let group_by_metrics = table.group_by_metrics(); - drop(table); - match self.reservation.try_resize(0) { - Ok(()) => (*spill_context).into_replay_stream( - &self.baseline_metrics, - group_by_metrics, - self.reservation.new_empty(), - ), - Err(e) => Err(e), - } - } - Err(e) => Err(e), + spill_context.spill_table(&mut table)?; + let replay = { + let group_by_metrics = table.group_by_metrics(); + drop(table); + self.reservation.try_resize(0)?; + (*spill_context).into_replay_stream( + &self.baseline_metrics, + group_by_metrics, + self.reservation.new_empty(), + )? }; timer.done(); - match replay { - Ok(stream) => { - ControlFlow::Continue(OrderedFinalAggregateState::MergingSpills { - stream, - }) - } - Err(e) => ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::Done, - )), - } + Ok(OrderedFinalAggregateState::MergingSpills { + stream, + }) } /// Forwards output from the fully ordered stream that consumes the merged @@ -669,32 +679,26 @@ impl OrderedFinalAggregateStream { /// See comments at `poll_next()` for details. /// /// Returns the next operator state with control flow decision. - fn handle_merging_spills( + async fn handle_merging_spills( &mut self, - cx: &mut Context<'_>, + emitter: &mut TryEmitter, original_state: OrderedFinalAggregateState, - ) -> OrderedFinalAggregateStateTransition { + ) -> Result { let OrderedFinalAggregateState::MergingSpills { mut stream } = original_state else { - return Self::break_with_internal_err( - "Ordered final aggregate stream expected MergingSpills state", + return internal_err!( + "Ordered final aggregate stream expected MergingSpills state" ); }; - match stream.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break(( - Poll::Pending, - OrderedFinalAggregateState::MergingSpills { stream }, - )), - Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( - Poll::Ready(Some(Ok(batch))), - OrderedFinalAggregateState::MergingSpills { stream }, - )), - Poll::Ready(Some(Err(e))) => ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::Done, - )), - Poll::Ready(None) => ControlFlow::Continue(OrderedFinalAggregateState::Done), + match stream.next().await.transpose()? { + Some(batch) => { + emitter.emit(batch).await; + Ok( + OrderedFinalAggregateState::MergingSpills { stream }, + ) + }, + None => Ok(OrderedFinalAggregateState::Done), } } @@ -706,178 +710,43 @@ impl OrderedFinalAggregateStream { /// See comments at `poll_next()` for details. /// /// Returns the next operator state with control flow decision. - fn handle_producing_output( + async fn handle_producing_output( &mut self, + emitter: &mut TryEmitter, original_state: OrderedFinalAggregateState, - ) -> OrderedFinalAggregateStateTransition { + ) -> Result { let OrderedFinalAggregateState::ProducingOutput { table } = original_state else { - return Self::break_with_internal_err( - "Ordered final aggregate stream expected ProducingOutput state", + return internal_err!( + "Ordered final aggregate stream expected ProducingOutput state" ); }; let mut table = table; let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = table.next_output_batch(); + let result = table.next_output_batch()?; timer.done(); match result { - Ok(Some(batch)) => { + Some(batch) => { let next_state = if table.is_empty() { drop(table); - if let Err(e) = self.reservation.try_resize(0) { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::Done, - )); - } + self.reservation.try_resize(0)?; OrderedFinalAggregateState::Done } else { - if let Err(e) = self.reservation.try_resize(table.memory_size()) { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::ProducingOutput { table }, - )); - } + self.reservation.try_resize(table.memory_size())?; OrderedFinalAggregateState::ProducingOutput { table } }; - ControlFlow::Break(( - Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), - next_state, - )) + emitter.emit(batch).await; + + Ok(next_state) } - Err(e) => ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::ProducingOutput { table }, - )), - Ok(None) => { + None => { drop(table); - let next_state = OrderedFinalAggregateState::Done; - if let Err(e) = self.reservation.try_resize(0) { - return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)); - } - ControlFlow::Continue(next_state) + self.reservation.try_resize(0)?; + Ok(OrderedFinalAggregateState::Done) } } } } - -impl Stream for OrderedFinalAggregateStream { - type Item = Result; - - /// Entry point for the ordered final aggregate state machine. - /// - /// See comments in [`OrderedFinalAggregateStream`] for high-level ideas. - /// - /// State transition graph: - /// - /// ```text - /// (start) - /// -> ReadingInput - /// The stream starts by polling ordered partial-state input and merging - /// those states into the ordered final aggregate table. - /// - /// ReadingInput - /// -> ReadingInput - /// Merge one input batch. If it fits in memory, optionally yield groups - /// proven complete by the input ordering, then read the next batch. - /// -> Spilling - /// The table cannot reserve enough memory. Move all current states into - /// one fully group-key-sorted spill run. - /// -> ProducingOutput - /// Input was exhausted without spilling. Mark every remaining group as - /// complete and produce its final result. - /// -> PreparingMergeInput - /// Input was exhausted after spilling. Spill the last in-memory run and - /// construct the ordered input used to merge all spill files. - /// - /// Spilling - /// -> ReadingInput - /// One sorted run was written; resume reading the original input. - /// - /// PreparingMergeInput - /// Spill the final in-memory run and build the input ordered replay stream. - /// -> MergingSpills - /// The final run was spilled and the ordered replay stream was built. - /// - /// MergingSpills - /// Aggregate the merged spill runs and emit final results. - /// -> MergingSpills - /// Forward one result batch from the fully ordered replay stream that - /// consumes the sort-preserving merge. - /// -> Done - /// The merged spill input was fully aggregated. - /// - /// ProducingOutput - /// -> ProducingOutput - /// One remaining final aggregate 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("OrderedFinalAggregateStream state should not be None"); - - let next_state = match cur_state { - state @ OrderedFinalAggregateState::ReadingInput { .. } => { - self.handle_reading_input(cx, state) - } - state @ OrderedFinalAggregateState::Spilling { .. } => { - self.handle_spilling(state) - } - state @ OrderedFinalAggregateState::PreparingMergeInput { .. } => { - self.handle_preparing_merge_input(state) - } - state @ OrderedFinalAggregateState::MergingSpills { .. } => { - self.handle_merging_spills(cx, state) - } - state @ OrderedFinalAggregateState::ProducingOutput { .. } => { - self.handle_producing_output(state) - } - state @ OrderedFinalAggregateState::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::Ready(Some(Err(e))), next_state)) => { - // Errors are terminal: discard all operator state and release - // its upstream input and memory reservation before returning. - drop(next_state); - self.close_input(); - self.reservation.free(); - self.state = Some(OrderedFinalAggregateState::Done); - return Poll::Ready(Some(Err(e))); - } - ControlFlow::Break((poll, next_state)) => { - self.state = Some(next_state); - return poll; - } - } - } - } -} - -impl RecordBatchStream for OrderedFinalAggregateStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} From 36a76e522a73fe32bfaaf0c5f7e1c059509eef64 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:04:34 +0300 Subject: [PATCH 2/6] chore(OrderedFinalAggregateStream): remove state --- .../physical-plan/src/aggregates/mod.rs | 4 +- .../src/aggregates/ordered_final_stream.rs | 415 ++++++------------ 2 files changed, 146 insertions(+), 273 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index fa91fc7de4435..f21e67543ea89 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -671,7 +671,7 @@ enum StreamType { /// Partial stage of aggregation for ordered input. OrderedPartialAggregate(OrderedPartialAggregateStream), /// Final stage of aggregation for ordered input. - OrderedFinalAggregate(OrderedFinalAggregateStream), + OrderedFinalAggregate(SendableRecordBatchStream), /// Hash aggregation reused for multiple stages /// /// Note this is being incrementally migrated to dedicated streams like @@ -698,7 +698,7 @@ impl From for SendableRecordBatchStream { StreamType::FinalHash(stream) => Box::pin(stream), StreamType::SingleHash(stream) => Box::pin(stream), StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), - StreamType::OrderedFinalAggregate(stream) => stream.into_stream(), + StreamType::OrderedFinalAggregate(stream) => stream, StreamType::GroupedHash(stream) => Box::pin(stream), StreamType::GroupedPriorityQueue(stream) => Box::pin(stream), } diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 2d1289b888f1f..4a9414c9bba34 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -22,8 +22,8 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result, internal_err}; -use datafusion_execution::{async_try_stream, TaskContext, TryEmitter}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -33,12 +33,12 @@ use super::AggregateExec; use super::aggregate_hash_table::{FinalMarker, OrderedAggregateTable}; use super::group_values::GroupByMetrics; use crate::aggregates::AggregateMode; -use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::sorts::IncrementalSortIterator; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::SpillManager; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; -use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; +use crate::{InputOrderMode, SendableRecordBatchStream}; /// Final aggregate stream for `InputOrderMode::Sorted` and /// `InputOrderMode::PartiallySorted`. @@ -64,7 +64,7 @@ pub(crate) struct OrderedFinalAggregateStream { input: SendableRecordBatchStream, reservation: MemoryReservation, baseline_metrics: BaselineMetrics, - state: Option, + spill_context: Option>, } /// Spill configuration and accumulated runs for partially ordered final @@ -91,29 +91,6 @@ struct OrderedFinalSpillContext { spills: Vec, } -/// See comments at `poll_next()` for details. -enum OrderedFinalAggregateState { - ReadingInput { - table: OrderedAggregateTable, - spill_context: Option>, - }, - Spilling { - table: OrderedAggregateTable, - spill_context: Box, - }, - ProducingOutput { - table: OrderedAggregateTable, - }, - PreparingMergeInput { - table: OrderedAggregateTable, - spill_context: Box, - }, - MergingSpills { - stream: SendableRecordBatchStream, - }, - Done, -} - impl OrderedFinalSpillContext { fn new( agg: &AggregateExec, @@ -228,7 +205,8 @@ impl OrderedFinalSpillContext { .with_batch_size(batch_size) .with_reservation(reservation) .build()?; - let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + + OrderedFinalAggregateStream::new_with_input_and_metrics( &agg, &context, partition, @@ -237,17 +215,16 @@ impl OrderedFinalSpillContext { baseline_metrics.clone(), group_by_metrics, None, - )?; - Ok(replay.into_stream()) + ) } } impl OrderedFinalAggregateStream { - pub fn new( + pub(crate) fn new( agg: &AggregateExec, context: &Arc, partition: usize, - ) -> Result { + ) -> Result { debug_assert!(matches!( agg.mode, AggregateMode::Final | AggregateMode::FinalPartitioned @@ -264,7 +241,7 @@ impl OrderedFinalAggregateStream { partition: usize, input: SendableRecordBatchStream, input_order_mode: &InputOrderMode, - ) -> Result { + ) -> Result { let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); let spill_metrics = SpillMetrics::new(&agg.metrics, partition); @@ -293,7 +270,7 @@ impl OrderedFinalAggregateStream { baseline_metrics: BaselineMetrics, group_by_metrics: GroupByMetrics, spill_metrics: Option, - ) -> Result { + ) -> Result { debug_assert!(matches!( agg.mode, AggregateMode::Final | AggregateMode::FinalPartitioned @@ -323,6 +300,11 @@ impl OrderedFinalAggregateStream { None }; + let reservation = + MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) + .with_can_spill(can_spill) + .register(context.memory_pool()); + let table = OrderedAggregateTable::::new_with_input_order( agg, &input_schema, @@ -331,33 +313,24 @@ impl OrderedFinalAggregateStream { input_order_mode, group_by_metrics, )?; - let reservation = - MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) - .with_can_spill(can_spill) - .register(context.memory_pool()); - Ok(Self { + let this = Self { schema, input, reservation, baseline_metrics, - state: Some(OrderedFinalAggregateState::ReadingInput { - table, - spill_context, - }), - }) - } + spill_context, + }; - pub(crate) fn into_stream(self) -> SendableRecordBatchStream { - let schema_clone = Arc::clone(&self.schema); + let schema_clone = Arc::clone(&this.schema); - let cloned_metrics = self.baseline_metrics.clone(); + let cloned_metrics = this.baseline_metrics.clone(); let stream = Box::pin(RecordBatchStreamAdapter::new( schema_clone, - self.create_stream(), + this.create_stream(table), )); - Box::pin(ObservedStream::new(stream, cloned_metrics, None)) + Ok(Box::pin(ObservedStream::new(stream, cloned_metrics, None))) } /// Entry point for the ordered final aggregate state machine. @@ -413,37 +386,38 @@ impl OrderedFinalAggregateStream { /// Done /// -> (end) /// ``` - fn create_stream(mut self) -> impl Stream> { + fn create_stream( + mut self, + mut table: OrderedAggregateTable, + ) -> impl Stream> { async_try_stream(|mut emitter| async move { - loop { - let cur_state = self - .state - .take() - .expect("OrderedFinalAggregateStream state should not be None"); - - self.state = Some(match cur_state { - state @ OrderedFinalAggregateState::ReadingInput { .. } => { - self.handle_reading_input(&mut emitter, state).await? - } - state @ OrderedFinalAggregateState::Spilling { .. } => { - self.handle_spilling(state).await? - } - state @ OrderedFinalAggregateState::PreparingMergeInput { .. } => { - self.handle_preparing_merge_input(state)? - } - state @ OrderedFinalAggregateState::MergingSpills { .. } => { - self.handle_merging_spills(&mut emitter, state).await? - } - state @ OrderedFinalAggregateState::ProducingOutput { .. } => { - self.handle_producing_output(&mut emitter, state).await? - } - state @ OrderedFinalAggregateState::Done => { - let _ = self.reservation.try_resize(0); - self.state = Some(state); - return Ok(()); - } - }); + let spilled = self.handle_reading_input(&mut table, &mut emitter).await?; + + let last_batch = if spilled { + let mut merging_spills_stream = + self.handle_preparing_merge_input(table)?; + + // Forwards output from the fully ordered stream that consumes the merged + // spill runs. + // + // See comments at `poll_next()` for details. + while let Some(batch) = merging_spills_stream.next().await.transpose()? { + emitter.emit(batch).await; + } + + None + } else { + table.input_done(); + self.handle_producing_output(table, &mut emitter).await? + }; + + self.reservation.try_resize(0)?; + + if let Some(last_batch) = last_batch { + emitter.emit(last_batch).await; } + + Ok(()) }) } @@ -474,113 +448,79 @@ impl OrderedFinalAggregateStream { /// Returns the next operator state with control flow decision. async fn handle_reading_input( &mut self, + table: &mut OrderedAggregateTable, emitter: &mut TryEmitter, - original_state: OrderedFinalAggregateState, - ) -> Result { - let OrderedFinalAggregateState::ReadingInput { - mut table, - spill_context, - } = original_state - else { - return internal_err!("Ordered final aggregate stream expected ReadingInput state"); - }; - - match self.input.next().await.transpose()? { - Some(batch) => { - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - table.aggregate_batch(&batch)?; - timer.done(); - - // Check memory reservation, and potentially spill. - let timer = elapsed_compute.timer(); - let resize_result = - self.reservation + ) -> Result { + while let Some(batch) = self.input.next().await.transpose()? { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + table.aggregate_batch(&batch)?; + timer.done(); + + // Check memory reservation, and potentially spill. + let timer = elapsed_compute.timer(); + let resize_result = + self.reservation .try_resize(Self::reservation_size_for_table( - &table, - spill_context.as_deref(), + table, + self.spill_context.as_deref(), )); - timer.done(); - match resize_result { - Ok(()) => {} - Err(e @ DataFusionError::ResourcesExhausted(_)) => { - let Some(spill_context) = spill_context else { - return Err(e); - }; - if table.is_empty() { - return Err(e); - } - return Ok(OrderedFinalAggregateState::Spilling { - table, - spill_context, - }); - } - Err(e) => { + timer.done(); + match resize_result { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + let Some(mut spill_context) = self.spill_context.take() else { + return Err(e); + }; + if table.is_empty() { return Err(e); } - } - let result = if spill_context - .as_ref() - .is_some_and(|spill_context| spill_context.has_spills()) - { - // Once one incomplete run is spilled, every remaining state - // must participate in replay so no group is finalized twice. - None - } else { - let timer = elapsed_compute.timer(); - let result = table.next_output_batch()?; - timer.done(); - result - }; - - match result { - // Some finalized groups can be emitted. Yield them, then - // continue aggregating input in the current state. - Some(batch) => { - self.reservation - .try_resize(Self::reservation_size_for_table( - &table, - spill_context.as_deref(), - ))?; - let next_state = OrderedFinalAggregateState::ReadingInput { - table, - spill_context, - }; - - emitter.emit(batch).await; - - Ok(next_state) - } - // Can't do early emit, continue aggregating. - None => { - Ok(OrderedFinalAggregateState::ReadingInput { - table, - spill_context, - }) - } + self.handle_spilling(table, &mut spill_context)?; + + self.spill_context = Some(spill_context); + + continue; } - } - None => { - self.close_input(); - match spill_context { - Some(spill_context) if spill_context.has_spills() => { - Ok( - OrderedFinalAggregateState::PreparingMergeInput { - table, - spill_context, - }, - ) - } - _ => { - table.input_done(); - Ok( - OrderedFinalAggregateState::ProducingOutput { table }, - ) - } + Err(e) => { + return Err(e); } } + + let result = if self + .spill_context + .as_ref() + .is_some_and(|spill_context| spill_context.has_spills()) + { + // Once one incomplete run is spilled, every remaining state + // must participate in replay so no group is finalized twice. + None + } else { + let timer = elapsed_compute.timer(); + let result = table.next_output_batch()?; + timer.done(); + result + }; + + // Some finalized groups can be emitted. Yield them, then + // continue aggregating input in the current state. + let Some(batch) = result else { + // Can't do early emit, continue aggregating. + + continue; + }; + self.reservation + .try_resize(Self::reservation_size_for_table( + table, + self.spill_context.as_deref(), + ))?; + + emitter.emit(batch).await; } + + self.close_input(); + + Ok(self.spill_context.as_ref().is_some_and(|c| c.has_spills())) } /// Sorts and spills one complete in-memory state run, then resumes input. @@ -590,18 +530,9 @@ impl OrderedFinalAggregateStream { /// Returns the next operator state with control flow decision. fn handle_spilling( &mut self, - original_state: OrderedFinalAggregateState, - ) -> Result { - let OrderedFinalAggregateState::Spilling { - mut table, - mut spill_context, - } = original_state - else { - return internal_err!( - "Ordered final aggregate stream expected Spilling state" - ); - }; - + table: &mut OrderedAggregateTable, + spill_context: &mut Box, + ) -> Result<()> { // Sanity check: it's impossible to OOM when the table is empty if table.is_empty() { return internal_err!( @@ -611,12 +542,12 @@ impl OrderedFinalAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let mut result = spill_context.spill_table(&mut table); + let result = spill_context.spill_table(table); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. if let Err(e) = self.reservation.try_resize(table.memory_size()) { - return Err(e.context("Decreasing allocation after spilling should succeed"))?; + return Err(e.context("Decreasing allocation after spilling should succeed")); } timer.done(); @@ -624,10 +555,7 @@ impl OrderedFinalAggregateStream { result?; // Finished spilling the aggregate table, continue aggregating from input - Ok(OrderedFinalAggregateState::ReadingInput { - table, - spill_context: Some(spill_context), - }) + Ok(()) } /// 1. Spills the last in-memory run. @@ -641,22 +569,16 @@ impl OrderedFinalAggregateStream { /// Returns the next operator state with control flow decision. fn handle_preparing_merge_input( &mut self, - original_state: OrderedFinalAggregateState, - ) -> Result { - let OrderedFinalAggregateState::PreparingMergeInput { - mut table, - mut spill_context, - } = original_state - else { - return internal_err!( - "Ordered final aggregate stream expected PreparingMergeInput state" - ); - }; - + mut table: OrderedAggregateTable, + ) -> Result { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); + let _timer = elapsed_compute.timer(); + let mut spill_context = self + .spill_context + .take() + .expect("must have spill context when merging input"); spill_context.spill_table(&mut table)?; - let replay = { + let replay = { let group_by_metrics = table.group_by_metrics(); drop(table); self.reservation.try_resize(0)?; @@ -666,40 +588,8 @@ impl OrderedFinalAggregateStream { self.reservation.new_empty(), )? }; - timer.done(); - - Ok(OrderedFinalAggregateState::MergingSpills { - stream, - }) - } - - /// Forwards output from the fully ordered stream that consumes the merged - /// spill runs. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - async fn handle_merging_spills( - &mut self, - emitter: &mut TryEmitter, - original_state: OrderedFinalAggregateState, - ) -> Result { - let OrderedFinalAggregateState::MergingSpills { mut stream } = original_state - else { - return internal_err!( - "Ordered final aggregate stream expected MergingSpills state" - ); - }; - match stream.next().await.transpose()? { - Some(batch) => { - emitter.emit(batch).await; - Ok( - OrderedFinalAggregateState::MergingSpills { stream }, - ) - }, - None => Ok(OrderedFinalAggregateState::Done), - } + Ok(replay) } /// Emits one batch after input is exhausted. @@ -712,41 +602,24 @@ impl OrderedFinalAggregateStream { /// Returns the next operator state with control flow decision. async fn handle_producing_output( &mut self, + mut table: OrderedAggregateTable, emitter: &mut TryEmitter, - original_state: OrderedFinalAggregateState, - ) -> Result { - let OrderedFinalAggregateState::ProducingOutput { table } = original_state else { - return internal_err!( - "Ordered final aggregate stream expected ProducingOutput state" - ); - }; - - let mut table = table; + ) -> Result> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = table.next_output_batch()?; - timer.done(); + let mut timer = elapsed_compute.timer(); - match result { - Some(batch) => { - let next_state = if table.is_empty() { - drop(table); - self.reservation.try_resize(0)?; - OrderedFinalAggregateState::Done - } else { - self.reservation.try_resize(table.memory_size())?; - OrderedFinalAggregateState::ProducingOutput { table } - }; + while let Some(batch) = table.next_output_batch()? { + if table.is_empty() { + return Ok(Some(batch)); + } - emitter.emit(batch).await; + self.reservation.try_resize(table.memory_size())?; - Ok(next_state) - } - None => { - drop(table); - self.reservation.try_resize(0)?; - Ok(OrderedFinalAggregateState::Done) - } + timer.done(); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); } + + Ok(None) } } From 325594fb56e4ce5c4907f9de582e8ae896615773 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:17:48 +0300 Subject: [PATCH 3/6] fix timer and comments --- .../src/aggregates/ordered_final_stream.rs | 91 +++++++++---------- 1 file changed, 41 insertions(+), 50 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 4a9414c9bba34..c27a808f902ca 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -333,7 +333,7 @@ impl OrderedFinalAggregateStream { Ok(Box::pin(ObservedStream::new(stream, cloned_metrics, None))) } - /// Entry point for the ordered final aggregate state machine. + /// Entry point for the ordered final aggregate stream /// /// See comments in [`OrderedFinalAggregateStream`] for high-level ideas. /// @@ -391,32 +391,27 @@ impl OrderedFinalAggregateStream { mut table: OrderedAggregateTable, ) -> impl Stream> { async_try_stream(|mut emitter| async move { - let spilled = self.handle_reading_input(&mut table, &mut emitter).await?; + let spilled = self.read_input(&mut table, &mut emitter).await?; - let last_batch = if spilled { + if spilled { let mut merging_spills_stream = - self.handle_preparing_merge_input(table)?; + self.prepare_merge_spills(table)?; // Forwards output from the fully ordered stream that consumes the merged // spill runs. // - // See comments at `poll_next()` for details. + // See comments at `create_stream()` for details. while let Some(batch) = merging_spills_stream.next().await.transpose()? { emitter.emit(batch).await; } - None + // Make sure empty + self.reservation.try_resize(0)?; } else { table.input_done(); - self.handle_producing_output(table, &mut emitter).await? + self.produce_output(table, &mut emitter).await?; }; - self.reservation.try_resize(0)?; - - if let Some(last_batch) = last_batch { - emitter.emit(last_batch).await; - } - Ok(()) }) } @@ -443,10 +438,10 @@ impl OrderedFinalAggregateStream { /// Consumes one ordered partial-state input batch, then immediately emits /// finalized groups if the ordering proves any group is ready. /// - /// See comments at `poll_next()` for details. + /// See comments at [`Self::create_stream`] for details. /// - /// Returns the next operator state with control flow decision. - async fn handle_reading_input( + /// Returns whether there are any spill files + async fn read_input( &mut self, table: &mut OrderedAggregateTable, emitter: &mut TryEmitter, @@ -455,17 +450,15 @@ impl OrderedFinalAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); table.aggregate_batch(&batch)?; - timer.done(); // Check memory reservation, and potentially spill. - let timer = elapsed_compute.timer(); let resize_result = self.reservation .try_resize(Self::reservation_size_for_table( table, self.spill_context.as_deref(), )); - timer.done(); + match resize_result { Ok(()) => {} Err(e @ DataFusionError::ResourcesExhausted(_)) => { @@ -476,15 +469,14 @@ impl OrderedFinalAggregateStream { return Err(e); } + timer.done(); self.handle_spilling(table, &mut spill_context)?; self.spill_context = Some(spill_context); continue; } - Err(e) => { - return Err(e); - } + Err(e) => return Err(e) } let result = if self @@ -496,25 +488,23 @@ impl OrderedFinalAggregateStream { // must participate in replay so no group is finalized twice. None } else { - let timer = elapsed_compute.timer(); - let result = table.next_output_batch()?; - timer.done(); - result + table.next_output_batch()? }; // Some finalized groups can be emitted. Yield them, then // continue aggregating input in the current state. let Some(batch) = result else { // Can't do early emit, continue aggregating. - continue; }; + self.reservation .try_resize(Self::reservation_size_for_table( table, self.spill_context.as_deref(), ))?; + drop(timer); emitter.emit(batch).await; } @@ -564,53 +554,54 @@ impl OrderedFinalAggregateStream { /// 3. Constructs a replay stream: an ordered aggregate stream over the fully /// ordered input constructed from the spills. /// - /// See comments at `poll_next()` for details. + /// See comments at [`Self::create_stream`] for details. /// - /// Returns the next operator state with control flow decision. - fn handle_preparing_merge_input( + /// Returns the merged spill stream + fn prepare_merge_spills( &mut self, mut table: OrderedAggregateTable, ) -> Result { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let _timer = elapsed_compute.timer(); + let mut spill_context = self .spill_context .take() .expect("must have spill context when merging input"); spill_context.spill_table(&mut table)?; - let replay = { - let group_by_metrics = table.group_by_metrics(); - drop(table); - self.reservation.try_resize(0)?; - (*spill_context).into_replay_stream( - &self.baseline_metrics, - group_by_metrics, - self.reservation.new_empty(), - )? - }; - - Ok(replay) + let group_by_metrics = table.group_by_metrics(); + drop(table); + self.reservation.try_resize(0)?; + (*spill_context).into_replay_stream( + &self.baseline_metrics, + group_by_metrics, + self.reservation.new_empty(), + ) } - /// Emits one batch after input is exhausted. + /// Emits all batches after input is exhausted. /// /// `table.input_done()` has already made every remaining group safe to emit, /// so this state keeps draining until the table is empty. /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - async fn handle_producing_output( + /// See comments at [`Self::create_stream`] for details. + async fn produce_output( &mut self, mut table: OrderedAggregateTable, emitter: &mut TryEmitter, - ) -> Result> { + ) -> Result<()> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let mut timer = elapsed_compute.timer(); while let Some(batch) = table.next_output_batch()? { if table.is_empty() { - return Ok(Some(batch)); + drop(table); + self.reservation.try_resize(0)?; + + drop(timer); + emitter.emit(batch).await; + + return Ok(()); } self.reservation.try_resize(table.memory_size())?; @@ -620,6 +611,6 @@ impl OrderedFinalAggregateStream { timer = elapsed_compute.timer(); } - Ok(None) + Ok(()) } } From e6f72a740a1d5f7ff20d502c8ff4a640b6fee2d0 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:24:37 +0300 Subject: [PATCH 4/6] cleanup --- .../physical-plan/src/aggregates/mod.rs | 4 +- .../src/aggregates/ordered_final_stream.rs | 58 ++++++++++--------- 2 files changed, 34 insertions(+), 28 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index f21e67543ea89..fa91fc7de4435 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -671,7 +671,7 @@ enum StreamType { /// Partial stage of aggregation for ordered input. OrderedPartialAggregate(OrderedPartialAggregateStream), /// Final stage of aggregation for ordered input. - OrderedFinalAggregate(SendableRecordBatchStream), + OrderedFinalAggregate(OrderedFinalAggregateStream), /// Hash aggregation reused for multiple stages /// /// Note this is being incrementally migrated to dedicated streams like @@ -698,7 +698,7 @@ impl From for SendableRecordBatchStream { StreamType::FinalHash(stream) => Box::pin(stream), StreamType::SingleHash(stream) => Box::pin(stream), StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), - StreamType::OrderedFinalAggregate(stream) => stream, + StreamType::OrderedFinalAggregate(stream) => stream.into_stream(), StreamType::GroupedHash(stream) => Box::pin(stream), StreamType::GroupedPriorityQueue(stream) => Box::pin(stream), } diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index c27a808f902ca..01f01351ba20d 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -64,6 +64,9 @@ pub(crate) struct OrderedFinalAggregateStream { input: SendableRecordBatchStream, reservation: MemoryReservation, baseline_metrics: BaselineMetrics, + + /// Will be taken on [`Self::aggregate`], we just keep it in the start until creating the stream itself + table: Option>, spill_context: Option>, } @@ -206,7 +209,7 @@ impl OrderedFinalSpillContext { .with_reservation(reservation) .build()?; - OrderedFinalAggregateStream::new_with_input_and_metrics( + Ok(OrderedFinalAggregateStream::new_with_input_and_metrics( &agg, &context, partition, @@ -215,7 +218,8 @@ impl OrderedFinalSpillContext { baseline_metrics.clone(), group_by_metrics, None, - ) + )? + .into_stream()) } } @@ -224,7 +228,7 @@ impl OrderedFinalAggregateStream { agg: &AggregateExec, context: &Arc, partition: usize, - ) -> Result { + ) -> Result { debug_assert!(matches!( agg.mode, AggregateMode::Final | AggregateMode::FinalPartitioned @@ -241,7 +245,7 @@ impl OrderedFinalAggregateStream { partition: usize, input: SendableRecordBatchStream, input_order_mode: &InputOrderMode, - ) -> Result { + ) -> Result { let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); let spill_metrics = SpillMetrics::new(&agg.metrics, partition); @@ -270,7 +274,7 @@ impl OrderedFinalAggregateStream { baseline_metrics: BaselineMetrics, group_by_metrics: GroupByMetrics, spill_metrics: Option, - ) -> Result { + ) -> Result { debug_assert!(matches!( agg.mode, AggregateMode::Final | AggregateMode::FinalPartitioned @@ -300,11 +304,6 @@ impl OrderedFinalAggregateStream { None }; - let reservation = - MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) - .with_can_spill(can_spill) - .register(context.memory_pool()); - let table = OrderedAggregateTable::::new_with_input_order( agg, &input_schema, @@ -314,23 +313,31 @@ impl OrderedFinalAggregateStream { group_by_metrics, )?; - let this = Self { + let reservation = + MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) + .with_can_spill(can_spill) + .register(context.memory_pool()); + + Ok(Self { schema, input, reservation, baseline_metrics, + table: Some(table), spill_context, - }; + }) + } - let schema_clone = Arc::clone(&this.schema); + pub(crate) fn into_stream(self) -> SendableRecordBatchStream { + let schema_clone = Arc::clone(&self.schema); - let cloned_metrics = this.baseline_metrics.clone(); + let cloned_metrics = self.baseline_metrics.clone(); let stream = Box::pin(RecordBatchStreamAdapter::new( schema_clone, - this.create_stream(table), + self.aggregate(), )); - Ok(Box::pin(ObservedStream::new(stream, cloned_metrics, None))) + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) } /// Entry point for the ordered final aggregate stream @@ -386,16 +393,15 @@ impl OrderedFinalAggregateStream { /// Done /// -> (end) /// ``` - fn create_stream( - mut self, - mut table: OrderedAggregateTable, - ) -> impl Stream> { + fn aggregate(mut self) -> impl Stream> { async_try_stream(|mut emitter| async move { + // Take the table on init since we want to control when the table is dropped and free memory. + let mut table = self.table.take().unwrap(); + let spilled = self.read_input(&mut table, &mut emitter).await?; if spilled { - let mut merging_spills_stream = - self.prepare_merge_spills(table)?; + let mut merging_spills_stream = self.prepare_merge_spills(table)?; // Forwards output from the fully ordered stream that consumes the merged // spill runs. @@ -438,7 +444,7 @@ impl OrderedFinalAggregateStream { /// Consumes one ordered partial-state input batch, then immediately emits /// finalized groups if the ordering proves any group is ready. /// - /// See comments at [`Self::create_stream`] for details. + /// See comments at [`Self::aggregate`] for details. /// /// Returns whether there are any spill files async fn read_input( @@ -476,7 +482,7 @@ impl OrderedFinalAggregateStream { continue; } - Err(e) => return Err(e) + Err(e) => return Err(e), } let result = if self @@ -554,7 +560,7 @@ impl OrderedFinalAggregateStream { /// 3. Constructs a replay stream: an ordered aggregate stream over the fully /// ordered input constructed from the spills. /// - /// See comments at [`Self::create_stream`] for details. + /// See comments at [`Self::aggregate`] for details. /// /// Returns the merged spill stream fn prepare_merge_spills( @@ -584,7 +590,7 @@ impl OrderedFinalAggregateStream { /// `table.input_done()` has already made every remaining group safe to emit, /// so this state keeps draining until the table is empty. /// - /// See comments at [`Self::create_stream`] for details. + /// See comments at [`Self::aggregate`] for details. async fn produce_output( &mut self, mut table: OrderedAggregateTable, From 478a772b4c87713a70f3cf49db701f4e7afa7871 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:27:47 +0300 Subject: [PATCH 5/6] rename --- .../physical-plan/src/aggregates/ordered_final_stream.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 01f01351ba20d..e3418b1937aea 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -398,7 +398,7 @@ impl OrderedFinalAggregateStream { // Take the table on init since we want to control when the table is dropped and free memory. let mut table = self.table.take().unwrap(); - let spilled = self.read_input(&mut table, &mut emitter).await?; + let spilled = self.consume_input(&mut table, &mut emitter).await?; if spilled { let mut merging_spills_stream = self.prepare_merge_spills(table)?; @@ -447,7 +447,7 @@ impl OrderedFinalAggregateStream { /// See comments at [`Self::aggregate`] for details. /// /// Returns whether there are any spill files - async fn read_input( + async fn consume_input( &mut self, table: &mut OrderedAggregateTable, emitter: &mut TryEmitter, From a7d2853606d174ff3311f4c1893fcb007f7bec53 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:29:24 +0300 Subject: [PATCH 6/6] remove done --- .../physical-plan/src/aggregates/ordered_final_stream.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index e3418b1937aea..9b66f1d899f7f 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -537,7 +537,7 @@ impl OrderedFinalAggregateStream { } let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); + let _timer = elapsed_compute.timer(); let result = spill_context.spill_table(table); // Spilling shrinks the aggregate table and releases its accumulated @@ -546,8 +546,6 @@ impl OrderedFinalAggregateStream { return Err(e.context("Decreasing allocation after spilling should succeed")); } - timer.done(); - result?; // Finished spilling the aggregate table, continue aggregating from input