diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index 9c29f39427805..9b50670219d3b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -13,6 +13,7 @@ use crate::aggregates::{ AccumulatorItem, AggrDynFilter, AggregateInputMode, AggregateMode, AggregateOutputMode, DynamicFilterAggregateType, aggregate_expressions, aggregate_metric_label, create_accumulators_with_metrics, + incoming_state_batch_reservation_bytes, }; use crate::metrics::{BaselineMetrics, RecordOutput}; use crate::stream::EmptyRecordBatchStream; @@ -386,6 +387,22 @@ impl AggregateStream { loop { let result = match this.input.next().await { Some(Ok(batch)) => { + // Partial states can own the entire input. Reserve the batch + // before merging it, then transfer that reservation to the + // accumulator's state below. + let previous_size = this.reservation.size(); + let batch_memory = + if this.mode.input_mode() == AggregateInputMode::Partial { + incoming_state_batch_reservation_bytes(&batch) + } else { + 0 + }; + if batch_memory > 0 + && let Err(e) = this.reservation.try_grow(batch_memory) + { + this.finished = true; + return Some((Err(e), this)); + } let result = { let elapsed_compute = this.baseline_metrics.elapsed_compute(); let _timer = elapsed_compute.timer(); // Stops on drop @@ -405,12 +422,11 @@ impl AggregateStream { Ok(allocated) }); - // allocate memory - // This happens AFTER we actually used the memory, but simplifies the whole accounting and we are OK with - // overshooting a bit. Also this means we either store the whole record batch or not. - match result - .and_then(|allocated| this.reservation.try_grow(allocated)) - { + // Once the batch is merged, its buffers are either part of + // the accumulator state or no longer retained. + match result.and_then(|allocated| { + this.reservation.try_resize(previous_size + allocated) + }) { Ok(_) => continue, Err(e) => Err(e), } @@ -596,8 +612,11 @@ mod tests { use crate::metrics::{MetricValue, MetricsSet}; use crate::test::TestMemoryExec; use crate::{ExecutionPlan, collect}; - use arrow::array::{Float64Array, Int32Array, UInt32Array}; + use arrow::array::{ArrayRef, Float64Array, Int32Array, Int64Array, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::DataFusionError; + use datafusion_execution::memory_pool::GreedyMemoryPool; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::{ array_agg::array_agg_udaf, min_max::{max_udaf, min_udaf}, @@ -1029,4 +1048,66 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn final_aggregate_reserves_incoming_state_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let aggregate = sum_aggregate(&schema, "v", "SUM(v)")?; + let group_by = PhysicalGroupBy::default(); + let empty = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![Arc::clone(&aggregate)], + vec![None], + empty, + Arc::clone(&schema), + )?; + // 64 rows that retain 1 MB, like the few rows of a `median` state. + let state_schema = partial.schema(); + let values: ArrayRef = + Arc::new(Int64Array::from(vec![1_i64; 131_072]).slice(0, 64)); + let batch = RecordBatch::try_new(Arc::clone(&state_schema), vec![values])?; + let input = TestMemoryExec::try_new_exec( + &[vec![batch]], + Arc::clone(&state_schema), + None, + )?; + let final_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![aggregate], + vec![None], + input, + state_schema, + )?); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(64 * 1024))) + .build_arc()?; + let context = Arc::new(TaskContext::default().with_runtime(runtime)); + let err = collect( + Arc::clone(&final_aggregate) as Arc, + context, + ) + .await + .unwrap_err(); + assert!(matches!( + err.find_root(), + DataFusionError::ResourcesExhausted(_) + )); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(2 * 1024 * 1024))) + .build_arc()?; + let context = Arc::new(TaskContext::default().with_runtime(runtime)); + let output = collect(final_aggregate as Arc, context).await?; + let result = output[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.value(0), 64); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs index 6eeb37a475ffe..9e0b14f548758 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs @@ -14,12 +14,12 @@ use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use futures::stream::StreamExt; -use super::AggregateExec; use super::aggregate_hash_table::{ ClusteredAggregateTable, ClusteredAggregateTableMetrics, FinalMarker, }; use super::order::GroupClusteringMode; use super::spill::AggregateSpill; +use super::{AggregateExec, incoming_state_batch_reservation_bytes}; use crate::SendableRecordBatchStream; use crate::aggregates::AggregateMode; use crate::metrics::{BaselineMetrics, SpillMetrics}; @@ -323,11 +323,32 @@ impl Aggregating { while let Some(batch) = self.input.next().await.transpose()? { let timer = elapsed_compute.timer(); + let batch_memory = incoming_state_batch_reservation_bytes(&batch); + let admission = if batch_memory > 0 { + reservation.try_grow(batch_memory) + } else { + Ok(()) + }; + match admission { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + let Some(spill_context) = self.spill_context.as_mut() else { + return Err(e); + }; + if self.table.is_empty() { + return Err(e); + } + spill_context.sort_and_spill(self.table.take_state_batch()?)?; + reservation.try_resize(self.table.memory_size())?; + reservation.try_grow(batch_memory)?; + } + Err(e) => return Err(e), + } self.table.aggregate_batch(&batch)?; - // Check the reservation before taking completed groups, as the partial - // stream does, so spilling is unchanged. Without spilling, the refusal - // is returned below only if no group is complete. + // Replace the input-batch reservation with the table footprint before + // taking completed groups, as the partial stream does. A table growth + // refusal can still be resolved by emitting completed groups. let oom = match reservation.try_resize(self.reservation_size()) { Ok(()) => None, Err(oom @ DataFusionError::ResourcesExhausted(_)) => { @@ -470,7 +491,7 @@ mod tests { }; use crate::common::collect; use crate::test::TestMemoryExec; - use arrow::array::{Int64Array, StringViewArray}; + use arrow::array::{ArrayRef, Int64Array, StringViewArray}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; @@ -768,4 +789,100 @@ mod tests { let limit = held_memory(&plan).await?; assert_resources_exhausted_under_limit(&plan, limit).await } + + /// The table stays small: only the incoming batch can exceed the limit. + #[tokio::test] + async fn clustered_final_refuses_large_state_batch() -> Result<()> { + let plan = + sum_over_sorted_input(AggregateMode::Final, &["k"], vec![vec![vec![1; 64]]])?; + // 64 rows that retain 1 MB per column, like the few rows of a `median` state. + let state = plan.input().schema(); + let values: ArrayRef = + Arc::new(Int64Array::from(vec![1_i64; 131_072]).slice(0, 64)); + let batch = + RecordBatch::try_new(Arc::clone(&state), vec![Arc::clone(&values), values])?; + let ordering = plan.input().properties().output_ordering().unwrap().clone(); + let input = TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&state), None)? + .try_with_sort_information(vec![ordering])?; + let plan = AggregateExec::try_new( + AggregateMode::Final, + plan.group_expr().clone(), + plan.aggr_expr().to_vec(), + vec![None], + Arc::new(TestMemoryExec::update_cache(&Arc::new(input))), + plan.input_schema(), + )?; + assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full); + assert_resources_exhausted_under_limit(&plan, 1024 * 1024).await?; + assert_output_unchanged_under_limit(&plan, 4 * 1024 * 1024).await + } + + /// Spilling the table makes room for a state batch that did not fit next to it. + #[tokio::test] + async fn clustered_final_spills_to_admit_large_state_batch() -> Result<()> { + let plan = sum_over_sorted_input( + AggregateMode::Final, + &["k1", "k2"], + vec![vec![vec![0; 4000], (0..4000).collect()]], + )?; + assert_eq!( + plan.group_clustering_mode(), + &GroupClusteringMode::Partial(vec![0]) + ); + let state = plan.input().schema(); + let batch = |rows: i64| { + RecordBatch::try_new( + Arc::clone(&state), + vec![ + Arc::new(Int64Array::from(vec![0_i64; rows as usize])) as ArrayRef, + Arc::new(Int64Array::from((0..rows).collect::>())), + Arc::new(Int64Array::from(vec![1_i64; rows as usize])), + ], + ) + }; + // 64 rows that retain 3 MB, like the few rows of a `median` state. + let large = batch(131_072)?.slice(0, 64); + let ordering = plan.input().properties().output_ordering().unwrap().clone(); + let input = TestMemoryExec::try_new( + &[vec![batch(4000)?, large]], + Arc::clone(&state), + None, + )? + .try_with_sort_information(vec![ordering])?; + let plan = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + plan.group_expr().clone(), + plan.aggr_expr().to_vec(), + vec![None], + Arc::new(TestMemoryExec::update_cache(&Arc::new(input))), + plan.input_schema(), + )?); + + // Fits the large batch, but not next to the table of 4000 groups. + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(3150 * 1024))) + .build_arc()?; + let context = + Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().set_bool( + "datafusion.execution.enable_migration_aggregate", + true, + )) + .with_runtime(runtime), + ); + let output = collect(plan.execute(0, context)?).await?; + assert!(plan.metrics().unwrap().spill_count().unwrap() > 0); + let rows: usize = output.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 4000); + let mut total = 0; + for b in &output { + total += arrow::compute::sum( + b.column(2).as_any().downcast_ref::().unwrap(), + ) + .unwrap_or(0); + } + assert_eq!(total, 4064); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index ca85ccb59f6bb..7aadd61127d9d 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -12,14 +12,11 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{ - DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, -}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use futures::stream::{Stream, StreamExt}; -use super::AggregateExec; use super::aggregate_hash_table::{ AggregateHashTable, ClusteredAggregateTableMetrics, FinalMarker, PartialMarker, PartialSkipMarker, @@ -27,6 +24,7 @@ use super::aggregate_hash_table::{ use super::order::GroupClusteringMode; use super::skip_partial::SkipAggregationProbe; use super::spill::AggregateSpill; +use super::{AggregateExec, incoming_state_batch_reservation_bytes}; use crate::metrics::{ BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput, SpillMetrics, }; @@ -667,6 +665,31 @@ impl FinalHashAggregateStream { } } + fn spill_table( + hash_table: &mut AggregateHashTable, + spill_context: Option<&mut Box>, + reservation: &MemoryReservation, + err: DataFusionError, + ) -> Result<()> { + let Some(spill_context) = spill_context else { + return Err(err.context( + "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", + )); + }; + if hash_table.building_group_count() == 0 { + return Err(err); + } + let result = hash_table + .take_state_batch() + .and_then(|state| spill_context.sort_and_spill(state)); + reservation + .try_resize(hash_table.memory_size()) + .map_err(|e| { + e.context("Decreasing allocation after spilling should succeed") + })?; + result + } + /// Read input stream, if no memory, then spill and continue reading - aggregate partial state batches into the hash table. /// /// Spilling: The table cannot reserve enough memory. @@ -680,6 +703,30 @@ impl FinalHashAggregateStream { while let Some(batch) = self.input.next().await.transpose()? { let _timer = elapsed_compute.timer(); + // Partial states may retain input-sized buffers. Account for the + // incoming batch before merging it into the table. + let batch_memory = incoming_state_batch_reservation_bytes(&batch); + let admission = if batch_memory > 0 { + self.reservation.try_grow(batch_memory) + } else { + Ok(()) + }; + match admission { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + // Spill the current table to make room for this batch, then + // retry admission. An individual batch that still cannot fit + // must fail rather than proceed without a reservation. + Self::spill_table( + hash_table, + spill_context.as_mut(), + &self.reservation, + e, + )?; + self.reservation.try_grow(batch_memory)?; + } + Err(e) => return Err(e), + } hash_table.aggregate_batch(&batch)?; // Soft group limits are usually small and rarely coincide with @@ -692,7 +739,8 @@ impl FinalHashAggregateStream { break; } - // Check memory reservation, and potentially spill. + // Replace the input-batch reservation with the table footprint, + // spilling if the table cannot be reserved. let resize_result = self.reservation .try_resize(Self::reservation_size_for_table( @@ -706,39 +754,12 @@ impl FinalHashAggregateStream { // The table cannot reserve enough memory. // Move all current states into one fully group-key-sorted spill run. Err(e @ DataFusionError::ResourcesExhausted(_)) => { - // OOM and don't support spilling from configuration - let spill_context = spill_context.as_mut().ok_or_else(|| 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. - assert_ne_or_internal_err!( - hash_table.building_group_count(), - 0, - "Final hash aggregate ran out of memory with no aggregated groups" - ); - - // Sorts and spills one complete in-memory state run - - // Go to the next state to perform spilling the aggregated - // groups so far. - 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. - self.reservation - .try_resize(hash_table.memory_size()) - .map_err(|e| { - e.context( - "Decreasing allocation after spilling should succeed", - ) - })?; - - result?; - - // One sorted run was written; resume reading the original input. + Self::spill_table( + hash_table, + spill_context.as_mut(), + &self.reservation, + e, + )?; } Err(e) => return Err(e), } @@ -1735,4 +1756,166 @@ mod tests { ); Ok(()) } + + #[tokio::test] + async fn final_hash_reserves_incoming_state_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]); + let aggregate = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?, + ); + let empty = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![Arc::clone(&aggregate)], + vec![None], + empty, + Arc::clone(&schema), + )?; + let state_schema = partial.schema(); + // 64 rows that retain 2 MB, like the few rows of a `median` state. + let values: ArrayRef = + Arc::new(Int64Array::from(vec![1_i64; 131_072]).slice(0, 64)); + let batch = RecordBatch::try_new( + Arc::clone(&state_schema), + vec![Arc::clone(&values), values], + )?; + let input = TestMemoryExec::try_new_exec( + &[vec![batch]], + Arc::clone(&state_schema), + None, + )?; + let final_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![aggregate], + vec![None], + input, + state_schema, + )?); + assert_eq!( + final_aggregate.group_clustering_mode(), + &GroupClusteringMode::None + ); + + let config = SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", true); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(64 * 1024))) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(config.clone()) + .with_runtime(runtime), + ); + let err = collect(final_aggregate.execute(0, context)?) + .await + .unwrap_err(); + assert!(matches!( + err.find_root(), + DataFusionError::ResourcesExhausted(_) + )); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(3 * 1024 * 1024))) + .build_arc()?; + let context = Arc::new( + TaskContext::default() + .with_session_config(config) + .with_runtime(runtime), + ); + let output = collect(final_aggregate.execute(0, context)?).await?; + let result = output[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.value(0), 64); + Ok(()) + } + + /// Spilling the table makes room for a state batch that did not fit next to it. + #[tokio::test] + async fn final_hash_spills_to_admit_large_state_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]); + let aggregate = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?, + ); + let empty = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let state = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![Arc::clone(&aggregate)], + vec![None], + empty, + Arc::clone(&schema), + )? + .schema(); + let batch = |rows: i64| { + RecordBatch::try_new( + Arc::clone(&state), + vec![ + Arc::new(Int64Array::from((0..rows).collect::>())) as ArrayRef, + Arc::new(Int64Array::from(vec![1_i64; rows as usize])), + ], + ) + }; + // 64 rows that retain 2 MB, like the few rows of a `median` state. + let large = batch(131_072)?.slice(0, 64); + let input = TestMemoryExec::try_new_exec( + &[vec![batch(4000)?, large]], + Arc::clone(&state), + None, + )?; + let plan = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![aggregate], + vec![None], + input, + Arc::clone(&state), + )?); + // Fits the large batch, but not next to the table of 4000 groups. + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(2200 * 1024))) + .build_arc()?; + let context = + Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().set_bool( + "datafusion.execution.enable_migration_aggregate", + true, + )) + .with_runtime(runtime), + ); + let output = collect(plan.execute(0, context)?).await?; + assert!(plan.metrics().unwrap().spill_count().unwrap() > 0); + let rows: usize = output.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 4000); + let mut total = 0; + for b in &output { + total += arrow::compute::sum( + b.column(1).as_primitive::(), + ) + .unwrap_or(0); + } + assert_eq!(total, 4064); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 0943c780e117b..471da9f11acc7 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -217,6 +217,24 @@ mod skip_partial; mod spill; mod topk; +/// Batches passing between operators are bounded by `batch_size` rows and are +/// not charged to the memory pool. A partial-state batch is charged only when +/// its rows carry more than this many bytes on average, which means a row holds +/// many input values (`median`, `array_agg`), so its size follows the input. +const MAX_IN_FLIGHT_STATE_BYTES_PER_ROW: usize = 1024; + +/// Below this size a batch is never charged, whatever its row count. +const MIN_CHARGED_STATE_BATCH_BYTES: usize = 64 * 1024; + +fn incoming_state_batch_reservation_bytes(batch: &RecordBatch) -> usize { + let bytes = datafusion_common::utils::memory::get_record_batch_memory_size(batch); + let in_flight_bytes = batch + .num_rows() + .saturating_mul(MAX_IN_FLIGHT_STATE_BYTES_PER_ROW) + .max(MIN_CHARGED_STATE_BATCH_BYTES); + if bytes > in_flight_bytes { bytes } else { 0 } +} + /// Hard-coded seed for aggregations to ensure hash values differ from `RepartitionExec`, avoiding collisions. const AGGREGATION_HASH_SEED: datafusion_common::hash_utils::RandomState = // This seed is chosen to be a large 64-bit number