This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25538-d67d8111c871d29658e2fd307b96343297619abd in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 4fab63d0dfc48b8dd641f534f0da367144a184f2 Author: Jay Zhan <[email protected]> AuthorDate: Mon Sep 21 13:25:17 2026 +0000 refactor(hash-aggr): share one spill context between the spilling aggregate streams (#25538) ## Which issue does this PR close? - Part of #25537 (unified spill-replay driver for grouped aggregation). Follow-up to #22710. ## Rationale for this change The four grouped aggregation streams that can spill — `FinalHashAggregateStream`, `SingleHashAggregateStream`, `OrderedSingleAggregateStream` and `OrderedFinalAggregateStream` — each define their own spill context (`FinalSpillContext`, `SingleSpillContext`, `OrderedSingleSpillContext`, `OrderedFinalSpillContext`). The four types have the same seven fields and the same three operations: - build the spill sort key and the aggregate configuration used for replay, - sort the table's state batch and write it as one spill file, - merge all runs with `StreamingMergeBuilder` and replay them through `OrderedFinalAggregateStream::new_with_input_and_metrics`. The bodies are copies of each other; the only real differences are constructor-time: | | sort key | replay aggregate | |---|---|---| | Final hash | group columns in natural order | same aggregate | | Single hash | group columns in natural order | `Single → Final`, `group_by.as_final()` | | Ordered final | already-ordered columns first, then the rest | same aggregate | | Ordered single | already-ordered columns first, then the rest | `Single → Final`, `group_by.as_final()` | Both axes are functions of values the constructor already receives (`AggregateMode` and `InputOrderMode`), and the natural order is just the ordered-columns-first formula with no ordered columns. ## What changes are included in this PR? No behaviour change. - Add `aggregates/spill.rs` with a single non-generic `AggregateSpill` (`try_new`, `has_spills`, `spill`, `into_replay_stream`). It is non-generic because the only thing a spill needs from a table is the batch returned by `take_state_batch()`, so the caller passes that batch in. - Delete the four per-stream spill context types and use `AggregateSpill` instead. - Spill request descriptions (`"FinalHashAggregateSpill"` etc.) and memory consumer names are unchanged. The text of four `internal_err!` messages that cannot be reached by users is now shared. Net: 4 stream files −665 lines, +1 new file of ~235 lines. This is the first step towards a single spill-replay driver for these streams (see the linked issue), but it stands on its own. ## What is the testing strategy for this PR? Existing tests: the aggregate unit tests in `aggregates/mod.rs` and `ordered_final_stream.rs` (spill + replay, OOM, drop/cancel, memory accounting), `aggregate_memory_spill.slt`, `ordered_aggregate_spill.slt`, the `memory_limit` integration tests and the aggregate fuzz tests. No new tests since there is no new behaviour. ## Are there any user-facing changes? No. --------- Co-authored-by: Yongting You <[email protected]> --- .../physical-plan/src/aggregates/hash_stream.rs | 190 ++----------- datafusion/physical-plan/src/aggregates/mod.rs | 1 + .../src/aggregates/ordered_final_stream.rs | 189 ++----------- .../src/aggregates/ordered_single_stream.rs | 209 ++------------- .../physical-plan/src/aggregates/single_stream.rs | 210 ++------------- datafusion/physical-plan/src/aggregates/spill.rs | 293 +++++++++++++++++++++ 6 files changed, 361 insertions(+), 731 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index d340bf5d43..43257acc14 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}; @@ -143,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, @@ -179,28 +172,6 @@ pub(crate) struct PartialHashAggregateStream { hash_table: Option<AggregateHashTable<PartialMarker>>, } -/// 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<TaskContext>, - /// 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<SortedSpillFile>, -} - /// 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<AggregateHashTable<FinalMarker>>, /// `None` if spilling is not supported by the configured `DiskManager`. - spill_context: Option<Box<FinalSpillContext>>, -} - -impl FinalSpillContext { - fn new( - agg: &AggregateExec, - context: &Arc<TaskContext>, - partition: usize, - batch_size: usize, - spill_schema: &SchemaRef, - spill_metrics: SpillMetrics, - ) -> Result<Self> { - 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<FinalMarker>, - ) -> 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<SendableRecordBatchStream> { - 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<Box<AggregateSpill>>, } #[derive(PartialEq)] @@ -762,11 +598,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, )?)) @@ -840,7 +678,7 @@ impl FinalHashAggregateStream { /// Reserve memory for the current aggregate table. fn reservation_size_for_table( hash_table: &AggregateHashTable<FinalMarker>, - spill_context: Option<&FinalSpillContext>, + spill_context: Option<&AggregateSpill>, ) -> usize { let table_size = hash_table.memory_size(); if spill_context.is_some() { @@ -864,7 +702,7 @@ impl FinalHashAggregateStream { async fn consume_input( &mut self, hash_table: &mut AggregateHashTable<FinalMarker>, - spill_context: &mut Option<Box<FinalSpillContext>>, + spill_context: &mut Option<Box<AggregateSpill>>, ) -> Result<()> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); @@ -912,7 +750,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.sort_and_spill(batch)); // Spilling shrinks the aggregate table and releases its accumulated // memory. Update the reservation accordingly. @@ -942,14 +782,16 @@ impl FinalHashAggregateStream { async fn produce_output_from_spills( &mut self, mut hash_table: AggregateHashTable<FinalMarker>, - mut spill_context: Box<FinalSpillContext>, + mut spill_context: Box<AggregateSpill>, mut emitter: TryEmitter<RecordBatch, DataFusionError>, ) -> 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.sort_and_spill(batch))?; // Construct the ordered input used to merge all spill files. let mut output_stream = @@ -978,7 +820,7 @@ impl FinalHashAggregateStream { fn switch_to_ordered_final_stream( &mut self, hash_table: AggregateHashTable<FinalMarker>, - spill_context: Box<FinalSpillContext>, + spill_context: Box<AggregateSpill>, ) -> Result<SendableRecordBatchStream> { 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 806eacb35d..5fa350b55d 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 cc992eaa51..3a09dbd97d 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}; @@ -57,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 @@ -70,30 +65,6 @@ pub(crate) struct OrderedFinalAggregateStream { state: Option<OrderedFinalAggregateState>, } -/// 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<TaskContext>, - /// 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<SortedSpillFile>, -} - /// 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<Box<OrderedFinalSpillContext>>, + spill_context: Option<Box<AggregateSpill>>, }, Spilling { table: OrderedAggregateTable<FinalMarker>, - spill_context: Box<OrderedFinalSpillContext>, + spill_context: Box<AggregateSpill>, }, ProducingOutput { table: OrderedAggregateTable<FinalMarker>, }, PreparingMergeInput { table: OrderedAggregateTable<FinalMarker>, - spill_context: Box<OrderedFinalSpillContext>, + spill_context: Box<AggregateSpill>, }, MergingSpills { stream: SendableRecordBatchStream, @@ -126,140 +97,6 @@ type OrderedFinalAggregateStateTransition = ControlFlow< OrderedFinalAggregateState, >; -impl OrderedFinalSpillContext { - fn new( - agg: &AggregateExec, - context: &Arc<TaskContext>, - partition: usize, - batch_size: usize, - input_order_mode: &InputOrderMode, - spill_schema: &SchemaRef, - spill_metrics: SpillMetrics, - ) -> Result<Self> { - 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<FinalMarker>, - ) -> 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<SendableRecordBatchStream> { - 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<FinalMarker>, - 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.sort_and_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.sort_and_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 40ba90729b..02c3b2a368 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<OrderedSingleAggregateState>, } -/// 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<TaskContext>, - /// 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<SortedSpillFile>, -} - /// 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<Box<OrderedSingleSpillContext>>, + spill_context: Option<Box<AggregateSpill>>, }, Spilling { table: OrderedAggregateTable<SingleMarker>, - spill_context: Box<OrderedSingleSpillContext>, + spill_context: Box<AggregateSpill>, }, ProducingOutput { table: OrderedAggregateTable<SingleMarker>, }, PreparingMergeInput { table: OrderedAggregateTable<SingleMarker>, - spill_context: Box<OrderedSingleSpillContext>, + spill_context: Box<AggregateSpill>, }, MergingSpills { stream: SendableRecordBatchStream, @@ -169,157 +137,6 @@ type OrderedSingleAggregateStateTransition = ControlFlow< OrderedSingleAggregateState, >; -impl OrderedSingleSpillContext { - fn new( - agg: &AggregateExec, - context: &Arc<TaskContext>, - partition: usize, - batch_size: usize, - input_order_mode: &InputOrderMode, - spill_schema: &SchemaRef, - spill_metrics: SpillMetrics, - ) -> Result<Self> { - 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<SingleMarker>, - ) -> 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<SendableRecordBatchStream> { - 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<SingleMarker>, - 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.sort_and_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.sort_and_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 091e9fb940..a185650e2f 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}; @@ -87,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 /// @@ -138,51 +132,22 @@ pub(crate) struct SingleHashAggregateStream { group_values_soft_limit: Option<usize>, } -/// 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<TaskContext>, - /// 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<SortedSpillFile>, -} - /// See comments at `poll_next()` for details. enum SingleHashAggregateState { ReadingInput { hash_table: AggregateHashTable<SingleMarker>, - spill_context: Option<Box<SingleSpillContext>>, + spill_context: Option<Box<AggregateSpill>>, }, Spilling { hash_table: AggregateHashTable<SingleMarker>, - spill_context: Box<SingleSpillContext>, + spill_context: Box<AggregateSpill>, }, ProducingOutput { hash_table: AggregateHashTable<SingleMarker>, }, PreparingMergeInput { hash_table: AggregateHashTable<SingleMarker>, - spill_context: Box<SingleSpillContext>, + spill_context: Box<AggregateSpill>, }, MergingSpills { stream: SendableRecordBatchStream, @@ -200,152 +165,6 @@ type SingleHashAggregateStateTransition = ControlFlow< SingleHashAggregateState, >; -impl SingleSpillContext { - fn new( - agg: &AggregateExec, - context: &Arc<TaskContext>, - partition: usize, - batch_size: usize, - spill_schema: &SchemaRef, - spill_metrics: SpillMetrics, - ) -> Result<Self> { - 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<SingleMarker>, - ) -> 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<SendableRecordBatchStream> { - 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<SingleMarker>, - 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<SingleMarker>, - spill_context: Option<Box<SingleSpillContext>>, + spill_context: Option<Box<AggregateSpill>>, ) -> 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.sort_and_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.sort_and_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 0000000000..5ad9a200e4 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -0,0 +1,293 @@ +// 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. + /// + /// # Example walkthrough + /// + /// This example walks through two key APIs of [`AggregateSpill`]: + /// - [`AggregateSpill::sort_and_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::sort_and_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<TaskContext>, + /// 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<SortedSpillFile>, + /// 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<TaskContext>, + partition: usize, + batch_size: usize, + input_order_mode: &InputOrderMode, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result<Self> { + 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 sort_and_spill( + &mut self, + state_batch: Option<RecordBatch>, + ) -> 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<SendableRecordBatchStream> { + 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)) + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
