From 62283d4a55bc32a8ae305d5dc78b62d679c79538 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Fri, 17 Jul 2026 11:41:18 +0800 Subject: [PATCH 1/3] cp --- .../aggregate_hash_table/common_ordered.rs | 59 +- .../ordered_final_table.rs | 8 +- .../ordered_partial_table.rs | 8 +- .../physical-plan/src/aggregates/mod.rs | 17 +- .../src/aggregates/order/full.rs | 5 + .../physical-plan/src/aggregates/order/mod.rs | 14 + .../src/aggregates/order/partial.rs | 6 + .../src/aggregates/ordered_final_stream.rs | 616 ++++++++++++++++-- .../src/aggregates/ordered_partial_stream.rs | 146 ++++- 9 files changed, 785 insertions(+), 94 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index c83303c51d6e8..2293e7b1b8e89 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -81,6 +81,10 @@ pub(in crate::aggregates) struct OrderedAggregateTable { /// Output schema: group columns followed by aggregate state or final values. pub(super) output_schema: SchemaRef, + /// Intermediate-state schema used when memory pressure requires the table + /// to pass through or spill its current state. + pub(super) state_schema: SchemaRef, + /// Maximum rows per emitted output batch, from config `batch_size`. pub(super) batch_size: usize, @@ -129,13 +133,14 @@ impl OrderedAggregateTable { )] pub(super) fn new_for_mode( agg: &AggregateExec, - partition: usize, input_schema: &SchemaRef, output_schema: SchemaRef, + state_schema: SchemaRef, batch_size: usize, input_order_mode: &InputOrderMode, aggregate_mode: &AggregateMode, filters: Vec>>, + group_by_metrics: GroupByMetrics, ) -> Result { assert_or_internal_err!( batch_size > 0, @@ -168,8 +173,9 @@ impl OrderedAggregateTable { Ok(Self { output_schema, + state_schema, batch_size, - group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), + group_by_metrics, buffer: OrderedAggregateTableBuffer { group_by: Arc::clone(&agg.group_by), group_ordering, @@ -217,9 +223,19 @@ impl OrderedAggregateTable { self.buffer.group_ordering.input_done(); } + /// Returns the ordering state used to decide how memory pressure is handled. + pub(in crate::aggregates) fn group_ordering(&self) -> &GroupOrdering { + &self.buffer.group_ordering + } + + /// Number of groups currently buffered. + pub(in crate::aggregates) fn num_groups(&self) -> usize { + self.buffer.group_values.len() + } + /// Check if there is zero groups accumulated so far. pub(in crate::aggregates) fn is_empty(&self) -> bool { - self.buffer.group_values.is_empty() + self.num_groups() == 0 } /// All internal buffer's memory size. @@ -234,6 +250,43 @@ impl OrderedAggregateTable { + self.buffer.group_indices.allocated_size() } + pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics { + self.group_by_metrics.clone() + } + + /// Takes every intermediate aggregate state and resets the table so it can + /// continue with a new ordered input segment. + /// + /// Unlike normal ordered emission, this operation is allowed to take the + /// active (incomplete) groups. Partial aggregation can pass those states to + /// its final stage, while final aggregation sorts and spills them before + /// replay. + pub(in crate::aggregates) fn take_state_batch( + &mut self, + ) -> Result> { + if self.buffer.group_values.is_empty() { + return Ok(None); + } + + let mut output = self.buffer.group_values.emit(EmitTo::All)?; + for acc in &mut self.buffer.accumulators { + output.extend(acc.state(EmitTo::All)?); + } + + let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the + // key/index buffers too so the memory reservation can be released + // before the batch is passed downstream or sorted for spilling. + self.buffer.group_values.clear_shrink(0); + self.buffer.group_indices.clear(); + self.buffer.group_indices.shrink_to_fit(); + self.buffer.group_ordering.reset(); + + Ok(Some(batch)) + } + /// Returns the [`EmitTo`], clamped to the specified batch size /// /// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index b7e3fd38edf25..fd064ebffec12 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -19,12 +19,15 @@ //! //! See comments in [`super::ordered_partial_table`] for details. +use std::sync::Arc; + use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; use crate::InputOrderMode; use crate::aggregates::aggregate_hash_table::FinalMarker; +use crate::aggregates::group_values::GroupByMetrics; use crate::aggregates::{AggregateExec, AggregateMode}; use super::common_ordered::OrderedAggregateTable; @@ -41,21 +44,22 @@ use super::common_ordered::OrderedAggregateTable; impl OrderedAggregateTable { pub(in crate::aggregates) fn new_with_input_order( agg: &AggregateExec, - partition: usize, input_schema: &SchemaRef, output_schema: SchemaRef, batch_size: usize, input_order_mode: &InputOrderMode, + group_by_metrics: GroupByMetrics, ) -> Result { Self::new_for_mode( agg, - partition, input_schema, output_schema, + Arc::clone(input_schema), batch_size, input_order_mode, &AggregateMode::Final, vec![None; agg.aggr_expr.len()], + group_by_metrics, ) } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs index 033c14056a419..a04e4dda8fb39 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -29,12 +29,15 @@ //! The implementation is separated from other aggregate tables because this //! execution path is likely to be optimized further in the future. +use std::sync::Arc; + use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; use crate::aggregates::{ AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker, + group_values::GroupByMetrics, }; use super::common_ordered::OrderedAggregateTable; @@ -56,15 +59,18 @@ impl OrderedAggregateTable { batch_size: usize, ) -> Result { let input_schema = agg.input().schema(); + let state_schema = Arc::clone(&output_schema); + let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); Self::new_for_mode( agg, - partition, &input_schema, output_schema, + state_schema, batch_size, &agg.input_order_mode, &AggregateMode::Partial, agg.filter_expr.iter().cloned().collect(), + group_by_metrics, ) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 48254edc6a5f0..4e8772645ded6 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1104,12 +1104,10 @@ impl AggregateExec { && self.limit_options_supported_by_hash_stream() } - fn should_use_ordered_partial_aggregate_stream(&self, context: &TaskContext) -> bool { - // TODO: implement memory-limited path and remove this limitation - if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { - return false; - } - + fn should_use_ordered_partial_aggregate_stream( + &self, + _context: &TaskContext, + ) -> bool { self.mode == AggregateMode::Partial && self.input_order_mode != InputOrderMode::Linear && !self.group_by.is_true_no_grouping() @@ -1160,12 +1158,7 @@ impl AggregateExec { && self.group_by.is_single() } - fn should_use_ordered_final_aggregate_stream(&self, context: &TaskContext) -> bool { - // TODO: implement memory-limited path and remove this limitation - if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { - return false; - } - + fn should_use_ordered_final_aggregate_stream(&self, _context: &TaskContext) -> bool { matches!( self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned diff --git a/datafusion/physical-plan/src/aggregates/order/full.rs b/datafusion/physical-plan/src/aggregates/order/full.rs index eb98611f79dfb..ca818d6a2d598 100644 --- a/datafusion/physical-plan/src/aggregates/order/full.rs +++ b/datafusion/physical-plan/src/aggregates/order/full.rs @@ -115,6 +115,11 @@ impl GroupOrderingFull { self.state = State::Complete; } + /// Starts tracking a new fully ordered input segment. + pub fn reset(&mut self) { + self.state = State::Start; + } + /// Called when new groups are added in a batch. See documentation /// on [`super::GroupOrdering::new_groups`] pub fn new_groups(&mut self, total_num_groups: usize) { diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 97fbd519c825c..259411b00b697 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -93,6 +93,20 @@ impl GroupOrdering { } } + /// Resets the ordering state while preserving the configured ordering mode. + /// + /// Ordered partial aggregation uses this after passing intermediate states + /// downstream, and ordered final aggregation uses it after spilling a run. + /// In both cases the hash table is empty and can start tracking the next + /// input batch from a fresh ordering state. + pub fn reset(&mut self) { + match self { + GroupOrdering::None => {} + GroupOrdering::Partial(partial) => partial.reset(), + GroupOrdering::Full(full) => full.reset(), + } + } + /// Removes the first `n` groups from the internal state, shifting all /// existing indexes down by `n`. pub fn remove_groups(&mut self, n: usize) { diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs index 476551a7ca210..1603bb6d079be 100644 --- a/datafusion/physical-plan/src/aggregates/order/partial.rs +++ b/datafusion/physical-plan/src/aggregates/order/partial.rs @@ -186,6 +186,12 @@ impl GroupOrderingPartial { }; } + /// Starts tracking a new ordered input segment with the same sort-key + /// columns. + pub fn reset(&mut self) { + self.state = State::Start; + } + fn updated_sort_key( current_sort: usize, sort_key: Option>, diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 89653e05ab4c7..51d90b613e312 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -23,15 +23,22 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +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}; +use super::group_values::GroupByMetrics; 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}; @@ -39,6 +46,21 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// `InputOrderMode::PartiallySorted`. /// /// See comments at [`super::ordered_partial_stream`] for details. +/// +/// # Spilling +/// +/// This section is only for implementation notes, for background, see [`super::ordered_partial_stream`] +/// +/// For partially sorted input, spilling works as follows: +/// +/// - 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 +/// 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 +/// ordered final aggregate stream. pub(crate) struct OrderedFinalAggregateStream { schema: SchemaRef, input: SendableRecordBatchStream, @@ -47,13 +69,46 @@ pub(crate) struct OrderedFinalAggregateStream { state: Option, } +/// Spill configuration and accumulated runs for partially ordered final +/// aggregation. Each file is one fully group-key-sorted intermediate-state run; +/// all runs are replayed together after the original input is exhausted. +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 { table: OrderedAggregateTable, + spill_context: Option>, }, - DrainingFinal { + Spilling { table: OrderedAggregateTable, + spill_context: Box, + }, + ProducingOutput { + table: OrderedAggregateTable, + }, + PreparingMergeInput { + table: OrderedAggregateTable, + spill_context: Box, + }, + MergingSpills { + stream: SendableRecordBatchStream, }, Done, } @@ -64,6 +119,135 @@ 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() + } + + /// Sort and spill the aggregated groups. Memory reservation should be cleared + /// by the caller of this function. + /// + /// See [`OrderedFinalAggregateStream`] for spilling details. + /// + /// It must not be called without rows to spill in the aggregate table, otherwise + /// it returns `Ok(false)`, and the caller propagates the error. + fn spill_table( + &mut self, + table: &mut OrderedAggregateTable, + ) -> Result { + let Some(batch) = table.take_state_batch()? else { + return Ok(false); + }; + + 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(true) + } + + /// Merges every sorted run and finalizes it through the fully ordered path. + fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + group_by_metrics: GroupByMetrics, + reservation: MemoryReservation, + ) -> Result { + let Self { + agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + 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(reservation) + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + group_by_metrics, + None, + )?; + Ok(Box::pin(replay)) + } +} + impl OrderedFinalAggregateStream { pub fn new( agg: &AggregateExec, @@ -86,6 +270,35 @@ impl OrderedFinalAggregateStream { partition: usize, input: SendableRecordBatchStream, input_order_mode: &InputOrderMode, + ) -> 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); + Self::new_with_input_and_metrics( + agg, + context, + partition, + input, + input_order_mode, + baseline_metrics, + group_by_metrics, + Some(spill_metrics), + ) + } + + #[expect( + clippy::too_many_arguments, + reason = "keeps replay metric reuse explicit" + )] + fn new_with_input_and_metrics( + agg: &AggregateExec, + context: &Arc, + partition: usize, + input: SendableRecordBatchStream, + input_order_mode: &InputOrderMode, + baseline_metrics: BaselineMetrics, + group_by_metrics: GroupByMetrics, + spill_metrics: Option, ) -> Result { debug_assert!(matches!( agg.mode, @@ -96,21 +309,37 @@ impl OrderedFinalAggregateStream { let schema = Arc::clone(&agg.schema); let input_schema = input.schema(); let batch_size = context.session_config().batch_size(); - let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); - // Preserve the existing aggregate metric surface for this plan node. - let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let can_spill = matches!(input_order_mode, InputOrderMode::PartiallySorted(_)) + && context.runtime_env().disk_manager.tmp_files_enabled(); + let spill_context = if can_spill { + let Some(spill_metrics) = spill_metrics else { + return internal_err!("Spillable ordered final stream requires metrics"); + }; + Some(Box::new(OrderedFinalSpillContext::new( + agg, + context, + partition, + batch_size, + input_order_mode, + &input_schema, + spill_metrics, + )?)) + } else { + None + }; let table = OrderedAggregateTable::::new_with_input_order( agg, - partition, &input_schema, Arc::clone(&schema), batch_size, input_order_mode, + group_by_metrics, )?; let reservation = MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) + .with_can_spill(can_spill) .register(context.memory_pool()); Ok(Self { @@ -118,7 +347,10 @@ impl OrderedFinalAggregateStream { input, reservation, baseline_metrics, - state: Some(OrderedFinalAggregateState::ReadingInput { table }), + state: Some(OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }), }) } @@ -127,6 +359,20 @@ impl OrderedFinalAggregateStream { self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); } + /// Reserve memory for the current aggregate table. + fn reservation_size_for_table( + table: &OrderedAggregateTable, + spill_context: Option<&OrderedFinalSpillContext>, + ) -> usize { + let table_size = table.memory_size(); + if spill_context.is_some() { + // See `OrderedFinalAggregateStream` comments for how is it estimated + table_size.saturating_add(table.num_groups().saturating_mul(size_of::())) + } else { + table_size + } + } + /// Consumes one ordered partial-state input batch, then immediately emits /// finalized groups if the ordering proves any group is ready. /// @@ -138,7 +384,10 @@ impl OrderedFinalAggregateStream { cx: &mut Context<'_>, original_state: OrderedFinalAggregateState, ) -> OrderedFinalAggregateStateTransition { - let OrderedFinalAggregateState::ReadingInput { mut table } = original_state + let OrderedFinalAggregateState::ReadingInput { + mut table, + spill_context, + } = original_state else { unreachable!("expected reading input state") }; @@ -146,7 +395,10 @@ impl OrderedFinalAggregateStream { match self.input.poll_next_unpin(cx) { Poll::Pending => ControlFlow::Break(( Poll::Pending, - OrderedFinalAggregateState::ReadingInput { table }, + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, )), Poll::Ready(Some(Ok(batch))) => { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); @@ -157,21 +409,86 @@ impl OrderedFinalAggregateStream { if let Err(e) = result { return ControlFlow::Break(( Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::ReadingInput { table }, + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, )); } + // Check memory reservation, and potentially spill. let timer = elapsed_compute.timer(); - let result = table.next_output_batch(); + let resize_result = + 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, + )); + }; + if table.is_empty() { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + return ControlFlow::Continue( + OrderedFinalAggregateState::Spilling { + table, + spill_context, + }, + ); + } + Err(e) => { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + } + + 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. + Ok(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. Ok(Some(batch)) => { - let next_state = - OrderedFinalAggregateState::ReadingInput { table }; - self.resize_reservation_for_state(&next_state); + 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, + )); + } + let next_state = OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }; ControlFlow::Break(( Poll::Ready(Some(Ok( @@ -180,36 +497,183 @@ impl OrderedFinalAggregateStream { next_state, )) } + // Can't do early emit, continue aggregating. Ok(None) => { - // Ordered variant doesn't support memory-limited - // execution, so it errors when memory reservation fails. - if let Err(e) = self.reservation.try_resize(table.memory_size()) { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::ReadingInput { table }, - )); - } - - // Can't do early emit, continue aggregating. ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput { table, + spill_context, }) } Err(e) => ControlFlow::Break(( Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::ReadingInput { table }, + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, )), } } Poll::Ready(Some(Err(e))) => ControlFlow::Break(( Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::ReadingInput { table }, + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, )), Poll::Ready(None) => { self.close_input(); - table.input_done(); - ControlFlow::Continue(OrderedFinalAggregateState::DrainingFinal { table }) + match spill_context { + Some(spill_context) if spill_context.has_spills() => { + ControlFlow::Continue( + OrderedFinalAggregateState::PreparingMergeInput { + table, + spill_context, + }, + ) + } + _ => { + table.input_done(); + ControlFlow::Continue( + OrderedFinalAggregateState::ProducingOutput { table }, + ) + } + } + } + } + } + + /// Sorts and spills one complete in-memory state run, then resumes input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_spilling( + &mut self, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::Spilling { + mut table, + mut spill_context, + } = original_state + else { + unreachable!("expected spilling state") + }; + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let mut result = spill_context.spill_table(&mut table); + + // Spilling shrinks the aggregate table and releases its accumulated + // memory. Update the reservation accordingly. + if self.reservation.try_resize(table.memory_size()).is_err() { + // Fold spill error and reservation error into one + result = internal_err!("Reservation after spilling should succeed") + } + + timer.done(); + + match result { + // Finished spilling the aggregate table, continue aggregating from input + Ok(true) => ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput { + table, + spill_context: Some(spill_context), + }), + Ok(false) => ControlFlow::Break(( + Poll::Ready(Some(internal_err!( + "Ordered final aggregation entered Spilling with an empty table" + ))), + OrderedFinalAggregateState::Done, + )), + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )), + } + } + + /// 1. Spills the last in-memory run. + /// 2. Constructs a globally ordered input stream by applying a sort-preserving + /// merge to all spills. + /// 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. + /// + /// Returns the next operator state with control flow decision. + fn handle_preparing_merge_input( + &mut self, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::PreparingMergeInput { + mut table, + mut spill_context, + } = original_state + else { + unreachable!("expected preparing merge input 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), + }; + timer.done(); + + match replay { + Ok(stream) => { + ControlFlow::Continue(OrderedFinalAggregateState::MergingSpills { + stream, + }) } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )), + } + } + + /// 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. + fn handle_merging_spills( + &mut self, + cx: &mut Context<'_>, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::MergingSpills { mut stream } = original_state + else { + unreachable!("expected merging spills 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), } } @@ -221,12 +685,12 @@ impl OrderedFinalAggregateStream { /// See comments at `poll_next()` for details. /// /// Returns the next operator state with control flow decision. - fn handle_draining_final( + fn handle_producing_output( &mut self, original_state: OrderedFinalAggregateState, ) -> OrderedFinalAggregateStateTransition { - let OrderedFinalAggregateState::DrainingFinal { table } = original_state else { - unreachable!("expected draining final state") + let OrderedFinalAggregateState::ProducingOutput { table } = original_state else { + unreachable!("expected producing output state") }; let mut table = table; @@ -238,11 +702,23 @@ impl OrderedFinalAggregateStream { match result { Ok(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, + )); + } OrderedFinalAggregateState::Done } else { - OrderedFinalAggregateState::DrainingFinal { table } + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ProducingOutput { table }, + )); + } + OrderedFinalAggregateState::ProducingOutput { table } }; - self.resize_reservation_for_state(&next_state); ControlFlow::Break(( Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), @@ -251,24 +727,18 @@ impl OrderedFinalAggregateStream { } Err(e) => ControlFlow::Break(( Poll::Ready(Some(Err(e))), - OrderedFinalAggregateState::DrainingFinal { table }, + OrderedFinalAggregateState::ProducingOutput { table }, )), Ok(None) => { + drop(table); let next_state = OrderedFinalAggregateState::Done; - self.resize_reservation_for_state(&next_state); + if let Err(e) = self.reservation.try_resize(0) { + return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)); + } ControlFlow::Continue(next_state) } } } - - fn resize_reservation_for_state(&mut self, state: &OrderedFinalAggregateState) { - let new_size = match state { - OrderedFinalAggregateState::ReadingInput { table } - | OrderedFinalAggregateState::DrainingFinal { table } => table.memory_size(), - OrderedFinalAggregateState::Done => 0, - }; - let _ = self.reservation.try_resize(new_size); - } } impl Stream for OrderedFinalAggregateStream { @@ -288,15 +758,35 @@ impl Stream for OrderedFinalAggregateStream { /// /// ReadingInput /// -> ReadingInput - /// Merge one input batch. If the ordering proves some groups are - /// complete, yield one final aggregate batch immediately, then continue - /// reading input. Otherwise continue directly with the next input batch. - /// -> DrainingFinal - /// Input was exhausted. Mark the table input as done so every remaining - /// group is safe to emit. + /// 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 + /// -> MergingSpills + /// The final run was spilled and the ordered replay stream was built. /// - /// DrainingFinal - /// -> DrainingFinal + /// MergingSpills + /// -> 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 @@ -319,8 +809,17 @@ impl Stream for OrderedFinalAggregateStream { state @ OrderedFinalAggregateState::ReadingInput { .. } => { self.handle_reading_input(cx, state) } - state @ OrderedFinalAggregateState::DrainingFinal { .. } => { - self.handle_draining_final(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); @@ -334,6 +833,15 @@ impl Stream for OrderedFinalAggregateStream { 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; diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index b4b7fa073aee0..deb69133a115e 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -23,7 +23,7 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::stream::{Stream, StreamExt}; @@ -31,6 +31,7 @@ use futures::stream::{Stream, StreamExt}; use super::AggregateExec; use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; use crate::aggregates::AggregateMode; +use crate::aggregates::order::GroupOrdering; use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput, SpillMetrics}; use crate::stream::EmptyRecordBatchStream; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; @@ -69,14 +70,50 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metric /// `k = 100`, it is safe to emit all groups with keys less than 100 because the /// input is ordered. /// +/// # Memory Pressure and Spilling +/// +/// ## Fully ordered case +/// +/// If the input is ordered by every group key, for example: +/// +/// - Input order: `a, b` +/// - `GROUP BY`: `a, b` +/// +/// Completed groups can be emitted as soon as the next group is observed. Thus, +/// only the current group remains active after completed groups are emitted, and +/// memory usage does not grow with the total number of groups. +/// +/// If a memory reservation nevertheless fails, the stream returns the error +/// directly, indicating an unexpected behavior. +/// +/// ## Partially ordered case +/// +/// If the input is ordered by only a subset of the group keys, for example: +/// +/// - Input order: `a` +/// - `GROUP BY`: `a, b` +/// +/// If one `a` value contains many distinct `b` values, the table may accumulate +/// enough groups to exceed the memory limit. +/// +/// - `OrderedPartialAggregateStream`: On reservation failure, it emits all current +/// intermediate states downstream and resets the table. The final stage can +/// merge repeated `(a, b)` state rows, so no disk spill is required. +/// - `OrderedFinalAggregateStream`: It cannot emit incomplete final results. On +/// reservation failure, it sorts the current intermediate states by the complete +/// group key and spills them as one run. After the input ends, it spills any +/// remaining states, performs a sort-preserving merge of all runs, and feeds the +/// merged input into a fully ordered final aggregate stream. +/// /// ## Implementation Note /// /// This is intentionally kept simple and closely maps to /// `GroupedHashAggregateStream` to finish the refactor sooner. /// +/// More applicable optimizations are left to future work. +/// /// See issue for details: /// -/// More applicable optimizations are left to future work. pub(crate) struct OrderedPartialAggregateStream { schema: SchemaRef, input: SendableRecordBatchStream, @@ -131,6 +168,10 @@ impl OrderedPartialAggregateStream { )?; let reservation = MemoryConsumer::new(format!("OrderedPartialAggregateStream[{partition}]")) + .with_can_spill(matches!( + table.group_ordering(), + GroupOrdering::Partial(_) + )) .register(context.memory_pool()); Ok(Self { @@ -185,6 +226,28 @@ impl OrderedPartialAggregateStream { )); } + // Check memory reservation. See function comments for details. + match self.resize_or_take_state_batch(&mut table) { + Ok(Some(batch)) => { + self.reduction_factor.add_part(batch.num_rows()); + return ControlFlow::Break(( + Poll::Ready(Some(Ok( + batch.record_output(&self.baseline_metrics) + ))), + OrderedPartialAggregateState::ReadingInput { table }, + )); + } + Ok(None) => {} + Err(e) => { + self.close_input(); + self.reservation.free(); + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedPartialAggregateState::Done, + )); + } + } + let timer = elapsed_compute.timer(); let result = table.next_output_batch(); timer.done(); @@ -195,9 +258,16 @@ impl OrderedPartialAggregateStream { // current state) Ok(Some(batch)) => { self.reduction_factor.add_part(batch.num_rows()); + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + self.close_input(); + self.reservation.free(); + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedPartialAggregateState::Done, + )); + } let next_state = OrderedPartialAggregateState::ReadingInput { table }; - self.resize_reservation_for_state(&next_state); ControlFlow::Break(( Poll::Ready(Some(Ok( @@ -206,21 +276,10 @@ impl OrderedPartialAggregateStream { next_state, )) } - Ok(None) => { - // Ordered variant don't support memory-limited execution, - // it have to error when OOM - if let Err(e) = self.reservation.try_resize(table.memory_size()) { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - OrderedPartialAggregateState::ReadingInput { table }, - )); - } - - // Can't do early emit, continue aggregating. - ControlFlow::Continue( - OrderedPartialAggregateState::ReadingInput { table }, - ) - } + // Can't do early emit, continue aggregating. + Ok(None) => ControlFlow::Continue( + OrderedPartialAggregateState::ReadingInput { table }, + ), Err(e) => ControlFlow::Break(( Poll::Ready(Some(Err(e))), OrderedPartialAggregateState::ReadingInput { table }, @@ -242,6 +301,42 @@ impl OrderedPartialAggregateStream { } } + /// Update the memory reservation, and: + /// - If memory reservation succeed, returns `Ok(None)` + /// - If memory reservation failed, + /// - If input is partially ordered, materialize all the output, and + /// directly send them to the final aggregation stage. + /// Returns `Ok(Some(batch))` + /// - If input is fully ordered, directly return error. It's not + /// expected to use more than constant memory. + /// Returns `Err(..)` + /// + /// # Implementation Note + /// Incrementally output it after the blocked state management is ready, keep + /// it simple for now. + /// + /// Issue: + fn resize_or_take_state_batch( + &mut self, + table: &mut OrderedAggregateTable, + ) -> Result> { + let oom = match self.reservation.try_resize(table.memory_size()) { + Ok(()) => return Ok(None), + Err(e @ DataFusionError::ResourcesExhausted(_)) => e, + Err(e) => return Err(e), + }; + + if matches!(table.group_ordering(), GroupOrdering::Full(_)) { + return Err(oom); + } + + let Some(batch) = table.take_state_batch()? else { + return Err(oom); + }; + self.reservation.try_resize(table.memory_size())?; + Ok(Some(batch)) + } + /// Emits one batch after input is exhausted. /// /// `table.input_done()` has already made every remaining group safe to emit, @@ -272,7 +367,9 @@ impl OrderedPartialAggregateStream { } else { OrderedPartialAggregateState::DrainingFinal { table } }; - self.resize_reservation_for_state(&next_state); + if let Err(e) = self.resize_reservation_for_state(&next_state) { + return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)); + } ControlFlow::Break(( Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), @@ -285,13 +382,18 @@ impl OrderedPartialAggregateStream { )), Ok(None) => { let next_state = OrderedPartialAggregateState::Done; - self.resize_reservation_for_state(&next_state); + if let Err(e) = self.resize_reservation_for_state(&next_state) { + return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)); + } ControlFlow::Continue(next_state) } } } - fn resize_reservation_for_state(&mut self, state: &OrderedPartialAggregateState) { + fn resize_reservation_for_state( + &mut self, + state: &OrderedPartialAggregateState, + ) -> Result<()> { let new_size = match state { OrderedPartialAggregateState::ReadingInput { table } | OrderedPartialAggregateState::DrainingFinal { table } => { @@ -299,7 +401,7 @@ impl OrderedPartialAggregateStream { } OrderedPartialAggregateState::Done => 0, }; - let _ = self.reservation.try_resize(new_size); + self.reservation.try_resize(new_size) } } From 02dfbbc667b74bb8bea7b380586b5a77d1542b59 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Fri, 17 Jul 2026 15:35:06 +0800 Subject: [PATCH 2/3] add basic tests --- .../test_files/ordered_aggregate_spill.slt | 201 ++++++++++++++++++ 1 file changed, 201 insertions(+) create mode 100644 datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt diff --git a/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt new file mode 100644 index 0000000000000..73f869467290c --- /dev/null +++ b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt @@ -0,0 +1,201 @@ +# 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. + +# End-to-end tests for ordered aggregation under finite memory. + +# Result set more than 100 lines will be hashed +hash-threshold 100 + +statement ok +SET datafusion.execution.target_partitions = 2 + +statement ok +SET datafusion.execution.batch_size = 128 + +statement ok +SET datafusion.optimizer.repartition_aggregations = true + +statement ok +SET datafusion.optimizer.prefer_existing_sort = true + +statement ok +SET datafusion.execution.enable_migration_aggregate = true + +statement ok +SET datafusion.runtime.memory_limit = '1M' + +# ================================================================================== +# Input is fully ordered by group keys (input order by (a,b), query is 'group by a,b') +# ================================================================================== + +# Fully ordered input uses the ordered partial and final streams without spill. +query TT +EXPLAIN ANALYZE +SELECT v1, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY v1 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,ordering_mode=Sorted, metrics=[spill_count=0,] +02)--RepartitionExec:preserve_order=true +03)----AggregateExec: mode=Partial,ordering_mode=Sorted, metrics=[spill_count=0,] + + +query II rowsort +SELECT v1, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY v1 +---- +40002 values hashing to 34c2b23730596cbd2489ed4a627c17d7 + +# The same fully ordered query cannot spill and reports OOM under tighter memory. +statement ok +SET datafusion.runtime.memory_limit = '64K' + +query error Resources exhausted +SELECT v1, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY v1 + +# ================================================================================== +# Input is partially ordered by group keys (input order by (a), query is 'group by a,b') +# +# Try different memory limits, ensure result is the same, but spill count differ + +# HACK: check `spilled_bytes=x KB` to ensure it has spilled. If it has not spilled, +# the it shows `spilled_bytes = 0B`. Should better check spill_count, but it's not +# stable due to ordered hash repartition, and `sqllogictest` don't support regex. +# ================================================================================== + +statement ok +SET datafusion.runtime.memory_limit = '2M' + +statement ok +SET datafusion.optimizer.enable_round_robin_repartition = false + +# Round 1: The input is partially ordered and does not spill with a 2 MB limit. +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,aggr=[sum(t1.v1 * Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] +02)--RepartitionExec:input_partitions=1, maintains_sort_order=true +03)----AggregateExec: mode=Partial,ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +# All rounds should have the same result hash +query III rowsort +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +45000 values hashing to e6ece4b4b86e6152a1c785e90ab5ba12 + +# Round 2: The same query spills five times with a 600 KB limit. +statement ok +SET datafusion.runtime.memory_limit = '600K' + +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,aggr=[sum(t1.v1 * Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spilled_bytes= KB,] +02)--RepartitionExec:input_partitions=1, maintains_sort_order=true +03)----AggregateExec: mode=Partial,ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +# All rounds should have the same result hash +query III rowsort +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +45000 values hashing to e6ece4b4b86e6152a1c785e90ab5ba12 + +# Round 3: The same query spills six times with a 500 KB limit. +statement ok +SET datafusion.runtime.memory_limit = '500K' + +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,aggr=[sum(t1.v1 * Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spilled_bytes= KB,] +02)--RepartitionExec:input_partitions=1, maintains_sort_order=true +03)----AggregateExec: mode=Partial,ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +# All rounds should have the same result hash +query III rowsort +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +45000 values hashing to e6ece4b4b86e6152a1c785e90ab5ba12 + +# Exercise the same spill path with a variable-width string aggregate state in +# the spilled payload. Keep one partial input partition so memory pressure is on +# the ordered aggregate rather than a repartition merge. +statement ok +SET datafusion.runtime.memory_limit = '600K' + +# Ensures final aggregate has spill_count > 0 +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, + sum(v1 * 2), min(CAST(v1 % 2 AS VARCHAR)) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,aggr=[sum(t1.v1 * Int64(2)), min(t1.v1 % Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spilled_bytes=KB,] +02)--RepartitionExec:input_partitions=1, maintains_sort_order=true +03)----AggregateExec: mode=Partial,aggr=[sum(t1.v1 * Int64(2)), min(t1.v1 % Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +statement ok +RESET datafusion.runtime.memory_limit + +statement ok +RESET datafusion.optimizer.enable_round_robin_repartition + +statement ok +RESET datafusion.execution.enable_migration_aggregate + +statement ok +RESET datafusion.optimizer.prefer_existing_sort + +statement ok +RESET datafusion.optimizer.repartition_aggregations + +statement ok +RESET datafusion.execution.batch_size + +statement ok +SET datafusion.execution.target_partitions = 4 + +statement ok +RESET datafusion.catalog.create_default_catalog_and_schema From a780de8773903645de7432bb6372c42ff8e549f4 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Fri, 17 Jul 2026 16:53:51 +0800 Subject: [PATCH 3/3] fix test --- datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt index 73f869467290c..a4d492ab82e12 100644 --- a/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt +++ b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt @@ -64,7 +64,7 @@ GROUP BY v1 # The same fully ordered query cannot spill and reports OOM under tighter memory. statement ok -SET datafusion.runtime.memory_limit = '64K' +SET datafusion.runtime.memory_limit = '1K' query error Resources exhausted SELECT v1, sum(v1 * 2)