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..4517594429f78 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,17 @@ impl ClusteredAggregateTable { + self.buffer.group_indices.allocated_size() } + /// Releases the capacity that the buffer holds beyond its stored groups, + /// such as hash table slots and the per-batch `group_indices` scratch, so + /// that [`Self::memory_size`] reflects the remaining groups. + /// + /// Accumulator capacity is not released. + 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(); + } + 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_final_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs index b5890e889167e..a0d5b4ec6710e 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs @@ -377,6 +377,11 @@ impl Aggregating { None => continue, } }; + if oom.is_some() { + // Release the capacity the table grew for this batch so that + // reserving the remaining groups can succeed. + self.table.shrink_to_fit(); + } timer.done(); return Ok(Some(ExecutionStage::Outputting(Outputting { batch, @@ -472,9 +477,9 @@ mod tests { use crate::aggregates::GroupClusteringMode; use crate::aggregates::PhysicalGroupBy; use crate::aggregates::clustered_test_util::{ - assert_completed_groups_emitted_under_limit, - assert_resources_exhausted_under_limit, full_batches, held_memory, - single_key_batches, sum_over_sorted_input, + assert_completed_groups_emitted_under_limit, assert_output_unchanged_under_limit, + assert_resources_exhausted_under_limit, distinct_key_batches, full_batches, + held_memory, single_key_batches, sum_over_sorted_input, }; use crate::common::collect; use crate::test::TestMemoryExec; @@ -756,6 +761,17 @@ mod tests { assert_completed_groups_emitted_under_limit(&plan).await } + /// The limit fits the active group, but not the capacity the table grows + /// for one batch, so the table must release it after the completed groups + /// are emitted. + #[tokio::test] + async fn completed_groups_release_batch_capacity() -> Result<()> { + let plan = + sum_over_sorted_input(AggregateMode::Final, &["k"], distinct_key_batches())?; + assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full); + assert_output_unchanged_under_limit(&plan, 100 * 1024).await + } + /// A single key never completes a group before the input ends. The limit /// fits only the initial table. #[tokio::test] diff --git a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs index 17ea53e29d2c4..72755bfbdb1b0 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,11 +431,140 @@ impl Outputting { #[cfg(test)] mod tests { + use std::ops::Range; + use super::*; use crate::aggregates::GroupClusteringMode; + use crate::aggregates::PhysicalGroupBy; use crate::aggregates::clustered_test_util::{ assert_completed_groups_emitted_under_limit, full_batches, sum_over_sorted_input, }; + use crate::test::TestMemoryExec; + 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_functions_aggregate::sum::sum_udaf; + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::sort_expr::LexOrdering; + + /// 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() + } + + /// 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 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![ordering])?; + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(keys), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum") + .build()?, + )], + vec![None], + 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(pool) + .build_arc()?, + ), + ); + let stream = + ClusteredPartialAggregateStream::new(&aggregate, &context, 0)?.into_stream(); + crate::common::collect(stream).await + } + + /// 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 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(()) + } #[tokio::test] async fn fully_clustered_completed_groups_under_memory_pressure() -> Result<()> { diff --git a/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs index c83f0f3887bc4..ad35babbcd038 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs @@ -377,7 +377,12 @@ impl ClusteredSingleAggregateStream { match (result, oom) { // Some finalized groups can be emitted. Yield them in // slices, then continue aggregating input. - (Ok(Some(batch)), _) => { + (Ok(Some(batch)), oom) => { + if oom.is_some() { + // Release the capacity the table grew for this batch + // so that reserving the remaining groups can succeed. + table.shrink_to_fit(); + } let table_memory = Self::reservation_size_for_table( &table, spill_context.as_deref(), @@ -770,9 +775,9 @@ mod tests { use super::*; use crate::aggregates::PhysicalGroupBy; use crate::aggregates::clustered_test_util::{ - assert_completed_groups_emitted_under_limit, - assert_resources_exhausted_under_limit, full_batches, held_memory, - single_key_batches, sum_over_sorted_input, + assert_completed_groups_emitted_under_limit, assert_output_unchanged_under_limit, + assert_resources_exhausted_under_limit, distinct_key_batches, full_batches, + held_memory, single_key_batches, sum_over_sorted_input, }; use crate::collect; use crate::test::TestMemoryExec; @@ -871,6 +876,17 @@ mod tests { assert_completed_groups_emitted_under_limit(&plan).await } + /// The limit fits the active group, but not the capacity the table grows + /// for one batch, so the table must release it after the completed groups + /// are emitted. + #[tokio::test] + async fn completed_groups_release_batch_capacity() -> Result<()> { + let plan = + sum_over_sorted_input(AggregateMode::Single, &["k"], distinct_key_batches())?; + assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full); + assert_output_unchanged_under_limit(&plan, 100 * 1024).await + } + /// A single key never completes a group before the input ends. The limit /// fits only the initial table. #[tokio::test] diff --git a/datafusion/physical-plan/src/aggregates/clustered_test_util.rs b/datafusion/physical-plan/src/aggregates/clustered_test_util.rs index ad2006bd9cc59..16185ad9a3e93 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_test_util.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_test_util.rs @@ -48,6 +48,17 @@ pub(super) fn full_batches() -> Vec>> { vec![vec![(0..=256).collect()], vec![(256..=512).collect()]] } +/// Key columns of three batches of 8192 distinct keys sorted on `k`, where each +/// batch shares its first key with the previous one. Aggregating a batch grows +/// the table for about 8192 groups, while the active group needs a few KB. +pub(super) fn distinct_key_batches() -> Vec>> { + vec![ + vec![(0..8193).collect()], + vec![(8192..16385).collect()], + vec![(16384..24576).collect()], + ] +} + /// Key column of two batches that hold the same single key, so no group /// completes before the input ends. pub(super) fn single_key_batches() -> Vec>> { @@ -181,6 +192,20 @@ pub(super) async fn assert_completed_groups_emitted_under_limit( Ok(()) } +/// Runs `plan` without a memory limit, then under `limit`. The result must not +/// change. +pub(super) async fn assert_output_unchanged_under_limit( + plan: &AggregateExec, + limit: usize, +) -> Result<()> { + let (expected, _) = + collect_with_held_memory(plan, Arc::new(UnboundedMemoryPool::default())).await?; + let (actual, _) = + collect_with_held_memory(plan, Arc::new(GreedyMemoryPool::new(limit))).await?; + assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + Ok(()) +} + /// Runs `plan` under `limit`, where emitting completed groups cannot bring the /// table under the limit, so the refused reservation must be returned. pub(super) async fn assert_resources_exhausted_under_limit( 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);