diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index 1c004cb70b931..3042e828d60a3 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -381,6 +381,14 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// /// May be expensive; check the implementation before calling on hot paths. fn size(&self) -> usize; + + /// Keeps the stored groups but releases capacity they do not need, such as + /// per-batch scratch buffers, so that [`Self::size`] reflects the stored + /// groups. For example, after [`EmitTo::First`] removes most groups, an + /// accumulator may keep a buffer it sized for the last input batch. + /// + /// The default does nothing. + fn shrink_to_fit(&mut self) {} } #[cfg(test)] diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index a2987be91835e..ee2e9b3b78646 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -666,6 +666,14 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { self.allocation_bytes } + fn shrink_to_fit(&mut self) { + let vec_size_pre = self.states.allocated_size(); + self.states.shrink_to_fit(); + self.adjust_allocation(vec_size_pre, self.states.allocated_size()); + let scratch = std::mem::take(&mut self.scratch); + self.free_allocation(scratch.allocated_size()); + } + fn convert_to_state( &self, values: &[ArrayRef], diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 7da68c47bcd2a..2d5f1db0d7cec 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -642,6 +642,11 @@ impl GroupsAccumulator for HllGroupsAccumulator { + self.allocated_bytes + self.hashes.capacity() * size_of::() } + + fn shrink_to_fit(&mut self) { + self.groups.shrink_to_fit(); + self.hashes = Vec::new(); + } } impl Debug for ApproxDistinct { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 698a3bddffe66..f4eec8faac15a 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -829,6 +829,11 @@ impl HashAggregateAccumulator { self.accumulator.size() } + /// See [`GroupsAccumulator::shrink_to_fit`]. + pub(super) fn shrink_to_fit(&mut self) { + self.accumulator.shrink_to_fit(); + } + pub(super) fn update_batch( &mut self, values: &CompactedAccumulatorArgs, diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs index e6c572677bcd3..08af0787cd7c6 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs @@ -299,6 +299,19 @@ impl ClusteredAggregateTable { + self.buffer.group_indices.allocated_size() } + /// Releases the capacity that the buffer holds beyond its stored groups, + /// such as hash table slots, the per-batch `group_indices` scratch and + /// per-batch accumulator scratch, so that [`Self::memory_size`] reflects + /// the remaining groups. + pub(in crate::aggregates) fn shrink_to_fit(&mut self) { + self.buffer.group_values.shrink_to_fit(); + self.buffer.group_indices.clear(); + self.buffer.group_indices.shrink_to_fit(); + for acc in &mut self.buffer.accumulators { + acc.shrink_to_fit(); + } + } + pub(in crate::aggregates) fn metrics(&self) -> ClusteredAggregateTableMetrics { ClusteredAggregateTableMetrics { group_by: self.group_by_metrics.clone(), diff --git a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs index f60f7ed30cdc7..05e3edb879397 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs @@ -87,7 +87,9 @@ use crate::{SendableRecordBatchStream, metrics}; /// /// If the reservation after an input batch fails, `ClusteredPartialAggregateStream` /// still emits the completed groups, and returns the error only when no group is -/// complete. +/// complete. After the emit it releases the capacity the table grew for the batch, +/// such as hash table slots and per-batch scratch buffers, so only the active +/// group stays reserved. /// /// ## Partial group clustering /// @@ -262,6 +264,8 @@ impl ClusteredPartialAggregateStream { /// 3. Prepare output: /// - A prefix of groups is complete: materialize the entire prefix once, /// retaining the input and active groups to resume aggregation. + /// - On memory pressure with full group clustering, materialize the + /// completed prefix and release the capacity the table grew for it. /// - On memory pressure with partial group clustering, materialize all /// current states, including incomplete groups, and reset the table. /// - At EOF, materialize all remaining states and prepare to output. @@ -319,11 +323,17 @@ impl Aggregating { let output = match reservation.try_resize(self.table.memory_size()) { Ok(()) => self.table.take_completed_state_batch()?, Err(oom @ DataFusionError::ResourcesExhausted(_)) => { - // Partial group clustering may have an unbounded active key range. - // The final stage can merge incomplete states emitted here. if matches!(self.table.group_clustering(), GroupClustering::Full(_)) { - Some(self.table.take_completed_state_batch()?.ok_or(oom)?) + // Only the active group remains after the emit. Release the + // capacity the table grew for this batch so that reserving + // the table can succeed. + let batch = + self.table.take_completed_state_batch()?.ok_or(oom)?; + self.table.shrink_to_fit(); + Some(batch) } else { + // Partial group clustering may have an unbounded active key range. + // The final stage can merge incomplete states emitted here. Some(self.table.take_state_batch()?.ok_or(oom)?) } } @@ -421,79 +431,211 @@ impl Outputting { #[cfg(test)] mod tests { + use std::ops::Range; + use super::*; use crate::aggregates::PhysicalGroupBy; use crate::test::TestMemoryExec; - use arrow::array::Int64Array; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::array::{ArrayRef, Float64Array, Int64Array, ListArray}; + use arrow::datatypes::{DataType, Field, Int64Type, Schema}; use datafusion_common::test_util::batches_to_string; use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_expr::AggregateUDF; + use datafusion_functions_aggregate::any_value::any_value_udaf; + use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf; use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::aggregate::{ + AggregateExprBuilder, AggregateFunctionExpr, + }; use datafusion_physical_expr::expressions::col; use datafusion_physical_expr_common::sort_expr::LexOrdering; - /// Runs a partial `SUM(k) GROUP BY k` over two batches sorted on `k` that - /// share a boundary key. Returns the output and the most memory the stream - /// held between batches. - async fn run(pool: Arc) -> Result<(Vec, usize)> { - let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); - let batches = [0..=256, 256..=512] - .into_iter() - .map(|keys| { - let k = Int64Array::from_iter_values(keys); - RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(k)]) + /// One batch per range. Every key column and the value column `v` hold the + /// range values, as the given key types. + fn batches( + key_types: &[DataType], + ranges: &[Range], + ) -> Result> { + let mut fields: Vec<_> = key_types + .iter() + .enumerate() + .map(|(i, data_type)| Field::new(format!("k{i}"), data_type.clone(), false)) + .collect(); + fields.push(Field::new("v", DataType::Int64, false)); + let schema = Arc::new(Schema::new(fields)); + ranges + .iter() + .map(|range| { + let int64: ArrayRef = + Arc::new(Int64Array::from_iter_values(range.clone())); + let mut columns = key_types + .iter() + .map(|data_type| match data_type { + DataType::Int64 => Arc::clone(&int64), + DataType::Float64 => Arc::new(Float64Array::from_iter_values( + range.clone().map(|v| v as f64), + )), + DataType::List(_) => { + Arc::new(ListArray::from_iter_primitive::( + range.clone().map(|v| Some([Some(v)])), + )) + } + _ => unreachable!("unsupported key type {data_type}"), + }) + .collect::>(); + columns.push(int64); + Ok(RecordBatch::try_new(Arc::clone(&schema), columns)?) }) - .collect::, _>>()?; + .collect() + } + + /// Runs a partial `SUM(v) GROUP BY k0, ..` over `batches`, which are sorted + /// on every key. + async fn run( + batches: Vec, + pool: Arc, + ) -> Result> { + let schema = batches[0].schema(); + let sum = AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?; + run_aggregates(batches, vec![Arc::new(sum)], pool).await + } + + /// Runs a partial aggregate of `aggregates` grouped by `k0, ..` over + /// `batches`, which are sorted on every key. + async fn run_aggregates( + batches: Vec, + aggregates: Vec>, + pool: Arc, + ) -> Result> { + let schema = batches[0].schema(); + let keys = (0..schema.fields().len() - 1) + .map(|i| Ok((col(&format!("k{i}"), &schema)?, format!("k{i}")))) + .collect::>>()?; + let ordering = LexOrdering::new( + keys.iter() + .map(|(expr, _)| PhysicalSortExpr::new_default(Arc::clone(expr))), + ) + .unwrap(); let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? - .try_with_sort_information(vec![ - LexOrdering::new([PhysicalSortExpr::new_default(col("k", &schema)?)]) - .unwrap(), - ])?; + .try_with_sort_information(vec![ordering])?; let aggregate = AggregateExec::try_new( AggregateMode::Partial, - PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]), - vec![Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![col("k", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("sum") - .build()?, - )], - vec![None], + PhysicalGroupBy::new_single(keys), + aggregates.clone(), + vec![None; aggregates.len()], Arc::new(input), Arc::clone(&schema), )?; + assert_eq!(aggregate.group_clustering_mode, GroupClusteringMode::Full); let context = Arc::new( TaskContext::default().with_runtime( RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(&pool)) + .with_memory_pool(pool) .build_arc()?, ), ); - let mut stream = + let stream = ClusteredPartialAggregateStream::new(&aggregate, &context, 0)?.into_stream(); - // Each output batch is a single slice, so the stream releases it before - // handing it off and holds only the table. - let mut held = pool.reserved(); - let mut output = vec![]; - while let Some(batch) = stream.next().await.transpose()? { - held = held.max(pool.reserved()); - output.push(batch); - } - Ok((output, held)) + crate::common::collect(stream).await } - /// The limit fits the table with only the active group, but not the table - /// before the completed groups are emitted. The result must not change. + /// Every row is a distinct group, and each batch shares its first key with + /// the previous one, so the active group spans the batches. Aggregating a + /// batch grows the table for about 8192 groups, for example the single-key + /// hash table (~272 KB) and the `group_indices` scratch (64 KB). After the + /// completed groups are emitted, the table must release that capacity, or + /// reserving it fails the limit although the active group needs a few KB. + /// The output must not change. #[tokio::test] async fn completed_groups_under_memory_pressure() -> Result<()> { - let (expected, held) = run(Arc::new(UnboundedMemoryPool::default())).await?; - let (actual, _) = run(Arc::new(GreedyMemoryPool::new(held))).await?; - assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + let list = DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))); + for key_types in [ + // `GroupValuesPrimitive` + vec![DataType::Int64], + // `GroupValuesClustered` + vec![DataType::Int64, DataType::Int64], + // `GroupValuesColumn` + vec![DataType::Float64, DataType::Float64], + // `GroupValuesRows` + vec![DataType::Int64, list], + ] { + let input = || batches(&key_types, &[0..8193, 8192..16385, 16384..24576]); + let expected = + run(input()?, Arc::new(UnboundedMemoryPool::default())).await?; + let actual = run(input()?, Arc::new(GreedyMemoryPool::new(100 * 1024))) + .await + .map_err(|e| e.context(format!("key types {key_types:?}")))?; + assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + } + Ok(()) + } + + /// Like [`completed_groups_under_memory_pressure`], but with accumulators + /// that keep per-batch scratch. For example, each `APPROX_DISTINCT` keeps + /// the hashes of the last batch (64 KB for 8192 rows), and the + /// `GroupsAccumulatorAdapter` of `ANY_VALUE` keeps its per-batch indices + /// (~160 KB). After the completed groups are emitted, the accumulators must + /// release that scratch too, or reserving the table fails the limit + /// although the active group needs a few KB. The output must not change. + #[tokio::test] + async fn accumulator_scratch_under_memory_pressure() -> Result<()> { + type Aggregates = fn(&SchemaRef) -> Result>>; + fn build( + udaf: Arc, + args: &[&str], + schema: &SchemaRef, + ) -> Result> { + let alias = format!("{}({})", udaf.name(), args.join(", ")); + let args = args + .iter() + .map(|name| col(name, schema)) + .collect::>>()?; + Ok(Arc::new( + AggregateExprBuilder::new(udaf, args) + .schema(Arc::clone(schema)) + .alias(alias) + .build()?, + )) + } + let cases: [(&str, Aggregates); 2] = [ + // `HllGroupsAccumulator` + ("approx_distinct", |schema| { + Ok(vec![ + build(approx_distinct_udaf(), &["k0"], schema)?, + build(approx_distinct_udaf(), &["v"], schema)?, + ]) + }), + // `GroupsAccumulatorAdapter` + ("any_value", |schema| { + Ok(vec![build(any_value_udaf(), &["v"], schema)?]) + }), + ]; + for (name, aggregates) in cases { + let input = + || batches(&[DataType::Int64], &[0..8193, 8192..16385, 16384..24576]); + let schema = input()?[0].schema(); + let expected = run_aggregates( + input()?, + aggregates(&schema)?, + Arc::new(UnboundedMemoryPool::default()), + ) + .await?; + let actual = run_aggregates( + input()?, + aggregates(&schema)?, + Arc::new(GreedyMemoryPool::new(100 * 1024)), + ) + .await + .map_err(|e| e.context(format!("aggregates {name}")))?; + assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + } Ok(()) } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 989b8b928beef..a097ed2adad17 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -133,6 +133,16 @@ pub trait GroupValues: Send { /// Clear the contents and shrink the capacity to the size of the batch (free up memory usage) fn clear_shrink(&mut self, num_rows: usize); + + /// Keeps the stored groups but releases capacity they do not need, such as + /// hash table slots and per-batch scratch buffers, so that [`Self::size`] + /// reflects the stored groups. For example, after [`EmitTo::First`] removes + /// most groups, the hash table may keep the slots it grew for them. + /// + /// The default does nothing. + /// + /// [`EmitTo::First`]: datafusion_expr::EmitTo::First + fn shrink_to_fit(&mut self) {} } /// Return a specialized implementation of [`GroupValues`] for the given schema. diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs index d331c510163cc..42cd63ffdb4ae 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/clustered.rs @@ -181,6 +181,12 @@ impl GroupValues for GroupValuesClustered { true } + fn shrink_to_fit(&mut self) { + // Scratch for one `intern` call. + self.new_groups.clear(); + self.new_groups.shrink_to_fit(); + } + fn clear_shrink(&mut self, num_rows: usize) { self.columns = GroupValuesColumn::::build_group_columns(&self.schema) .expect("schema validated by GroupValuesClustered::try_new"); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index 2836b3177eb05..aecbc6ca8b94e 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -1235,6 +1235,19 @@ impl GroupValues for GroupValuesColumn { true } + fn shrink_to_fit(&mut self) { + self.map.shrink_to_fit(|(hash, _)| *hash); + self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); + // Scratch for one `intern` call. + self.hashes_buffer.clear(); + self.hashes_buffer.shrink_to_fit(); + if !STREAMING { + self.emit_group_index_list_buffer.clear(); + self.emit_group_index_list_buffer.shrink_to_fit(); + self.vectorized_operation_buffers.clear_shrink(0); + } + } + fn clear_shrink(&mut self, num_rows: usize) { // Reset to a fresh column-builder vector. The schema was validated // in `try_new`, so rebuilding cannot fail unless something else diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index ee3f36088668d..9caf3270679a7 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -292,6 +292,16 @@ impl GroupValues for GroupValuesRows { true } + fn shrink_to_fit(&mut self) { + self.map.shrink_to_fit(|(hash, _)| *hash); + // Scratch for one `intern` call. + self.hashes_buffer.clear(); + self.hashes_buffer.shrink_to_fit(); + if self.rows_buffer.size() > self.initial_rows_buffer_size { + self.rows_buffer = Self::new_rows_buffer(&self.row_converter); + } + } + fn clear_shrink(&mut self, num_rows: usize) { self.group_values = if num_rows == 0 { None diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index 05f2a26ca07f6..17ae25269b134 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -265,6 +265,11 @@ where true } + fn shrink_to_fit(&mut self) { + self.map.shrink_to_fit(|&(_, hash)| hash); + self.values.shrink_to_fit(); + } + fn clear_shrink(&mut self, num_rows: usize) { self.values.clear(); self.values.shrink_to(num_rows);