From ea35a4d6ab6f74e18b325754ba33aacbe55f6fb5 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 20 Sep 2026 18:17:46 +0800 Subject: [PATCH 1/4] refactor(hash-aggr): share one spill context between the spilling aggregate streams --- .../src/aggregates/hash_stream.rs | 188 ++------------ .../physical-plan/src/aggregates/mod.rs | 1 + .../src/aggregates/ordered_final_stream.rs | 187 ++------------ .../src/aggregates/ordered_single_stream.rs | 209 ++------------- .../src/aggregates/single_stream.rs | 208 ++------------- .../physical-plan/src/aggregates/spill.rs | 240 ++++++++++++++++++ 6 files changed, 305 insertions(+), 728 deletions(-) create mode 100644 datafusion/physical-plan/src/aggregates/spill.rs diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 8e74e7bec78a1..bcb80c00cfaf9 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -27,13 +27,9 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{ DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, - internal_err, }; 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; use futures::stream::{Stream, StreamExt}; use super::AggregateExec; @@ -41,14 +37,11 @@ use super::aggregate_hash_table::{ AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker, PartialSkipMarker, }; -use super::ordered_final_stream::OrderedFinalAggregateStream; use super::skip_partial::SkipAggregationProbe; +use super::spill::AggregateSpill; use crate::metrics::{ BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput, SpillMetrics, }; -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, SendableRecordBatchStream, metrics}; @@ -179,28 +172,6 @@ pub(crate) struct PartialHashAggregateStream { hash_table: Option>, } -/// Spill configuration and accumulated runs for final hash aggregation. -/// -/// Each spill event drains all currently buffered groups, sorts their intermediate -/// states by the full group key, and writes them to one spill file. All files are -/// merged and replayed after the original input ends. -struct FinalSpillContext { - /// Aggregate configuration used to construct the final replay stream. - final_agg: AggregateExec, - /// Task context. - context: Arc, - /// Original partition index. - partition: usize, - /// Target batch size from configuration. - batch_size: usize, - /// Full group-key ordering kept by every spill file and the merged input. - spill_expr: LexOrdering, - /// Spill I/O and metrics manager. - spill_manager: SpillManager, - /// Spill runs waiting to be merged, they're all sorted by full group-by keys. - spills: Vec, -} - /// Hash aggregation is implemented in two stages: partial and final. This /// stream implements the final stage. /// @@ -227,142 +198,7 @@ pub(crate) struct FinalHashAggregateStream { /// This will be None when creating the stream hash_table: Option>, /// `None` if spilling is not supported by the configured `DiskManager`. - spill_context: Option>, -} - -impl FinalSpillContext { - fn new( - agg: &AggregateExec, - context: &Arc, - partition: usize, - batch_size: usize, - spill_schema: &SchemaRef, - spill_metrics: SpillMetrics, - ) -> Result { - let group_schema = agg.group_by.group_schema(&agg.input().schema())?; - let output_ordering = agg.cache.output_ordering(); - let spill_sort_exprs = - group_schema - .fields() - .iter() - .enumerate() - .map(|(idx, field)| { - let output_expr = Column::new(field.name(), idx); - let sort_options = output_ordering - .and_then(|ordering| ordering.get_sort_options(&output_expr)) - .unwrap_or_default(); - PhysicalSortExpr::new(Arc::new(output_expr), sort_options) - }); - let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { - return internal_err!("Final hash aggregate spill expression is empty"); - }; - - let spill_manager = SpillManager::new( - context.runtime_env(), - spill_metrics, - Arc::clone(spill_schema), - ) - .with_compression_type(context.session_config().spill_compression()); - - let mut final_agg = agg.clone(); - final_agg.input_order_mode = InputOrderMode::Sorted; - - Ok(Self { - final_agg, - context: Arc::clone(context), - partition, - batch_size, - spill_expr, - spill_manager, - spills: vec![], - }) - } - - fn has_spills(&self) -> bool { - !self.spills.is_empty() - } - - /// Sorts and spills the aggregated groups. Memory reservation should be updated - /// by the caller. - /// - /// Individual spill files are ordered by the `group by` keys. - /// - /// See [`FinalHashAggregateStream`] for spilling details. - fn spill_table( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { - let Some(batch) = hash_table.take_state_batch()? else { - return Ok(()); - }; - - let sorted_iter = - IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); - let spill_file = self - .spill_manager - .spill_record_batch_iter_and_return_max_batch_memory( - sorted_iter, - "FinalHashAggregateSpill", - )?; - - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Final hash aggregation produced an empty spill"); - }; - - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); - - Ok(()) - } - - /// Merges every sorted run, and do the aggregate evaluation with - /// [`OrderedFinalAggregateStream`] - fn into_replay_stream( - self, - baseline_metrics: &BaselineMetrics, - metrics: OrderedAggregateTableMetrics, - reservation: MemoryReservation, - ) -> Result { - let Self { - final_agg, - context, - partition, - batch_size, - spill_expr, - spill_manager, - spills, - } = self; - - let spill_schema = Arc::clone(spill_manager.schema()); - // The merge and replay table are two components of the same aggregate - // operator. Keep them under one consumer registration so a fair memory - // pool does not divide this operator's quota between its own phases. - let merge_reservation = reservation.new_empty(); - let merged = StreamingMergeBuilder::new() - .with_schema(spill_schema) - .with_spill_manager(spill_manager) - .with_sorted_spill_files(spills) - .with_expressions(&spill_expr) - .with_metrics(baseline_metrics.intermediate()) - .with_batch_size(batch_size) - .with_reservation(merge_reservation) - .with_replay_headroom() - .build()?; - let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( - &final_agg, - &context, - partition, - merged, - &InputOrderMode::Sorted, - baseline_metrics.clone(), - metrics, - None, - reservation, - )?; - Ok(Box::pin(replay)) - } + spill_context: Option>, } #[derive(PartialEq)] @@ -752,11 +588,13 @@ impl FinalHashAggregateStream { let can_spill = context.runtime_env().disk_manager.tmp_files_enabled(); let spill_context = if can_spill { - Some(Box::new(FinalSpillContext::new( + Some(Box::new(AggregateSpill::try_new( + "FinalHashAggregateSpill", agg, context, partition, batch_size, + &InputOrderMode::Linear, &input_schema, spill_metrics, )?)) @@ -830,7 +668,7 @@ impl FinalHashAggregateStream { /// Reserve memory for the current aggregate table. fn reservation_size_for_table( hash_table: &AggregateHashTable, - spill_context: Option<&FinalSpillContext>, + spill_context: Option<&AggregateSpill>, ) -> usize { let table_size = hash_table.memory_size(); if spill_context.is_some() { @@ -854,7 +692,7 @@ impl FinalHashAggregateStream { async fn consume_input( &mut self, hash_table: &mut AggregateHashTable, - spill_context: &mut Option>, + spill_context: &mut Option>, ) -> Result<()> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); @@ -902,7 +740,9 @@ impl FinalHashAggregateStream { // Go to the next state to perform spilling the aggregated // groups so far. - let result = spill_context.spill_table(hash_table); + let result = hash_table + .take_state_batch() + .and_then(|batch| spill_context.spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -932,14 +772,16 @@ impl FinalHashAggregateStream { async fn produce_output_from_spills( &mut self, mut hash_table: AggregateHashTable, - mut spill_context: Box, + mut spill_context: Box, mut emitter: TryEmitter, ) -> Result<()> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); // Input was exhausted after spilling. Spill the last in-memory run - spill_context.spill_table(&mut hash_table)?; + hash_table + .take_state_batch() + .and_then(|batch| spill_context.spill(batch))?; // Construct the ordered input used to merge all spill files. let mut output_stream = @@ -968,7 +810,7 @@ impl FinalHashAggregateStream { fn switch_to_ordered_final_stream( &mut self, hash_table: AggregateHashTable, - spill_context: Box, + spill_context: Box, ) -> Result { let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); drop(hash_table); diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 806eacb35de2d..5fa350b55d202 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -222,6 +222,7 @@ mod ordered_single_stream; mod partial_reduce_stream; mod single_stream; mod skip_partial; +mod spill; mod topk; /// Returns true if TopK aggregation data structures support the provided key and value types. diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index cc992eaa51181..24409a42a9c98 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -26,20 +26,15 @@ use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result, internal_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_physical_expr::PhysicalSortExpr; -use datafusion_physical_expr::expressions::Column; -use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::stream::{Stream, StreamExt}; use super::AggregateExec; use super::aggregate_hash_table::{ FinalMarker, OrderedAggregateTable, OrderedAggregateTableMetrics, }; +use super::spill::AggregateSpill; use crate::aggregates::AggregateMode; 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::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; @@ -70,30 +65,6 @@ pub(crate) struct OrderedFinalAggregateStream { state: Option, } -/// Spill configuration and accumulated runs for partially ordered final -/// aggregation. -/// -/// Each spill event drains all currently buffered groups, sorts their intermediate -/// states by the full group key, and writes them to one spill file. All files are -/// merged and replayed after the original input ends. -struct OrderedFinalSpillContext { - /// Aggregate configuration - agg: AggregateExec, - /// Task context - context: Arc, - /// Original partition index - partition: usize, - /// Target batch size from configuration - batch_size: usize, - /// Full group-key ordering, such ordering with be kept in: a) individual spill - /// files, b) order after final merging and streaming aggregate - spill_expr: LexOrdering, - /// Spill I/O and metrics manager. - spill_manager: SpillManager, - /// Fully sorted spill runs waiting to be merged. - spills: Vec, -} - /// See comments at `poll_next()` for details. enum OrderedFinalAggregateState { ReadingInput { @@ -101,18 +72,18 @@ enum OrderedFinalAggregateState { /// None if either /// - Disk Manager doesn't enable temporary file creation /// - The group keys are fully ordered, it's expected to use bounded memory - spill_context: Option>, + spill_context: Option>, }, Spilling { table: OrderedAggregateTable, - spill_context: Box, + spill_context: Box, }, ProducingOutput { table: OrderedAggregateTable, }, PreparingMergeInput { table: OrderedAggregateTable, - spill_context: Box, + spill_context: Box, }, MergingSpills { stream: SendableRecordBatchStream, @@ -126,140 +97,6 @@ type OrderedFinalAggregateStateTransition = ControlFlow< OrderedFinalAggregateState, >; -impl OrderedFinalSpillContext { - fn new( - agg: &AggregateExec, - context: &Arc, - partition: usize, - batch_size: usize, - input_order_mode: &InputOrderMode, - spill_schema: &SchemaRef, - spill_metrics: SpillMetrics, - ) -> Result { - let group_schema = agg.group_by.group_schema(spill_schema)?; - let output_ordering = agg.cache.output_ordering(); - let InputOrderMode::PartiallySorted(order_indices) = input_order_mode else { - return internal_err!("Ordered final spill requires partially ordered input"); - }; - let spill_indices = order_indices.iter().copied().chain( - (0..group_schema.fields().len()).filter(|idx| !order_indices.contains(idx)), - ); - let spill_sort_exprs = spill_indices.map(|idx| { - let field = group_schema.field(idx); - let output_expr = Column::new(field.name(), idx); - let sort_options = output_ordering - .and_then(|ordering| ordering.get_sort_options(&output_expr)) - .unwrap_or_default(); - PhysicalSortExpr::new(Arc::new(output_expr), sort_options) - }); - let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { - return internal_err!("Ordered final spill expression is empty"); - }; - - let spill_manager = SpillManager::new( - context.runtime_env(), - spill_metrics, - Arc::clone(spill_schema), - ) - .with_compression_type(context.session_config().spill_compression()); - - Ok(Self { - agg: agg.clone(), - context: Arc::clone(context), - partition, - batch_size, - spill_expr, - spill_manager, - spills: vec![], - }) - } - - fn has_spills(&self) -> bool { - !self.spills.is_empty() - } - - /// Sorts and spills the aggregated groups. Memory reservation should be updated - /// by the caller. - /// - /// Individual spill files are ordered by the `group by` keys. - /// - /// See [`OrderedFinalAggregateStream`] for spilling details. - fn spill_table( - &mut self, - table: &mut OrderedAggregateTable, - ) -> Result<()> { - let Some(batch) = table.take_state_batch()? else { - return Ok(()); - }; - - let sorted_iter = - IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); - let spill_file = self - .spill_manager - .spill_record_batch_iter_and_return_max_batch_memory( - sorted_iter, - "OrderedFinalAggregateSpill", - )?; - - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Ordered final aggregation produced an empty spill"); - }; - - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); - - Ok(()) - } - - /// Merges every sorted run and finalizes it through the fully ordered path. - fn into_replay_stream( - self, - baseline_metrics: &BaselineMetrics, - metrics: OrderedAggregateTableMetrics, - reservation: MemoryReservation, - ) -> Result { - let Self { - agg, - context, - partition, - batch_size, - spill_expr, - spill_manager, - spills, - } = self; - - let spill_schema = Arc::clone(spill_manager.schema()); - // The merge and replay table are two components of the same aggregate - // operator. Keep them under one consumer registration so a fair memory - // pool does not divide this operator's quota between its own phases. - let merge_reservation = reservation.new_empty(); - let merged = StreamingMergeBuilder::new() - .with_schema(spill_schema) - .with_spill_manager(spill_manager) - .with_sorted_spill_files(spills) - .with_expressions(&spill_expr) - .with_metrics(baseline_metrics.intermediate()) - .with_batch_size(batch_size) - .with_reservation(merge_reservation) - .with_replay_headroom() - .build()?; - let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( - &agg, - &context, - partition, - merged, - &InputOrderMode::Sorted, - baseline_metrics.clone(), - metrics, - None, - reservation, - )?; - Ok(Box::pin(replay)) - } -} - impl OrderedFinalAggregateStream { pub fn new( agg: &AggregateExec, @@ -342,7 +179,8 @@ impl OrderedFinalAggregateStream { let Some(spill_metrics) = spill_metrics else { return internal_err!("Spillable ordered final stream requires metrics"); }; - Some(Box::new(OrderedFinalSpillContext::new( + Some(Box::new(AggregateSpill::try_new( + "OrderedFinalAggregateSpill", agg, context, partition, @@ -390,7 +228,7 @@ impl OrderedFinalAggregateStream { /// Reserve memory for the current aggregate table. fn reservation_size_for_table( table: &OrderedAggregateTable, - spill_context: Option<&OrderedFinalSpillContext>, + spill_context: Option<&AggregateSpill>, ) -> usize { let table_size = table.memory_size(); if spill_context.is_some() { @@ -605,7 +443,9 @@ 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 mut result = table + .take_state_batch() + .and_then(|batch| spill_context.spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -654,7 +494,10 @@ impl OrderedFinalAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let replay = match spill_context.spill_table(&mut table) { + let replay = match table + .take_state_batch() + .and_then(|batch| spill_context.spill(batch)) + { Ok(()) => { let metrics = table.metrics(); drop(table); @@ -916,8 +759,10 @@ mod tests { use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::{min_max::min_udaf, sum::sum_udaf}; + use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::FutureExt; use futures::channel::mpsc; use std::collections::BTreeMap; diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index 40ba90729b55f..a88ebb99c7dae 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -23,24 +23,16 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_physical_expr::PhysicalSortExpr; -use datafusion_physical_expr::expressions::Column; -use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::stream::{Stream, StreamExt}; -use super::aggregate_hash_table::{ - OrderedAggregateTable, OrderedAggregateTableMetrics, SingleMarker, -}; -use super::ordered_final_stream::OrderedFinalAggregateStream; +use super::aggregate_hash_table::{OrderedAggregateTable, SingleMarker}; +use super::spill::AggregateSpill; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; 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::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; @@ -109,30 +101,6 @@ pub(crate) struct OrderedSingleAggregateStream { state: Option, } -/// Spill configuration and accumulated runs for partially ordered single -/// aggregation. -/// -/// Each spill event drains all currently buffered groups, sorts their intermediate -/// states by the full group key, and writes them to one spill file. All files are -/// merged and replayed after the original input ends. -struct OrderedSingleSpillContext { - /// Aggregate configuration used to construct the final replay stream. - final_agg: AggregateExec, - /// Task context - context: Arc, - /// Original partition index - partition: usize, - /// Target batch size from configuration - batch_size: usize, - /// Full group-key ordering, such ordering with be kept in: a) individual spill - /// files, b) order after final merging and streaming aggregate - spill_expr: LexOrdering, - /// Spill I/O and metrics manager. - spill_manager: SpillManager, - /// Fully sorted spill runs waiting to be merged. - spills: Vec, -} - /// See comments at `poll_next()` for details. enum OrderedSingleAggregateState { ReadingInput { @@ -140,18 +108,18 @@ enum OrderedSingleAggregateState { /// None if either /// - Disk Manager doesn't enable temporary file creation /// - The group keys are fully ordered, it's expected to use bounded memory - spill_context: Option>, + spill_context: Option>, }, Spilling { table: OrderedAggregateTable, - spill_context: Box, + spill_context: Box, }, ProducingOutput { table: OrderedAggregateTable, }, PreparingMergeInput { table: OrderedAggregateTable, - spill_context: Box, + spill_context: Box, }, MergingSpills { stream: SendableRecordBatchStream, @@ -169,157 +137,6 @@ type OrderedSingleAggregateStateTransition = ControlFlow< OrderedSingleAggregateState, >; -impl OrderedSingleSpillContext { - fn new( - agg: &AggregateExec, - context: &Arc, - partition: usize, - batch_size: usize, - input_order_mode: &InputOrderMode, - spill_schema: &SchemaRef, - spill_metrics: SpillMetrics, - ) -> Result { - let group_schema = agg.group_by.group_schema(&agg.input().schema())?; - let output_ordering = agg.cache.output_ordering(); - let InputOrderMode::PartiallySorted(order_indices) = input_order_mode else { - return internal_err!( - "Ordered single spill requires partially ordered input" - ); - }; - let spill_indices = order_indices.iter().copied().chain( - (0..group_schema.fields().len()).filter(|idx| !order_indices.contains(idx)), - ); - let spill_sort_exprs = spill_indices.map(|idx| { - let field = group_schema.field(idx); - let output_expr = Column::new(field.name(), idx); - let sort_options = output_ordering - .and_then(|ordering| ordering.get_sort_options(&output_expr)) - .unwrap_or_default(); - PhysicalSortExpr::new(Arc::new(output_expr), sort_options) - }); - let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { - return internal_err!("Ordered single spill expression is empty"); - }; - - let spill_manager = SpillManager::new( - context.runtime_env(), - spill_metrics, - Arc::clone(spill_schema), - ) - .with_compression_type(context.session_config().spill_compression()); - - // Spilled rows contain group keys and intermediate states. Replay must - // merge those states and evaluate the final aggregate values. - let mut final_agg = agg.clone(); - final_agg.mode = match agg.mode { - AggregateMode::Single => AggregateMode::Final, - AggregateMode::SinglePartitioned => AggregateMode::FinalPartitioned, - mode => { - return internal_err!( - "Ordered single aggregate spill cannot replay aggregate mode {mode:?}" - ); - } - }; - final_agg.group_by = Arc::new(agg.group_by.as_final()); - final_agg.input_order_mode = InputOrderMode::Sorted; - - Ok(Self { - final_agg, - context: Arc::clone(context), - partition, - batch_size, - spill_expr, - spill_manager, - spills: vec![], - }) - } - - fn has_spills(&self) -> bool { - !self.spills.is_empty() - } - - /// Sorts and spills the aggregated groups. Memory reservation should be updated - /// by the caller. - /// - /// Individual spill files are ordered by the `group by` keys. - /// - /// See [`OrderedSingleAggregateStream`] for spilling details. - fn spill_table( - &mut self, - table: &mut OrderedAggregateTable, - ) -> Result<()> { - let Some(batch) = table.take_state_batch()? else { - return Ok(()); - }; - - let sorted_iter = - IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); - let spill_file = self - .spill_manager - .spill_record_batch_iter_and_return_max_batch_memory( - sorted_iter, - "OrderedSingleAggregateSpill", - )?; - - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Ordered single aggregation produced an empty spill"); - }; - - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); - - Ok(()) - } - - /// Merges every sorted run and finalizes it through the fully ordered path. - fn into_replay_stream( - self, - baseline_metrics: &BaselineMetrics, - metrics: OrderedAggregateTableMetrics, - reservation: MemoryReservation, - ) -> Result { - let Self { - final_agg, - context, - partition, - batch_size, - spill_expr, - spill_manager, - spills, - } = self; - - let spill_schema = Arc::clone(spill_manager.schema()); - // The merge and replay table are two components of the same aggregate - // operator. Keep them under one consumer registration so a fair memory - // pool does not divide this operator's quota between its own phases. - let merge_reservation = reservation.new_empty(); - let merged = StreamingMergeBuilder::new() - .with_schema(spill_schema) - .with_spill_manager(spill_manager) - .with_sorted_spill_files(spills) - .with_expressions(&spill_expr) - .with_metrics(baseline_metrics.intermediate()) - .with_batch_size(batch_size) - .with_reservation(merge_reservation) - .with_replay_headroom() - .build()?; - let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( - &final_agg, - &context, - partition, - merged, - &InputOrderMode::Sorted, - baseline_metrics.clone(), - metrics, - None, - reservation, - )?; - Ok(Box::pin(replay)) - } -} - impl OrderedSingleAggregateStream { pub fn new( agg: &AggregateExec, @@ -357,7 +174,8 @@ impl OrderedSingleAggregateStream { matches!(agg.input_order_mode, InputOrderMode::PartiallySorted(_)) && context.runtime_env().disk_manager.tmp_files_enabled(); let spill_context = if can_spill { - Some(Box::new(OrderedSingleSpillContext::new( + Some(Box::new(AggregateSpill::try_new( + "OrderedSingleAggregateSpill", agg, context, partition, @@ -406,7 +224,7 @@ impl OrderedSingleAggregateStream { /// Reserve memory for the current aggregate table. fn reservation_size_for_table( table: &OrderedAggregateTable, - spill_context: Option<&OrderedSingleSpillContext>, + spill_context: Option<&AggregateSpill>, ) -> usize { let table_size = table.memory_size(); if spill_context.is_some() { @@ -590,7 +408,9 @@ impl OrderedSingleAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let mut result = spill_context.spill_table(&mut table); + let mut result = table + .take_state_batch() + .and_then(|batch| spill_context.spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -635,7 +455,10 @@ impl OrderedSingleAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let replay = match spill_context.spill_table(&mut table) { + let replay = match table + .take_state_batch() + .and_then(|batch| spill_context.spill(batch)) + { Ok(()) => { let metrics = table.metrics(); drop(table); diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 091e9fb940bce..919c25fc27c11 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -23,24 +23,18 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_physical_expr::PhysicalSortExpr; -use datafusion_physical_expr::expressions::Column; -use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::stream::{Stream, StreamExt}; use super::aggregate_hash_table::{ AggregateHashTable, OrderedAggregateTableMetrics, SingleMarker, }; -use super::ordered_final_stream::OrderedFinalAggregateStream; +use super::spill::AggregateSpill; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; 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::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; @@ -138,51 +132,22 @@ pub(crate) struct SingleHashAggregateStream { group_values_soft_limit: Option, } -/// Spill configuration and accumulated runs for single hash aggregation. -/// -/// Each spill event drains all currently buffered groups, sorts their intermediate -/// states by the full group key, and writes them to one spill file. All files are -/// merged and replayed after the original input ends. -struct SingleSpillContext { - /// Aggregate configuration used to construct the final replay stream. - /// - /// Spilled rows already contain evaluated group keys and intermediate - /// aggregate states. Replay must therefore use final aggregation semantics - /// and column-based group expressions rather than evaluating the raw input - /// expressions a second time. After the spill files are merged into ordered - /// input, this configuration is used to construct an - /// [`OrderedFinalAggregateStream`], and perform the final evaluation step. - final_agg: AggregateExec, - /// Task context. - context: Arc, - /// Original partition index. - partition: usize, - /// Target batch size from configuration. - batch_size: usize, - /// Full group-key ordering kept by every spill file and the merged input. - spill_expr: LexOrdering, - /// Spill I/O and metrics manager. - spill_manager: SpillManager, - /// Spill runs waiting to be merged, they're all sorted by full group-by keys. - spills: Vec, -} - /// See comments at `poll_next()` for details. enum SingleHashAggregateState { ReadingInput { hash_table: AggregateHashTable, - spill_context: Option>, + spill_context: Option>, }, Spilling { hash_table: AggregateHashTable, - spill_context: Box, + spill_context: Box, }, ProducingOutput { hash_table: AggregateHashTable, }, PreparingMergeInput { hash_table: AggregateHashTable, - spill_context: Box, + spill_context: Box, }, MergingSpills { stream: SendableRecordBatchStream, @@ -200,152 +165,6 @@ type SingleHashAggregateStateTransition = ControlFlow< SingleHashAggregateState, >; -impl SingleSpillContext { - fn new( - agg: &AggregateExec, - context: &Arc, - partition: usize, - batch_size: usize, - spill_schema: &SchemaRef, - spill_metrics: SpillMetrics, - ) -> Result { - let group_schema = agg.group_by.group_schema(&agg.input().schema())?; - let output_ordering = agg.cache.output_ordering(); - let spill_sort_exprs = - group_schema - .fields() - .iter() - .enumerate() - .map(|(idx, field)| { - let output_expr = Column::new(field.name(), idx); - let sort_options = output_ordering - .and_then(|ordering| ordering.get_sort_options(&output_expr)) - .unwrap_or_default(); - PhysicalSortExpr::new(Arc::new(output_expr), sort_options) - }); - let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { - return internal_err!("Single hash aggregate spill expression is empty"); - }; - - let spill_manager = SpillManager::new( - context.runtime_env(), - spill_metrics, - Arc::clone(spill_schema), - ) - .with_compression_type(context.session_config().spill_compression()); - - // See `SingleSpillContext::final_agg` comments for `final_agg`'s usage - let mut final_agg = agg.clone(); - final_agg.mode = match agg.mode { - AggregateMode::Single => AggregateMode::Final, - AggregateMode::SinglePartitioned => AggregateMode::FinalPartitioned, - mode => { - return internal_err!( - "Single hash aggregate spill cannot replay aggregate mode {mode:?}" - ); - } - }; - final_agg.group_by = Arc::new(agg.group_by.as_final()); - final_agg.input_order_mode = InputOrderMode::Sorted; - - Ok(Self { - final_agg, - context: Arc::clone(context), - partition, - batch_size, - spill_expr, - spill_manager, - spills: vec![], - }) - } - - fn has_spills(&self) -> bool { - !self.spills.is_empty() - } - - /// Sorts and spills the aggregated groups. Memory reservation should be updated - /// by the caller. - /// - /// Individual spill files are ordered by the `group by` keys. - /// - /// See [`SingleHashAggregateStream`] for spilling details. - fn spill_table( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { - let Some(batch) = hash_table.take_state_batch()? else { - return Ok(()); - }; - - let sorted_iter = - IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); - let spill_file = self - .spill_manager - .spill_record_batch_iter_and_return_max_batch_memory( - sorted_iter, - "SingleHashAggregateSpill", - )?; - - let Some((file, max_record_batch_memory)) = spill_file else { - return internal_err!("Single hash aggregation produced an empty spill"); - }; - - self.spills.push(SortedSpillFile { - file, - max_record_batch_memory, - }); - - Ok(()) - } - - /// Merges every sorted run, and do the aggregate evaluation with - /// [`OrderedFinalAggregateStream`] - fn into_replay_stream( - self, - baseline_metrics: &BaselineMetrics, - metrics: OrderedAggregateTableMetrics, - reservation: MemoryReservation, - ) -> Result { - let Self { - final_agg, - context, - partition, - batch_size, - spill_expr, - spill_manager, - spills, - } = self; - - let spill_schema = Arc::clone(spill_manager.schema()); - // The merge and replay table are two components of the same aggregate - // operator. Keep them under one consumer registration so a fair memory - // pool does not divide this operator's quota between its own phases. - let merge_reservation = reservation.new_empty(); - let merged = StreamingMergeBuilder::new() - .with_schema(spill_schema) - .with_spill_manager(spill_manager) - .with_sorted_spill_files(spills) - .with_expressions(&spill_expr) - .with_metrics(baseline_metrics.intermediate()) - .with_batch_size(batch_size) - .with_reservation(merge_reservation) - .with_replay_headroom() - .build()?; - let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( - &final_agg, - &context, - partition, - merged, - &InputOrderMode::Sorted, - baseline_metrics.clone(), - metrics, - None, - reservation, - )?; - Ok(Box::pin(replay)) - } -} - impl SingleHashAggregateStream { pub fn new( agg: &AggregateExec, @@ -381,11 +200,13 @@ impl SingleHashAggregateStream { let can_spill = context.runtime_env().disk_manager.tmp_files_enabled(); let spill_context = if can_spill { - Some(Box::new(SingleSpillContext::new( + Some(Box::new(AggregateSpill::try_new( + "SingleHashAggregateSpill", agg, context, partition, batch_size, + &InputOrderMode::Linear, &state_schema, spill_metrics, )?)) @@ -430,7 +251,7 @@ impl SingleHashAggregateStream { /// Reserve memory for the current aggregate table. fn reservation_size_for_table( hash_table: &AggregateHashTable, - spill_context: Option<&SingleSpillContext>, + spill_context: Option<&AggregateSpill>, ) -> usize { let table_size = hash_table.memory_size(); if spill_context.is_some() { @@ -562,7 +383,7 @@ impl SingleHashAggregateStream { fn close_input_and_prepare_output( &mut self, mut hash_table: AggregateHashTable, - spill_context: Option>, + spill_context: Option>, ) -> SingleHashAggregateStateTransition { self.close_input(); match spill_context { @@ -618,7 +439,9 @@ impl SingleHashAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let mut result = spill_context.spill_table(&mut hash_table); + let mut result = hash_table + .take_state_batch() + .and_then(|batch| spill_context.spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -664,7 +487,10 @@ impl SingleHashAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let replay = match spill_context.spill_table(&mut hash_table) { + let replay = match hash_table + .take_state_batch() + .and_then(|batch| spill_context.spill(batch)) + { Ok(()) => { let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); drop(hash_table); diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs new file mode 100644 index 0000000000000..8e21d343c2e65 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -0,0 +1,240 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spill and replay support shared by the grouped aggregation streams. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::sort_expr::LexOrdering; + +use super::aggregate_hash_table::OrderedAggregateTableMetrics; +use super::ordered_final_stream::OrderedFinalAggregateStream; +use super::{AggregateExec, AggregateMode}; +use crate::metrics::{BaselineMetrics, SpillMetrics}; +use crate::sorts::IncrementalSortIterator; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::spill_manager::SpillManager; +use crate::{InputOrderMode, SendableRecordBatchStream}; + +/// Spill configuration and accumulated runs of one grouped aggregation stream. +/// +/// Every aggregation stream that spills does so the same way. Each spill event +/// drains all currently buffered groups as intermediate state (see +/// `take_state_batch` on the aggregate tables), sorts them by the full group +/// key, and writes them to one spill file. After the original input ends, all +/// files are merged and replayed through an [`OrderedFinalAggregateStream`], +/// which merges the states and evaluates the final aggregate values. +pub(super) struct AggregateSpill { + /// Aggregate configuration used to construct the replay stream. + /// + /// Spilled rows already contain evaluated group keys and intermediate + /// aggregate states. Replay must therefore use final aggregation semantics + /// and column-based group expressions rather than evaluating the raw input + /// expressions a second time, so single-stage aggregates are rewritten to + /// their final counterpart here. + replay_agg: AggregateExec, + /// Task context. + context: Arc, + /// Original partition index. + partition: usize, + /// Target batch size from configuration. + batch_size: usize, + /// Full group-key ordering kept by every spill file and the merged input. + spill_expr: LexOrdering, + /// Spill I/O and metrics manager. + spill_manager: SpillManager, + /// Spill runs waiting to be merged, all sorted by `spill_expr`. + spills: Vec, + /// Describes this stream's spill requests, and prefixes its internal errors. + label: &'static str, +} + +impl AggregateSpill { + /// Creates the spill context of a stream, whose spill requests are described + /// as `label`. + /// + /// `input_order_mode` is the order of the stream's input: spill files are + /// sorted by the already ordered group columns first, followed by the + /// remaining ones, so that replay keeps the ordering the stream promised. + /// Fully sorted input aggregates in bounded memory and never spills. + /// + /// `spill_schema` is the schema of the intermediate state batches. + #[expect(clippy::too_many_arguments)] + pub(super) fn try_new( + label: &'static str, + agg: &AggregateExec, + context: &Arc, + partition: usize, + batch_size: usize, + input_order_mode: &InputOrderMode, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result { + let mut replay_agg = agg.clone(); + replay_agg.input_order_mode = InputOrderMode::Sorted; + let group_schema = match agg.mode { + AggregateMode::Final | AggregateMode::FinalPartitioned => { + agg.group_by.group_schema(spill_schema)? + } + AggregateMode::Single | AggregateMode::SinglePartitioned => { + replay_agg.mode = if agg.mode == AggregateMode::Single { + AggregateMode::Final + } else { + AggregateMode::FinalPartitioned + }; + replay_agg.group_by = Arc::new(agg.group_by.as_final()); + agg.group_by.group_schema(&agg.input().schema())? + } + mode => { + return internal_err!("{label}: cannot replay aggregate mode {mode:?}"); + } + }; + + let num_group_columns = group_schema.fields().len(); + let ordered_indices: &[usize] = match input_order_mode { + InputOrderMode::Linear => &[], + InputOrderMode::PartiallySorted(ordered_indices) => ordered_indices, + InputOrderMode::Sorted => { + return internal_err!("{label}: fully ordered input does not spill"); + } + }; + let spill_indices = ordered_indices + .iter() + .copied() + .chain((0..num_group_columns).filter(|idx| !ordered_indices.contains(idx))); + let output_ordering = agg.cache.output_ordering(); + let spill_sort_exprs = spill_indices.map(|idx| { + let output_expr = Column::new(group_schema.field(idx).name(), idx); + let sort_options = output_ordering + .and_then(|ordering| ordering.get_sort_options(&output_expr)) + .unwrap_or_default(); + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("{label}: spill expression is empty"); + }; + + let spill_manager = SpillManager::new( + context.runtime_env(), + spill_metrics, + Arc::clone(spill_schema), + ) + .with_compression_type(context.session_config().spill_compression()); + + Ok(Self { + replay_agg, + context: Arc::clone(context), + partition, + batch_size, + spill_expr, + spill_manager, + spills: vec![], + label, + }) + } + + pub(super) fn has_spills(&self) -> bool { + !self.spills.is_empty() + } + + /// Sorts `state_batch`, the intermediate state of all currently buffered + /// groups (`None` if there are no groups), and writes it as one spill file. + /// Memory reservation should be updated by the caller. + pub(super) fn spill(&mut self, state_batch: Option) -> Result<()> { + let Some(state_batch) = state_batch else { + return Ok(()); + }; + + let sorted_iter = IncrementalSortIterator::new( + state_batch, + self.spill_expr.clone(), + self.batch_size, + ); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + self.label, + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("{}: produced an empty spill", self.label); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) + } + + /// Merges every sorted run, and does the aggregate evaluation with + /// [`OrderedFinalAggregateStream`]. + pub(super) fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + metrics: OrderedAggregateTableMetrics, + reservation: MemoryReservation, + ) -> Result { + let Self { + replay_agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + label: _, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + // The merge and replay table are two components of the same aggregate + // operator. Keep them under one consumer registration so a fair memory + // pool does not divide this operator's quota between its own phases. + let merge_reservation = reservation.new_empty(); + let merged = StreamingMergeBuilder::new() + .with_schema(spill_schema) + .with_spill_manager(spill_manager) + .with_sorted_spill_files(spills) + .with_expressions(&spill_expr) + .with_metrics(baseline_metrics.intermediate()) + .with_batch_size(batch_size) + .with_reservation(merge_reservation) + .with_replay_headroom() + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &replay_agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + metrics, + None, + reservation, + )?; + Ok(Box::pin(replay)) + } +} From c70ddc2f8b88723dd5ca783bbfc03d326086501f Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 20 Sep 2026 19:36:32 +0800 Subject: [PATCH 2/4] docs: fix broken intra-doc links in aggregate streams --- datafusion/physical-plan/src/aggregates/hash_stream.rs | 2 +- datafusion/physical-plan/src/aggregates/ordered_final_stream.rs | 2 +- datafusion/physical-plan/src/aggregates/single_stream.rs | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index bcb80c00cfaf9..f7682e4ffeddf 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -136,7 +136,7 @@ use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// 3. Perform a sort-preserving merge of all spill files and feed the merged output /// into an ordered streaming aggregation, which ensures bounded memory usage and /// evaluates the final result. -/// - [`OrderedFinalAggregateStream`] is reused for the streaming aggregation. +/// - [`OrderedFinalAggregateStream`](super::ordered_final_stream::OrderedFinalAggregateStream) is reused for the streaming aggregation. pub(crate) struct PartialHashAggregateStream { /// Output schema: group columns followed by partial aggregate state columns. schema: SchemaRef, diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 24409a42a9c98..09e177d898bf6 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -52,7 +52,7 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// - Reserve the table footprint plus one `u32` sort index per buffered group. The /// extra index array is used in later sorting before spilling. /// - On memory pressure, materialize all group states into one batch. -/// - Use [`IncrementalSortIterator`] to compute the full-batch index, then +/// - Use [`IncrementalSortIterator`](crate::sorts::IncrementalSortIterator) to compute the full-batch index, then /// materialize and write one sorted `batch_size` slice at a time. The original /// batch and full index remain live until the run is written. /// - After input ends, merge the sorted runs and replay them through a fully diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 919c25fc27c11..1ca7da2cbf825 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -81,7 +81,7 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// 3. Perform a sort-preserving merge of all spill files and feed the merged output /// into an ordered streaming aggregation, which ensures bounded memory usage and /// evaluates the final result. -/// - [`OrderedFinalAggregateStream`] is reused for the streaming aggregation. +/// - [`OrderedFinalAggregateStream`](super::ordered_final_stream::OrderedFinalAggregateStream) is reused for the streaming aggregation. /// /// # Optimization: DISTINCT LIMIT Soft Limit /// From 6b2c5ce9bf3b615e686b9c575f688d5a52709c26 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Mon, 21 Sep 2026 19:52:48 +0800 Subject: [PATCH 3/4] Update datafusion/physical-plan/src/aggregates/spill.rs Co-authored-by: Yongting You <2010youy01@gmail.com> --- .../physical-plan/src/aggregates/spill.rs | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index 8e21d343c2e65..2158e7ca76c1b 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -53,6 +53,56 @@ pub(super) struct AggregateSpill { /// and column-based group expressions rather than evaluating the raw input /// expressions a second time, so single-stage aggregates are rewritten to /// their final counterpart here. + /// + /// # Example walkthrough + /// + /// This example walks through two key APIs of [`AggregateSpill`]: + /// - [`AggregateSpill::spill`] + /// - [`AggregateSpill::into_replay_stream`] + /// + /// ```txt + /// SELECT k, SUM(v) FROM t GROUP BY k + /// + /// -------------------- + /// Step 1: OOM round 1 + /// -------------------- + /// + /// First OOM: sort by k and write spill file 1 using `AggregateSpill::spill`. + /// + /// Buffered batch Spill file 1 (sorted) + /// k partial_sum k partial_sum + /// 1 3 1 3 + /// 3 4 -> 2 5 + /// 2 5 3 4 + /// + /// -------------------- + /// Step 2: OOM round 2 + /// -------------------- + /// After more input, a second OOM occurs: sort and spill similarly. + /// + /// Buffered batch Spill file 2 (sorted) + /// k partial_sum k partial_sum + /// 3 6 -> 1 2 + /// 1 2 3 6 + /// + /// ------------------------------------------ + /// Step 3: Global sort and final aggregation + /// ------------------------------------------ + /// 1. Construct a globally sorted aggregate stream via `SortPreservingMergeStream` + /// using the two previously sorted spill files. + /// 2. Build a final aggregation stream: + /// - The input is the SPM stream. + /// - It reuses `OrderedFinalAggregateStream` for processing. + /// - It returns the final aggregation result directly. + /// + /// SPM output Final aggregate output + /// k partial_sum k SUM(v) + /// 1 3 1 5 + /// 1 2 -> 2 5 + /// 2 5 3 10 + /// 3 4 + /// 3 6 + /// ``` replay_agg: AggregateExec, /// Task context. context: Arc, From 4a6cfd676b88229149c9657b5ecdb9ae59d2f0dd Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Mon, 21 Sep 2026 20:03:43 +0800 Subject: [PATCH 4/4] refactor: rename AggregateSpill::spill to sort_and_spill --- datafusion/physical-plan/src/aggregates/hash_stream.rs | 4 ++-- .../physical-plan/src/aggregates/ordered_final_stream.rs | 4 ++-- .../src/aggregates/ordered_single_stream.rs | 4 ++-- datafusion/physical-plan/src/aggregates/single_stream.rs | 4 ++-- datafusion/physical-plan/src/aggregates/spill.rs | 9 ++++++--- 5 files changed, 14 insertions(+), 11 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index f7682e4ffeddf..e7dc7ce8b1e7f 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -742,7 +742,7 @@ impl FinalHashAggregateStream { // groups so far. let result = hash_table .take_state_batch() - .and_then(|batch| spill_context.spill(batch)); + .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -781,7 +781,7 @@ impl FinalHashAggregateStream { // Input was exhausted after spilling. Spill the last in-memory run hash_table .take_state_batch() - .and_then(|batch| spill_context.spill(batch))?; + .and_then(|batch| spill_context.sort_and_spill(batch))?; // Construct the ordered input used to merge all spill files. let mut output_stream = diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 09e177d898bf6..3a09dbd97d704 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -445,7 +445,7 @@ impl OrderedFinalAggregateStream { let timer = elapsed_compute.timer(); let mut result = table .take_state_batch() - .and_then(|batch| spill_context.spill(batch)); + .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -496,7 +496,7 @@ impl OrderedFinalAggregateStream { let timer = elapsed_compute.timer(); let replay = match table .take_state_batch() - .and_then(|batch| spill_context.spill(batch)) + .and_then(|batch| spill_context.sort_and_spill(batch)) { Ok(()) => { let metrics = table.metrics(); diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index a88ebb99c7dae..02c3b2a368150 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -410,7 +410,7 @@ impl OrderedSingleAggregateStream { let timer = elapsed_compute.timer(); let mut result = table .take_state_batch() - .and_then(|batch| spill_context.spill(batch)); + .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -457,7 +457,7 @@ impl OrderedSingleAggregateStream { let timer = elapsed_compute.timer(); let replay = match table .take_state_batch() - .and_then(|batch| spill_context.spill(batch)) + .and_then(|batch| spill_context.sort_and_spill(batch)) { Ok(()) => { let metrics = table.metrics(); diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 1ca7da2cbf825..a185650e2f044 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -441,7 +441,7 @@ impl SingleHashAggregateStream { let timer = elapsed_compute.timer(); let mut result = hash_table .take_state_batch() - .and_then(|batch| spill_context.spill(batch)); + .and_then(|batch| spill_context.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -489,7 +489,7 @@ impl SingleHashAggregateStream { let timer = elapsed_compute.timer(); let replay = match hash_table .take_state_batch() - .and_then(|batch| spill_context.spill(batch)) + .and_then(|batch| spill_context.sort_and_spill(batch)) { Ok(()) => { let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index 2158e7ca76c1b..5ad9a200e4d8d 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -57,7 +57,7 @@ pub(super) struct AggregateSpill { /// # Example walkthrough /// /// This example walks through two key APIs of [`AggregateSpill`]: - /// - [`AggregateSpill::spill`] + /// - [`AggregateSpill::sort_and_spill`] /// - [`AggregateSpill::into_replay_stream`] /// /// ```txt @@ -67,7 +67,7 @@ pub(super) struct AggregateSpill { /// Step 1: OOM round 1 /// -------------------- /// - /// First OOM: sort by k and write spill file 1 using `AggregateSpill::spill`. + /// First OOM: sort by k and write spill file 1 using `AggregateSpill::sort_and_spill`. /// /// Buffered batch Spill file 1 (sorted) /// k partial_sum k partial_sum @@ -211,7 +211,10 @@ impl AggregateSpill { /// Sorts `state_batch`, the intermediate state of all currently buffered /// groups (`None` if there are no groups), and writes it as one spill file. /// Memory reservation should be updated by the caller. - pub(super) fn spill(&mut self, state_batch: Option) -> Result<()> { + pub(super) fn sort_and_spill( + &mut self, + state_batch: Option, + ) -> Result<()> { let Some(state_batch) = state_batch else { return Ok(()); };