diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index d6e38b5d01995..84d7e9c4508b5 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -101,7 +101,8 @@ async fn group_by_row_hash() { TestCase::new() .with_query("select count(*) from t GROUP BY response_bytes") .with_expected_errors(vec![ - "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) as:\n GroupedHashAggregateStream" + "Resources exhausted: Additional allocation failed", + "for FinalHashAggregateStream[0]", ]) .with_memory_limit(2_000) .run() @@ -114,7 +115,8 @@ async fn group_by_hash() { // group by dict column .with_query("select count(*) from t GROUP BY service, host, pod, container") .with_expected_errors(vec![ - "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) as:\n GroupedHashAggregateStream" + "Resources exhausted: Additional allocation failed", + "for PartialHashAggregateStream[0]", ]) .with_memory_limit(1_000) .run() @@ -425,7 +427,7 @@ async fn oom_grouped_hash_aggregate() { .with_query("SELECT COUNT(*), SUM(request_bytes) FROM t GROUP BY host") .with_expected_errors(vec![ "Failed to allocate additional", - "GroupedHashAggregateStream[0] (count(1), sum(t.request_bytes))", + "for PartialHashAggregateStream[0]", ]) .with_memory_limit(1_000) .run() diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index e7f0f075b33a5..2f855b4e3ec90 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -25,25 +25,34 @@ //! //! See issue for details: +use std::mem::size_of; use std::ops::ControlFlow; use std::sync::Arc; use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err, 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::{ AggregateHashTable, FinalMarker, PartialMarker, PartialSkipMarker, }; +use super::group_values::GroupByMetrics; +use super::ordered_final_stream::OrderedFinalAggregateStream; use super::skip_partial::SkipAggregationProbe; 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; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; @@ -107,6 +116,28 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metric /// accumulated groups, then switches to a skip state. In that state, each /// remaining input batch is converted directly to partial aggregate state rows /// without inserting the rows into the grouped hash table. +/// +/// # Feature: Memory-limited Execution +/// +/// ## Partial Aggregation +/// +/// Partial aggregation can emit incomplete results because the final stage merges +/// all intermediate states for the same group. If the memory reservation exceeds +/// its limit after aggregating an input batch, this stream emits all accumulated +/// states and continues aggregating the remaining input with an empty table. +/// +/// ## Final Aggregation +/// +/// During final aggregation, group keys and states accumulate. If memory usage +/// exceeds the budget, spilling is triggered as follows: +/// 1. After aggregating a new input batch, if the memory reservation exceeds its +/// limit, spill all accumulated groups and states. +/// - Sort all groups by the group keys before spilling. +/// 2. Repeat until the input is exhausted. +/// 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. pub(crate) struct PartialHashAggregateStream { /// Output schema: group columns followed by partial aggregate state columns. schema: SchemaRef, @@ -114,6 +145,9 @@ pub(crate) struct PartialHashAggregateStream { /// Input batches containing raw rows, not partial aggregate state. input: SendableRecordBatchStream, + /// Target output batch size from configuration. + batch_size: usize, + /// Memory reservation for group keys and accumulators. reservation: MemoryReservation, @@ -142,6 +176,13 @@ enum PartialHashAggregateState { ReadingInput { hash_table: AggregateHashTable, }, + /// A fully materialized partial-state batch being emitted incrementally. + EmittingOnMemoryPressure { + hash_table: AggregateHashTable, + // After each incremental emitting step, the `remaining_groups` will be updated + // with batch slicing. + remaining_groups: RecordBatch, + }, ProducingOutput { hash_table: AggregateHashTable, /// If `None`, partial skip was never triggered and this state will @@ -154,6 +195,10 @@ enum PartialHashAggregateState { hash_table: AggregateHashTable, }, Done, + /// Sentinel state to use when returning error from any other states, because: + /// - It explicitly releases state-owned resources immediately + /// - More defensive against accidentally resuming execution after error + Error, } type PartialHashAggregatePoll = Poll>>; @@ -162,26 +207,26 @@ type PartialHashAggregateStateTransition = ControlFlow< PartialHashAggregateState, >; -impl PartialHashAggregateState { - fn hash_table(&self) -> &AggregateHashTable { - match self { - Self::ReadingInput { hash_table } - | Self::ProducingOutput { hash_table, .. } => hash_table, - Self::SkippingAggregation { .. } | Self::Done => { - unreachable!("state does not hold a partial hash table") - } - } - } - - fn hash_table_mut(&mut self) -> &mut AggregateHashTable { - match self { - Self::ReadingInput { hash_table } - | Self::ProducingOutput { hash_table, .. } => hash_table, - Self::SkippingAggregation { .. } | Self::Done => { - unreachable!("state does not hold a partial hash table") - } - } - } +/// Spill configuration and accumulated runs for final hash aggregation. +/// +/// Each spill event drains all currently buffered groups, sorts their intermediate +/// states by the full group key, and writes them to one spill file. All files are +/// merged and replayed after the original input ends. +struct FinalSpillContext { + /// Aggregate configuration used to construct the final replay stream. + final_agg: AggregateExec, + /// Task context. + context: Arc, + /// Original partition index. + partition: usize, + /// Target batch size from configuration. + batch_size: usize, + /// Full group-key ordering kept by every spill file and the merged input. + spill_expr: LexOrdering, + /// Spill I/O and metrics manager. + spill_manager: SpillManager, + /// Spill runs waiting to be merged, they're all sorted by full group-by keys. + spills: Vec, } /// Hash aggregation is implemented in two stages: partial and final. This @@ -198,7 +243,7 @@ pub(crate) struct FinalHashAggregateStream { /// Execution metrics shared with the aggregate plan node. baseline_metrics: BaselineMetrics, - /// Memory reservation for group keys and accumulators. + /// Memory reservation for group keys, accumulators, and spill sorting. reservation: MemoryReservation, /// See comments for the same variable in [`PartialHashAggregateStream`]. @@ -215,11 +260,28 @@ pub(crate) struct FinalHashAggregateStream { enum FinalHashAggregateState { ReadingInput { hash_table: AggregateHashTable, + /// `None` if spilling is not supported by the configured `DiskManager`. + spill_context: Option>, + }, + Spilling { + hash_table: AggregateHashTable, + spill_context: Box, }, ProducingOutput { hash_table: AggregateHashTable, }, + PreparingMergeInput { + hash_table: AggregateHashTable, + spill_context: Box, + }, + MergingSpills { + stream: SendableRecordBatchStream, + }, Done, + /// Sentinel state to use when returning error from any other states, because: + /// - It explicitly releases state-owned resources immediately + /// - More defensive against accidentally resuming execution after error + Error, } type FinalHashAggregatePoll = Poll>>; @@ -228,42 +290,137 @@ type FinalHashAggregateStateTransition = ControlFlow< FinalHashAggregateState, >; -impl FinalHashAggregateState { - fn hash_table(&self) -> &AggregateHashTable { - match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table - } - Self::Done => unreachable!("Done state does not hold a hash table"), - } - } +impl FinalSpillContext { + fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + batch_size: usize, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result { + let group_schema = agg.group_by.group_schema(&agg.input().schema())?; + let output_ordering = agg.cache.output_ordering(); + let spill_sort_exprs = + group_schema + .fields() + .iter() + .enumerate() + .map(|(idx, field)| { + let output_expr = Column::new(field.name(), idx); + let sort_options = output_ordering + .and_then(|ordering| ordering.get_sort_options(&output_expr)) + .unwrap_or_default(); + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("Final hash aggregate spill expression is empty"); + }; - fn hash_table_mut(&mut self) -> &mut AggregateHashTable { - match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table - } - Self::Done => unreachable!("Done state does not hold a hash table"), - } + 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 into_hash_table(self) -> AggregateHashTable { - match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table - } - Self::Done => unreachable!("Done state does not hold a hash table"), - } + fn has_spills(&self) -> bool { + !self.spills.is_empty() } - fn into_producing_output(self) -> Self { - Self::ProducingOutput { - hash_table: self.into_hash_table(), - } + /// Sorts and spills the aggregated groups. Memory reservation should be updated + /// by the caller. + /// + /// Individual spill files are ordered by the `group by` keys. + /// + /// See [`FinalHashAggregateStream`] for spilling details. + fn spill_table( + &mut self, + hash_table: &mut AggregateHashTable, + ) -> Result<()> { + let Some(batch) = hash_table.take_state_batch()? else { + return Ok(()); + }; + + let sorted_iter = + IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "FinalHashAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Final hash aggregation produced an empty spill"); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) } - fn into_done(self) -> Self { - Self::Done + /// Merges every sorted run, and do the aggregate evaluation with + /// [`OrderedFinalAggregateStream`] + fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + group_by_metrics: GroupByMetrics, + reservation: MemoryReservation, + ) -> Result { + let Self { + final_agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + // The merge and replay table are two components of the same aggregate + // operator. Keep them under one consumer registration so a fair memory + // pool does not divide this operator's quota between its own phases. + let merge_reservation = reservation.new_empty(); + let merged = StreamingMergeBuilder::new() + .with_schema(spill_schema) + .with_spill_manager(spill_manager) + .with_sorted_spill_files(spills) + .with_expressions(&spill_expr) + .with_metrics(baseline_metrics.intermediate()) + .with_batch_size(batch_size) + .with_reservation(merge_reservation) + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &final_agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + group_by_metrics, + None, + reservation, + )?; + Ok(Box::pin(replay)) } } @@ -317,11 +474,13 @@ impl PartialHashAggregateStream { let reservation = MemoryConsumer::new(format!("PartialHashAggregateStream[{partition}]")) + .with_can_spill(true) .register(context.memory_pool()); Ok(Self { schema, input, + batch_size, baseline_metrics, reservation, reduction_factor, @@ -331,6 +490,24 @@ impl PartialHashAggregateStream { }) } + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_err(error: DataFusionError) -> PartialHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + PartialHashAggregateState::Error, + )) + } + + fn break_with_internal_err( + message: impl std::fmt::Display, + ) -> PartialHashAggregateStateTransition { + Self::break_with_err(internal_datafusion_err!("{message}")) + } + /// See comments in [`Self::group_values_soft_limit`] for details. fn hit_soft_group_limit( &self, @@ -375,48 +552,48 @@ impl PartialHashAggregateStream { fn handle_reading_input( &mut self, cx: &mut Context<'_>, - mut original_state: PartialHashAggregateState, + original_state: PartialHashAggregateState, ) -> PartialHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - PartialHashAggregateState::ReadingInput { .. } - )); - debug_assert!(original_state.hash_table().is_building()); + let PartialHashAggregateState::ReadingInput { mut hash_table } = original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected ReadingInput state", + ); + }; + debug_assert!(hash_table.is_building()); match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)), + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + PartialHashAggregateState::ReadingInput { hash_table }, + )), Poll::Ready(Some(Ok(batch))) => { + // ---------------------------------- + // Step 1: Aggregate the input batch + // ---------------------------------- let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); let input_rows = batch.num_rows(); self.reduction_factor.add_total(input_rows); - let result = original_state.hash_table_mut().aggregate_batch(&batch); + let result = hash_table.aggregate_batch(&batch); timer.done(); if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + return Self::break_with_err(e); } - if self.hit_soft_group_limit(original_state.hash_table()) { + // -------------------------------- + // Step 2: Soft limit optimization + // -------------------------------- + if self.hit_soft_group_limit(&hash_table) { let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut(), true); + let result = self.start_output(&mut hash_table, true); timer.done(); if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + return Self::break_with_err(e); } - let PartialHashAggregateState::ReadingInput { hash_table } = - original_state - else { - unreachable!("expected reading input state") - }; return ControlFlow::Continue( PartialHashAggregateState::ProducingOutput { hash_table, @@ -425,17 +602,20 @@ impl PartialHashAggregateStream { ); } + // ---------------------------------------------- + // Step 3: Skip partial aggregation optimization + // ---------------------------------------------- self.update_skip_aggregation_probe( input_rows, - original_state.hash_table().building_group_count(), + hash_table.building_group_count(), ); // True branch: a decision has been made to skip partial aggregation. if self.should_skip_aggregation() { let timer = elapsed_compute.timer(); - let result = match original_state.hash_table().partial_skip_table() { + let result = match hash_table.partial_skip_table() { Ok(skip_hash_table) => self - .start_output(original_state.hash_table_mut(), false) + .start_output(&mut hash_table, false) .map(|()| skip_hash_table), Err(e) => Err(e), }; @@ -443,12 +623,6 @@ impl PartialHashAggregateStream { match result { Ok(skip_hash_table) => { - let PartialHashAggregateState::ReadingInput { hash_table } = - original_state - else { - unreachable!("expected reading input state") - }; - // Move to `ProducingOutput` first. Its `skip_hash_table` // field moves the stream to skip-partial aggregation after // the accumulated batches have been output. @@ -459,60 +633,126 @@ impl PartialHashAggregateStream { }, ); } - Err(e) => { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } + Err(e) => return Self::break_with_err(e), } } - // TODO: impl memory-limited aggr, when OOM directly send - // partial state to final aggregate stage - if let Err(e) = self - .reservation - .try_resize(original_state.hash_table().memory_size()) - { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + // ------------------------------------------------- + // Step 4: Larger-than-memory execution (early emit) + // ------------------------------------------------- + let timer = elapsed_compute.timer(); + let resize_result = self.reservation.try_resize(hash_table.memory_size()); + timer.done(); + match resize_result { + Ok(()) => {} + Err(DataFusionError::ResourcesExhausted(_)) => { + let elapsed_compute = + self.baseline_metrics.elapsed_compute().clone(); + // Stops on drop + let _timer = elapsed_compute.timer(); + let state_batch_result = hash_table.take_state_batch(); + + // Emitting clears the aggregate table and releases its + // accumulated memory. Update the reservation accordingly. + let resize_result = + self.reservation.try_resize(hash_table.memory_size()); + + if let Err(e) = resize_result { + return Self::break_with_err(e); + } + + let materialized_group_states = match state_batch_result { + Ok(Some(batch)) => batch, + Ok(None) => { + return Self::break_with_err(internal_datafusion_err!( + "Partial hash aggregate ran out of memory with no aggregated groups" + )); + } + Err(e) => return Self::break_with_err(e), + }; + + return ControlFlow::Continue( + PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: materialized_group_states, + }, + ); + } + Err(e) => return Self::break_with_err(e), } - ControlFlow::Continue(original_state) - } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + ControlFlow::Continue(PartialHashAggregateState::ReadingInput { + hash_table, + }) } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), Poll::Ready(None) => { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut(), true); + let result = self.start_output(&mut hash_table, true); timer.done(); match result { - Ok(()) => { - let PartialHashAggregateState::ReadingInput { hash_table } = - original_state - else { - unreachable!("expected reading input state") - }; - ControlFlow::Continue( - PartialHashAggregateState::ProducingOutput { - hash_table, - skip_hash_table: None, - }, - ) - } - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + Ok(()) => ControlFlow::Continue( + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table: None, + }, + ), + Err(e) => Self::break_with_err(e), } } } } + /// Handle EmittingOnMemoryPressure state - emit a materialized partial-state + /// batch in `batch_size`(from configuration) slices, then resume reading input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_emitting_on_memory_pressure( + &mut self, + original_state: PartialHashAggregateState, + ) -> PartialHashAggregateStateTransition { + let PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: batch, + } = original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected EmittingOnMemoryPressure state", + ); + }; + + let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { + // Last batch to output, go back to `ReadingInput` + ( + batch, + PartialHashAggregateState::ReadingInput { hash_table }, + ) + } else { + // More batch to output, continue in the current state. + let remaining = + batch.slice(self.batch_size, batch.num_rows() - self.batch_size); + let output = batch.slice(0, self.batch_size); + ( + output, + PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: remaining, + }, + ) + }; + + self.reduction_factor.add_part(output_batch.num_rows()); + debug_assert!(output_batch.num_rows() > 0); + ControlFlow::Break(( + Poll::Ready(Some(Ok(output_batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + /// Handle ProducingOutput state - emit partial aggregate state batches. /// /// See comments at `poll_next()` for details. @@ -520,42 +760,41 @@ impl PartialHashAggregateStream { /// Returns the next operator state with control flow decision. fn handle_producing_output( &mut self, - mut original_state: PartialHashAggregateState, + original_state: PartialHashAggregateState, ) -> PartialHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - PartialHashAggregateState::ProducingOutput { .. } - )); - debug_assert!(!original_state.hash_table().is_building()); + let PartialHashAggregateState::ProducingOutput { + mut hash_table, + skip_hash_table, + } = original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected ProducingOutput state", + ); + }; + debug_assert!(!hash_table.is_building()); let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = original_state.hash_table_mut().next_output_batch(); + let result = hash_table.next_output_batch(); timer.done(); match result { Ok(Some(batch)) => { - let _ = self - .reservation - .try_resize(original_state.hash_table().memory_size()); + let _ = self.reservation.try_resize(hash_table.memory_size()); self.reduction_factor.add_part(batch.num_rows()); debug_assert!(batch.num_rows() > 0); - let next_state = if original_state.hash_table().is_done() { - match original_state { - PartialHashAggregateState::ProducingOutput { - skip_hash_table: Some(hash_table), - .. - } => { + let next_state = if hash_table.is_done() { + match skip_hash_table { + Some(hash_table) => { PartialHashAggregateState::SkippingAggregation { hash_table } } - PartialHashAggregateState::ProducingOutput { - skip_hash_table: None, - .. - } => PartialHashAggregateState::Done, - _ => unreachable!("expected producing output state"), + None => PartialHashAggregateState::Done, } } else { - original_state + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table, + } }; ControlFlow::Break(( @@ -567,20 +806,15 @@ impl PartialHashAggregateStream { let _ = self.reservation.try_resize(0); // If the previous `Aggregating` stage decided to skip partial // aggregation, go to the `SkippingAggregation` stage; otherwise finish. - let next_state = match original_state { - PartialHashAggregateState::ProducingOutput { - skip_hash_table: Some(hash_table), - .. - } => PartialHashAggregateState::SkippingAggregation { hash_table }, - PartialHashAggregateState::ProducingOutput { - skip_hash_table: None, - .. - } => PartialHashAggregateState::Done, - _ => unreachable!("expected producing output state"), + let next_state = match skip_hash_table { + Some(hash_table) => { + PartialHashAggregateState::SkippingAggregation { hash_table } + } + None => PartialHashAggregateState::Done, }; ControlFlow::Continue(next_state) } - Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)), + Err(e) => Self::break_with_err(e), } } @@ -592,15 +826,21 @@ impl PartialHashAggregateStream { fn handle_skipping_aggregation( &mut self, cx: &mut Context<'_>, - mut original_state: PartialHashAggregateState, + original_state: PartialHashAggregateState, ) -> PartialHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - PartialHashAggregateState::SkippingAggregation { .. } - )); + let PartialHashAggregateState::SkippingAggregation { mut hash_table } = + original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected SkippingAggregation state", + ); + }; match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)), + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + PartialHashAggregateState::SkippingAggregation { hash_table }, + )), Poll::Ready(Some(Ok(batch))) => { if let Some(probe) = self.skip_aggregation_probe.as_mut() { probe.record_skipped(&batch); @@ -608,12 +848,7 @@ impl PartialHashAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = match &mut original_state { - PartialHashAggregateState::SkippingAggregation { hash_table } => { - hash_table.convert_batch_to_state(&batch) - } - _ => unreachable!("expected skipping aggregation state"), - }; + let result = hash_table.convert_batch_to_state(&batch); timer.done(); match result { @@ -621,16 +856,12 @@ impl PartialHashAggregateStream { Poll::Ready(Some( Ok(batch.record_output(&self.baseline_metrics)), )), - original_state, + PartialHashAggregateState::SkippingAggregation { hash_table }, )), - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + Err(e) => Self::break_with_err(e), } } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), Poll::Ready(None) => { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); @@ -659,6 +890,9 @@ impl Stream for PartialHashAggregateStream { /// -> ReadingInput /// Aggregate one batch, update the inner aggregate hash table, and /// continue with the next input batch. + /// -> EmittingOnMemoryPressure + /// The table cannot reserve enough memory. Materialize all accumulated + /// partial states and begin emitting them incrementally. /// -> ProducingOutput(skip=None) /// Input was exhausted, or the soft group limit was reached. Move to /// the next state to start outputting. @@ -668,6 +902,13 @@ impl Stream for PartialHashAggregateStream { /// the `SkippingAggregation` state to convert input directly to partial /// state without aggregation. /// + /// EmittingOnMemoryPressure + /// -> EmittingOnMemoryPressure + /// One batch-sized slice was yielded; repeat until all materialized + /// partial states are emitted. + /// -> ReadingInput + /// The materialized states were emitted; continue with the empty table. + /// /// ProducingOutput(skip=None) /// -> ProducingOutput(skip=None) /// One accumulated output batch was yielded, repeat to continue producing @@ -690,6 +931,13 @@ impl Stream for PartialHashAggregateStream { /// -> Done /// Input was exhausted. /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// /// Done /// -> (end) /// ``` @@ -707,12 +955,21 @@ impl Stream for PartialHashAggregateStream { state @ PartialHashAggregateState::ReadingInput { .. } => { self.handle_reading_input(cx, state) } + state @ PartialHashAggregateState::EmittingOnMemoryPressure { .. } => { + self.handle_emitting_on_memory_pressure(state) + } state @ PartialHashAggregateState::ProducingOutput { .. } => { self.handle_producing_output(state) } state @ PartialHashAggregateState::SkippingAggregation { .. } => { self.handle_skipping_aggregation(cx, state) } + state @ PartialHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } state @ PartialHashAggregateState::Done => { let _ = self.reservation.try_resize(0); self.state = Some(state); @@ -725,6 +982,16 @@ impl Stream for PartialHashAggregateStream { self.state = Some(next_state); continue; } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!(next_state, PartialHashAggregateState::Error)); + + // The handler has already discarded its state-owned resources. + // Release the remaining stream-owned resources before returning. + self.close_input(); + self.reservation.free(); + self.state = Some(PartialHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } ControlFlow::Break((poll, next_state)) => { self.state = Some(next_state); return poll; @@ -754,11 +1021,10 @@ impl FinalHashAggregateStream { let schema = Arc::clone(&agg.schema); let input = agg.input.execute(partition, Arc::clone(context))?; + 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 spill_metrics = SpillMetrics::new(&agg.metrics, partition); let hash_table = AggregateHashTable::::new( agg, @@ -767,8 +1033,23 @@ impl FinalHashAggregateStream { batch_size, )?; + let can_spill = context.runtime_env().disk_manager.tmp_files_enabled(); + let spill_context = if can_spill { + Some(Box::new(FinalSpillContext::new( + agg, + context, + partition, + batch_size, + &input_schema, + spill_metrics, + )?)) + } else { + None + }; + let reservation = MemoryConsumer::new(format!("FinalHashAggregateStream[{partition}]")) + .with_can_spill(can_spill) .register(context.memory_pool()); Ok(Self { @@ -777,10 +1058,31 @@ impl FinalHashAggregateStream { baseline_metrics, reservation, group_values_soft_limit: agg.limit_options().map(|config| config.limit()), - state: Some(FinalHashAggregateState::ReadingInput { hash_table }), + state: Some(FinalHashAggregateState::ReadingInput { + hash_table, + spill_context, + }), }) } + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_err(error: DataFusionError) -> FinalHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + FinalHashAggregateState::Error, + )) + } + + fn break_with_internal_err( + message: impl std::fmt::Display, + ) -> FinalHashAggregateStateTransition { + Self::break_with_err(internal_datafusion_err!("{message}")) + } + /// See comments in [`Self::group_values_soft_limit`] for details. fn hit_soft_group_limit(&self, hash_table: &AggregateHashTable) -> bool { self.group_values_soft_limit @@ -791,11 +1093,30 @@ impl FinalHashAggregateStream { &mut self, hash_table: &mut AggregateHashTable, ) -> Result<()> { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + self.close_input(); hash_table.start_output() } + /// Reserve memory for the current aggregate table. + fn reservation_size_for_table( + hash_table: &AggregateHashTable, + spill_context: Option<&FinalSpillContext>, + ) -> usize { + let table_size = hash_table.memory_size(); + if spill_context.is_some() { + // Count extra space needed for in-memory sorting and spilling. Only + // count memory for indices, the payload will be materialize incrementally + // in smaller chunks. + table_size.saturating_add( + hash_table + .building_group_count() + .saturating_mul(size_of::()), + ) + } else { + table_size + } + } + /// Handle ReadingInput state - aggregate partial state batches into the hash table. /// /// See comments at `poll_next()` for details. @@ -804,77 +1125,254 @@ impl FinalHashAggregateStream { fn handle_reading_input( &mut self, cx: &mut Context<'_>, - mut original_state: FinalHashAggregateState, + original_state: FinalHashAggregateState, ) -> FinalHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - FinalHashAggregateState::ReadingInput { .. } - )); - debug_assert!(original_state.hash_table().is_building()); + let FinalHashAggregateState::ReadingInput { + mut hash_table, + spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected ReadingInput state", + ); + }; match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)), + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + FinalHashAggregateState::ReadingInput { + hash_table, + spill_context, + }, + )), Poll::Ready(Some(Ok(batch))) => { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = original_state.hash_table_mut().aggregate_batch(&batch); + let result = hash_table.aggregate_batch(&batch); timer.done(); if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + return Self::break_with_err(e); } - if self.hit_soft_group_limit(original_state.hash_table()) { + if self.hit_soft_group_limit(&hash_table) { let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut()); + let result = self.start_output(&mut hash_table); timer.done(); - if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } - - return ControlFlow::Continue(original_state.into_producing_output()); + return match result { + Ok(()) => ControlFlow::Continue( + FinalHashAggregateState::ProducingOutput { hash_table }, + ), + Err(e) => Self::break_with_err(e), + }; } - if let Err(e) = self - .reservation - .try_resize(original_state.hash_table().memory_size()) - { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + // Check memory reservation, and potentially spill. + let timer = elapsed_compute.timer(); + let resize_result = + self.reservation + .try_resize(Self::reservation_size_for_table( + &hash_table, + spill_context.as_deref(), + )); + timer.done(); + match resize_result { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + // OOM and don't support spilling from configuration + let Some(spill_context) = spill_context else { + return Self::break_with_err(e.context( + "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", + )); + }; + // Sanity check: impossible to OOM when there is no group aggregated. + if hash_table.building_group_count() == 0 { + return Self::break_with_internal_err( + "Final hash aggregate ran out of memory with no aggregated groups", + ); + } + // Go to the next state to perform spilling the aggregated + // groups so far. + return ControlFlow::Continue( + FinalHashAggregateState::Spilling { + hash_table, + spill_context, + }, + ); + } + Err(e) => return Self::break_with_err(e), } - ControlFlow::Continue(original_state) - } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + ControlFlow::Continue(FinalHashAggregateState::ReadingInput { + hash_table, + spill_context, + }) } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + // Input done, move to next state: + // - If spilled before, perform merging spill runs + // - If not spilled, start producing outputs Poll::Ready(None) => { - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut()); - timer.done(); - - match result { - Ok(()) => { - ControlFlow::Continue(original_state.into_producing_output()) + self.close_input(); + match spill_context { + Some(spill_context) if spill_context.has_spills() => { + ControlFlow::Continue( + FinalHashAggregateState::PreparingMergeInput { + hash_table, + spill_context, + }, + ) } - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + _ => { + let elapsed_compute = + self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.start_output(); + timer.done(); + + match result { + Ok(()) => ControlFlow::Continue( + FinalHashAggregateState::ProducingOutput { hash_table }, + ), + Err(e) => Self::break_with_err(e), + } } } } } } + /// 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: FinalHashAggregateState, + ) -> FinalHashAggregateStateTransition { + let FinalHashAggregateState::Spilling { + mut hash_table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected Spilling state", + ); + }; + + // Sanity check: it is impossible to OOM when the table is empty. + if hash_table.building_group_count() == 0 { + return Self::break_with_internal_err( + "Final hash aggregation entered Spilling with an empty table", + ); + } + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let mut result = spill_context.spill_table(&mut hash_table); + + // Spilling shrinks the aggregate table and releases its accumulated + // memory. Update the reservation accordingly. + if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) { + result = + Err(e.context("Decreasing allocation after spilling should succeed")); + } + + timer.done(); + + match result { + // Finished spilling the aggregate table, continue aggregating from input. + Ok(()) => ControlFlow::Continue(FinalHashAggregateState::ReadingInput { + hash_table, + spill_context: Some(spill_context), + }), + Err(e) => Self::break_with_err(e), + } + } + + /// 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 final 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: FinalHashAggregateState, + ) -> FinalHashAggregateStateTransition { + let FinalHashAggregateState::PreparingMergeInput { + mut hash_table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected PreparingMergeInput state", + ); + }; + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let replay = match spill_context.spill_table(&mut hash_table) { + Ok(()) => { + let group_by_metrics = hash_table.group_by_metrics().clone(); + drop(hash_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(FinalHashAggregateState::MergingSpills { stream }) + } + Err(e) => Self::break_with_err(e), + } + } + + /// 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: FinalHashAggregateState, + ) -> FinalHashAggregateStateTransition { + let FinalHashAggregateState::MergingSpills { mut stream } = original_state else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected MergingSpills state", + ); + }; + + match stream.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + FinalHashAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( + Poll::Ready(Some(Ok(batch))), + FinalHashAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => ControlFlow::Continue(FinalHashAggregateState::Done), + } + } + /// Handle ProducingOutput state - emit final aggregate value batches. /// /// See comments at `poll_next()` for details. @@ -882,29 +1380,34 @@ impl FinalHashAggregateStream { /// Returns the next operator state with control flow decision. fn handle_producing_output( &mut self, - mut original_state: FinalHashAggregateState, + original_state: FinalHashAggregateState, ) -> FinalHashAggregateStateTransition { - debug_assert!(matches!( - &original_state, - FinalHashAggregateState::ProducingOutput { .. } - )); - debug_assert!(!original_state.hash_table().is_building()); + let FinalHashAggregateState::ProducingOutput { mut hash_table } = original_state + else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected ProducingOutput state", + ); + }; let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = original_state.hash_table_mut().next_output_batch(); + let result = hash_table.next_output_batch(); timer.done(); match result { Ok(Some(batch)) => { - let _ = self - .reservation - .try_resize(original_state.hash_table().memory_size()); - debug_assert!(batch.num_rows() > 0); - let next_state = if original_state.hash_table().is_done() { - original_state.into_done() + let next_state = if hash_table.is_done() { + drop(hash_table); + if let Err(e) = self.reservation.try_resize(0) { + return Self::break_with_err(e); + } + FinalHashAggregateState::Done } else { - original_state + if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) + { + return Self::break_with_err(e); + } + FinalHashAggregateState::ProducingOutput { hash_table } }; ControlFlow::Break(( @@ -912,11 +1415,15 @@ impl FinalHashAggregateStream { next_state, )) } + Err(e) => Self::break_with_err(e), Ok(None) => { - let _ = self.reservation.try_resize(0); - ControlFlow::Continue(original_state.into_done()) + drop(hash_table); + let next_state = FinalHashAggregateState::Done; + if let Err(e) = self.reservation.try_resize(0) { + return Self::break_with_err(e); + } + ControlFlow::Continue(next_state) } - Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)), } } } @@ -938,21 +1445,49 @@ impl Stream for FinalHashAggregateStream { /// /// ReadingInput /// -> ReadingInput - /// Aggregate one partial-state input batch, update the inner aggregate - /// hash table, and continue with the next input batch. - /// + /// Aggregate one partial-state input batch. If it fits in memory, + /// continue with the next input batch. + /// -> Spilling + /// The table cannot reserve enough memory. Move all current states into + /// one fully group-key-sorted spill run. /// -> ProducingOutput - /// Input was exhausted, or the soft group limit was reached. Move to - /// the next state to start outputting final aggregate values. + /// Input was exhausted without spilling, or the soft group limit was + /// reached. Start outputting final aggregate values. + /// -> PreparingMergeInput + /// Input was exhausted after spilling. Spill the last in-memory run and + /// construct the ordered input used to merge all spill files. + /// + /// Spilling + /// -> ReadingInput + /// One sorted run was written; resume reading the original input. + /// + /// PreparingMergeInput + /// Spill the final in-memory run and build the input ordered replay stream. + /// -> MergingSpills + /// The final run was spilled and the ordered replay stream was built. + /// + /// MergingSpills + /// Aggregate the merged spill runs and emit final results. + /// -> MergingSpills + /// Forward one result batch from the fully ordered replay stream that + /// consumes the sort-preserving merge. + /// -> Done + /// The merged spill input was fully aggregated. /// /// ProducingOutput /// -> ProducingOutput /// One final output batch was yielded; repeat to continue producing /// output incrementally. - /// /// -> Done /// All final output was emitted. /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// /// Done /// -> (end) /// ``` @@ -970,9 +1505,24 @@ impl Stream for FinalHashAggregateStream { state @ FinalHashAggregateState::ReadingInput { .. } => { self.handle_reading_input(cx, state) } + state @ FinalHashAggregateState::Spilling { .. } => { + self.handle_spilling(state) + } + state @ FinalHashAggregateState::PreparingMergeInput { .. } => { + self.handle_preparing_merge_input(state) + } + state @ FinalHashAggregateState::MergingSpills { .. } => { + self.handle_merging_spills(cx, state) + } state @ FinalHashAggregateState::ProducingOutput { .. } => { self.handle_producing_output(state) } + state @ FinalHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } state @ FinalHashAggregateState::Done => { let _ = self.reservation.try_resize(0); self.state = Some(state); @@ -985,6 +1535,16 @@ impl Stream for FinalHashAggregateStream { self.state = Some(next_state); continue; } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!(next_state, FinalHashAggregateState::Error)); + + // The handler has already discarded its state-owned resources. + // Release the remaining stream-owned resources before returning. + self.close_input(); + self.reservation.free(); + self.state = Some(FinalHashAggregateState::Error); + 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/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 33860d3f51c0b..ad5ee3db17969 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1221,12 +1221,7 @@ impl AggregateExec { )?)) } - fn should_use_partial_hash_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_partial_hash_stream(&self, _context: &TaskContext) -> bool { self.mode == AggregateMode::Partial && self.input_order_mode == InputOrderMode::Linear && !self.group_by.is_true_no_grouping() @@ -1245,12 +1240,7 @@ impl AggregateExec { && self.limit_options_supported_by_hash_stream() } - fn should_use_final_hash_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_final_hash_stream(&self, _context: &TaskContext) -> bool { matches!( self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned @@ -3390,7 +3380,8 @@ mod tests { | 2 | 1 | 1.0 | | 3 | 1 | 2.0 | | 3 | 2 | 5.0 | - | 4 | 3 | 11.0 | + | 4 | 1 | 4.0 | + | 4 | 2 | 7.0 | +---+---------------+-------------+ "); } @@ -3428,7 +3419,7 @@ mod tests { let task_ctx = if spill { // enlarge memory limit to let the final aggregation finish - new_spill_ctx(2, 2600) + new_spill_ctx(2, 4640) } else { Arc::clone(&task_ctx) }; @@ -3457,17 +3448,12 @@ mod tests { let spilled_bytes = metrics.spilled_bytes().unwrap(); let spilled_rows = metrics.spilled_rows().unwrap(); + assert_eq!(3, output_rows); if spill { - // When spilling, the output rows metrics become partial output size + final output size - // This is because final aggregation starts while partial aggregation is still emitting - assert_eq!(8, output_rows); - assert!(spill_count > 0); assert!(spilled_bytes > 0); assert!(spilled_rows > 0); } else { - assert_eq!(3, output_rows); - assert_eq!(0, spill_count); assert_eq!(0, spilled_bytes); assert_eq!(0, spilled_rows); @@ -4495,7 +4481,7 @@ mod tests { async fn run_first_last_multi_partitions() -> Result<()> { for is_first_acc in [false, true] { for spill in [false, true] { - first_last_multi_partitions(is_first_acc, spill, 4200).await? + first_last_multi_partitions(is_first_acc, spill, 5000).await? } } Ok(()) diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 26f644d8b62e2..19deedc258c46 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -230,6 +230,10 @@ impl OrderedFinalSpillContext { } = 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) @@ -237,7 +241,7 @@ impl OrderedFinalSpillContext { .with_expressions(&spill_expr) .with_metrics(baseline_metrics.intermediate()) .with_batch_size(batch_size) - .with_reservation(reservation) + .with_reservation(merge_reservation) .build()?; let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( &agg, @@ -248,6 +252,7 @@ impl OrderedFinalSpillContext { baseline_metrics.clone(), group_by_metrics, None, + reservation, )?; Ok(Box::pin(replay)) } @@ -279,6 +284,15 @@ impl OrderedFinalAggregateStream { 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); + let reservation = + MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) + // HACK: Technically, fully ordered aggregate is a non-spillable + // consumer, since it uses bounded memory. There is a known race + // condition bug, and we set it to spillable to let it have larger + // memory budget to suppress the bug. + // Bug issue: https://github.com/apache/datafusion/issues/17334 + .with_can_spill(true) + .register(context.memory_pool()); Self::new_with_input_and_metrics( agg, context, @@ -288,6 +302,7 @@ impl OrderedFinalAggregateStream { baseline_metrics, group_by_metrics, Some(spill_metrics), + reservation, ) } @@ -295,6 +310,9 @@ impl OrderedFinalAggregateStream { clippy::too_many_arguments, reason = "keeps replay metric reuse explicit" )] + /// Builds the stream with the reservation of its logical aggregate operator. + /// Replay callers pass a sibling of the reservation used by the merge input, + /// keeping both components under one memory-consumer registration. pub(in crate::aggregates) fn new_with_input_and_metrics( agg: &AggregateExec, context: &Arc, @@ -304,6 +322,7 @@ impl OrderedFinalAggregateStream { baseline_metrics: BaselineMetrics, group_by_metrics: GroupByMetrics, spill_metrics: Option, + reservation: MemoryReservation, ) -> Result { debug_assert!(matches!( agg.mode, @@ -342,11 +361,6 @@ impl OrderedFinalAggregateStream { input_order_mode, group_by_metrics, )?; - let reservation = - MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) - .with_can_spill(can_spill) - .register(context.memory_pool()); - Ok(Self { schema, input, diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 2917b960f6431..c6f25dc2cf28b 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -281,6 +281,10 @@ impl SingleSpillContext { } = 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) @@ -288,7 +292,7 @@ impl SingleSpillContext { .with_expressions(&spill_expr) .with_metrics(baseline_metrics.intermediate()) .with_batch_size(batch_size) - .with_reservation(reservation) + .with_reservation(merge_reservation) .build()?; let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( &final_agg, @@ -299,6 +303,7 @@ impl SingleSpillContext { baseline_metrics.clone(), group_by_metrics, None, + reservation, )?; Ok(Box::pin(replay)) } diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index cce3a3e903cdf..3dbf880fd1fa9 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -194,6 +194,36 @@ FROM ( 04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[array_agg(t.v)], metrics=[spill_count=7,] +# --- Case G: partial/final aggregation under memory limit --- +statement ok +SET datafusion.execution.target_partitions = 4 + +query II +SELECT count(*), sum(total) +FROM ( + SELECT (v * 7) % 100000 AS k, sum(v) AS total + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- +100000 5000050000 + +# Assert spill happened in the final aggregation. +# In multi-partitions configuration, 'spilled_rows' is not deterministic, so assert +# the unit to be 'K' +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(total) +FROM ( + SELECT (v * 7) % 100000 AS k, sum(v) AS total + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- + +06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v * Int64(7) % Int64(100000)@0 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[spilled_rows=K,] + + # Restore settings to slt runner defaults statement ok RESET datafusion.runtime.memory_limit @@ -201,8 +231,5 @@ RESET datafusion.runtime.memory_limit 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