From 5901c570c8f85dbb1359d2962192602e9360b664 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 24 Oct 2026 11:49:34 -0700 Subject: [PATCH 1/6] test: show clustered partial OOM fallback fails on retained table capacity After the completed groups are emitted under memory pressure with full group clustering, only one group remains, but the table keeps the capacity it grew for the whole batch (the single-key hash table, the clustered new_groups scratch and group_indices). Re-reserving that retained size fails in a 100 KiB pool for 8192 distinct keys per batch. This test fails on purpose; it demonstrates the issue and is not a fix. Co-authored-by: Isaac --- .../aggregates/clustered_partial_stream.rs | 87 +++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs index f60f7ed30cdc7..5c870d2028977 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs @@ -496,4 +496,91 @@ mod tests { assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); Ok(()) } + + use arrow::array::ArrayRef; + + const ROWS: usize = 8192; + + /// Runs a partial `SUM(v) GROUP BY ` over three batches sorted on all + /// keys, where every row is a distinct group. Each batch completes all but + /// its last group. + async fn run_distinct_keys( + num_keys: usize, + pool: Arc, + ) -> Result> { + let mut fields: Vec<_> = (0..num_keys) + .map(|i| Field::new(format!("k{i}"), DataType::Int64, false)) + .collect(); + fields.push(Field::new("v", DataType::Int64, false)); + let schema = Arc::new(Schema::new(fields)); + let batches = (0..3) + .map(|b| { + let ids = (b * ROWS) as i64..((b + 1) * ROWS) as i64; + let cols: Vec = (0..=num_keys) + .map(|_| { + Arc::new(Int64Array::from_iter_values(ids.clone())) as ArrayRef + }) + .collect(); + RecordBatch::try_new(Arc::clone(&schema), cols) + }) + .collect::, _>>()?; + let keys = (0..num_keys) + .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 mut stream = + ClusteredPartialAggregateStream::new(&aggregate, &context, 0)?.into_stream(); + let mut output = vec![]; + while let Some(batch) = stream.next().await.transpose()? { + output.push(batch); + } + Ok(output) + } + + /// After the completed groups are emitted on memory pressure only one group + /// remains, but the table keeps the capacity it grew for the whole batch: + /// the single-key hash table (~272 KB for 8192 keys), the clustered + /// `new_groups` scratch and `group_indices` (64 KB each). Re-reserving that + /// retained size fails, although the active group fits easily. + #[tokio::test] + async fn completed_groups_under_memory_pressure_retained_capacity() -> Result<()> { + for num_keys in [1, 2] { + let expected = + run_distinct_keys(num_keys, Arc::new(UnboundedMemoryPool::default())) + .await?; + let actual = + run_distinct_keys(num_keys, Arc::new(GreedyMemoryPool::new(100 * 1024))) + .await + .map_err(|e| e.context(format!("num_keys = {num_keys}")))?; + assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + } + Ok(()) + } } From c7a46b17ace4cb6da08fb7886b86393896bd2d3c Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 24 Oct 2026 12:06:27 -0700 Subject: [PATCH 2/6] test: SQL repro for clustered partial OOM fallback on retained capacity Run a sorted GROUP BY end to end with a pool that limits only the clustered partial aggregate streams, since the downstream merge and final aggregate cannot spill and a shared limit cannot single out the partial stream. With 8192 distinct keys per batch the retained table needs ~344 KB (one key) or ~131 KB (two keys) after the completed groups are emitted, and a 100 KiB limit fails. This test fails on purpose; it demonstrates the issue and is not a fix. Co-authored-by: Isaac --- .../memory_limit/clustered_partial_oom.rs | 171 ++++++++++++++++++ datafusion/core/tests/memory_limit/mod.rs | 1 + 2 files changed, 172 insertions(+) create mode 100644 datafusion/core/tests/memory_limit/clustered_partial_oom.rs diff --git a/datafusion/core/tests/memory_limit/clustered_partial_oom.rs b/datafusion/core/tests/memory_limit/clustered_partial_oom.rs new file mode 100644 index 0000000000000..dba0600e01cc8 --- /dev/null +++ b/datafusion/core/tests/memory_limit/clustered_partial_oom.rs @@ -0,0 +1,171 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A partial aggregate over input sorted on every grouping key must finish +//! under memory pressure once its completed groups are emitted. +//! +//! On a refused reservation the clustered partial stream emits the completed +//! groups, leaving one active group. The table still holds the capacity it +//! grew for the whole batch, so re-reserving it can fail although the active +//! group alone fits easily. +//! +//! The downstream merge and final aggregate cannot spill and need at least +//! a batch of memory themselves, so a shared limit cannot single out the +//! partial stream. The pool here limits only the partial streams. + +use std::fmt::{self, Display}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::MemTable; +use datafusion::physical_plan::displayable; +use datafusion::prelude::*; +use datafusion_common::test_util::batches_to_sort_string; +use datafusion_common::{Result, resources_datafusion_err}; +use datafusion_execution::memory_pool::{ + MemoryPool, MemoryReservation, UnboundedMemoryPool, +}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; + +const ROWS: usize = 8192; +const PARTITIONS: usize = 4; +const PARTIAL_STREAM: &str = "ClusteredPartialAggregateStream"; + +/// Unbounded pool except that each partial aggregate stream may reserve at +/// most `limit` bytes. +#[derive(Debug)] +struct PartialStreamLimitPool { + limit: usize, + used: AtomicUsize, +} + +impl Display for PartialStreamLimitPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}(limit: {})", self.name(), self.limit) + } +} + +impl MemoryPool for PartialStreamLimitPool { + fn name(&self) -> &str { + "partial_stream_limit" + } + + fn grow(&self, _reservation: &MemoryReservation, additional: usize) { + self.used.fetch_add(additional, Ordering::Relaxed); + } + + fn shrink(&self, _reservation: &MemoryReservation, shrink: usize) { + self.used.fetch_sub(shrink, Ordering::Relaxed); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + if reservation.consumer().name().starts_with(PARTIAL_STREAM) + && reservation.size() + additional > self.limit + { + return Err(resources_datafusion_err!( + "Failed to allocate additional {additional} bytes for {} with {} bytes \ + already allocated, limit {}", + reservation.consumer().name(), + reservation.size(), + self.limit + )); + } + self.grow(reservation, additional); + Ok(()) + } + + fn reserved(&self) -> usize { + self.used.load(Ordering::Relaxed) + } +} + +/// Each partition holds three batches sorted on every key, and every row is +/// a distinct group. +fn table(num_keys: usize) -> Result { + let mut fields: Vec<_> = (0..num_keys) + .map(|i| Field::new(format!("k{i}"), DataType::Int64, false)) + .collect(); + fields.push(Field::new("v", DataType::Int64, false)); + let schema = Arc::new(Schema::new(fields)); + let partitions = (0..PARTITIONS) + .map(|p| { + (0..3) + .map(|b| { + let start = ((p * 3 + b) * ROWS) as i64; + let ids = start..start + ROWS as i64; + let columns = (0..=num_keys) + .map(|_| { + Arc::new(Int64Array::from_iter_values(ids.clone())) + as ArrayRef + }) + .collect(); + Ok(RecordBatch::try_new(Arc::clone(&schema), columns)?) + }) + .collect::>>() + }) + .collect::>>()?; + let order = (0..num_keys) + .map(|i| col(format!("k{i}")).sort(true, false)) + .collect(); + Ok(MemTable::try_new(schema, partitions)?.with_sort_order(vec![order])) +} + +async fn run(num_keys: usize, pool: Arc) -> Result> { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(pool) + .build_arc()?; + let config = SessionConfig::new().with_target_partitions(PARTITIONS); + let ctx = SessionContext::new_with_config_rt(config, runtime); + ctx.register_table("t", Arc::new(table(num_keys)?))?; + let keys = (0..num_keys) + .map(|i| format!("k{i}")) + .collect::>() + .join(", "); + let df = ctx + .sql(&format!("select {keys}, sum(v) from t group by {keys}")) + .await?; + let plan = df.clone().create_physical_plan().await?; + let plan = displayable(plan.as_ref()).indent(true).to_string(); + assert!( + plan.contains("mode=Partial") && plan.contains("group_clustering_mode=Full"), + "{plan}" + ); + df.collect().await +} + +/// One retained table needs about 344 KB for a single key and 131 KB for two +/// keys, while the active group needs a few KB. +#[tokio::test] +async fn sorted_partial_aggregate_completes_under_memory_pressure() -> Result<()> { + for num_keys in [1, 2] { + let pool = Arc::new(PartialStreamLimitPool { + limit: 100 * 1024, + used: AtomicUsize::new(0), + }); + let actual = run(num_keys, pool) + .await + .map_err(|e| e.context(format!("num_keys = {num_keys}")))?; + let expected = run(num_keys, Arc::new(UnboundedMemoryPool::default())).await?; + assert_eq!( + batches_to_sort_string(&actual), + batches_to_sort_string(&expected) + ); + } + Ok(()) +} diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index d8814eaa86307..07ac4b55c0ee1 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -20,6 +20,7 @@ use std::num::NonZeroUsize; use std::sync::{Arc, LazyLock}; +mod clustered_partial_oom; #[cfg(feature = "extended_tests")] mod memory_limit_validation; mod nlj_spill_unmatched; From b8d2dd698ae4424d650af1449c1b32c4d7c933e1 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 24 Oct 2026 15:01:19 -0700 Subject: [PATCH 3/6] fix: release the clustered partial aggregate table on memory pressure under full group clustering On a refused reservation under full group clustering, the stream emitted only the completed groups. The remaining table holds one group but keeps the capacity it grew for the batch (the single-key hash table, the clustered new_groups scratch and group_indices), so re-reserving it could still fail with ResourcesExhausted. Use take_state_batch as partial group clustering already does: emit every state, including the active group, and reset the table. The final stage merges the split active group, as it merges states of the same key from different partitions. Co-authored-by: Isaac --- .../memory_limit/clustered_partial_oom.rs | 14 +- .../aggregates/clustered_partial_stream.rs | 204 +++++++++--------- 2 files changed, 103 insertions(+), 115 deletions(-) diff --git a/datafusion/core/tests/memory_limit/clustered_partial_oom.rs b/datafusion/core/tests/memory_limit/clustered_partial_oom.rs index dba0600e01cc8..4cb0e76746ddb 100644 --- a/datafusion/core/tests/memory_limit/clustered_partial_oom.rs +++ b/datafusion/core/tests/memory_limit/clustered_partial_oom.rs @@ -16,12 +16,12 @@ // under the License. //! A partial aggregate over input sorted on every grouping key must finish -//! under memory pressure once its completed groups are emitted. +//! when its table needs only a few KB between batches. //! -//! On a refused reservation the clustered partial stream emits the completed -//! groups, leaving one active group. The table still holds the capacity it -//! grew for the whole batch, so re-reserving it can fail although the active -//! group alone fits easily. +//! On a refused reservation the clustered partial stream emits its states and +//! resets the table. Emitting only the completed groups would leave one active +//! group but keep the capacity the table grew for the whole batch, so +//! re-reserving it would fail although the active group alone fits easily. //! //! The downstream merge and final aggregate cannot spill and need at least //! a batch of memory themselves, so a shared limit cannot single out the @@ -149,8 +149,8 @@ async fn run(num_keys: usize, pool: Arc) -> Result Result<()> { for num_keys in [1, 2] { diff --git a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs index 5c870d2028977..36aab486a9cb2 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs @@ -86,8 +86,10 @@ use crate::{SendableRecordBatchStream, metrics}; /// memory usage does not grow with the total number of groups. /// /// If the reservation after an input batch fails, `ClusteredPartialAggregateStream` -/// still emits the completed groups, and returns the error only when no group is -/// complete. +/// emits all current intermediate states, including the active group, and resets +/// the table, as with partial group clustering below. Emitting only the completed +/// groups would leave the capacity the table grew for the batch reserved. The +/// final stage merges the two state rows of the split active group. /// /// ## Partial group clustering /// @@ -262,8 +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 partial group clustering, materialize all - /// current states, including incomplete groups, and reset the table. + /// - On memory pressure, materialize all current states, including + /// incomplete groups, and reset the table. /// - At EOF, materialize all remaining states and prepare to output. /// 4. Input was exhausted with no remaining groups, directly end. /// 5. Yield one slice without materializing the table again. Keep the shared @@ -319,13 +321,10 @@ 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)?) - } else { - Some(self.table.take_state_batch()?.ok_or(oom)?) - } + // Emit every state, including the active groups, and reset the + // table to release the capacity it grew for this batch. The + // final stage merges the states of a group split across outputs. + Some(self.table.take_state_batch()?.ok_or(oom)?) } Err(e) => return Err(e), }; @@ -421,12 +420,14 @@ impl Outputting { #[cfg(test)] mod tests { + use std::collections::BTreeMap; + 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 datafusion_common::test_util::batches_to_string; + use arrow::array::{ArrayRef, AsArray, Int64Array}; + use arrow::datatypes::{DataType, Field, Int64Type, Schema}; use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, }; @@ -437,93 +438,36 @@ mod tests { 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)]) - }) - .collect::, _>>()?; - let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)? - .try_with_sort_information(vec![ - LexOrdering::new([PhysicalSortExpr::new_default(col("k", &schema)?)]) - .unwrap(), - ])?; - 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], - Arc::new(input), - Arc::clone(&schema), - )?; - let context = Arc::new( - TaskContext::default().with_runtime( - RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(&pool)) - .build_arc()?, - ), - ); - let mut 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)) - } - - /// 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. - #[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)); - Ok(()) - } - - use arrow::array::ArrayRef; - - const ROWS: usize = 8192; - - /// Runs a partial `SUM(v) GROUP BY ` over three batches sorted on all - /// keys, where every row is a distinct group. Each batch completes all but - /// its last group. - async fn run_distinct_keys( - num_keys: usize, - pool: Arc, - ) -> Result> { + /// One batch per range, with `num_keys` key columns and a value column + /// that all hold the range values. + fn batches(num_keys: usize, ranges: &[Range]) -> Result> { let mut fields: Vec<_> = (0..num_keys) .map(|i| Field::new(format!("k{i}"), DataType::Int64, false)) .collect(); fields.push(Field::new("v", DataType::Int64, false)); let schema = Arc::new(Schema::new(fields)); - let batches = (0..3) - .map(|b| { - let ids = (b * ROWS) as i64..((b + 1) * ROWS) as i64; - let cols: Vec = (0..=num_keys) + ranges + .iter() + .map(|range| { + let columns = (0..=num_keys) .map(|_| { - Arc::new(Int64Array::from_iter_values(ids.clone())) as ArrayRef + Arc::new(Int64Array::from_iter_values(range.clone())) as ArrayRef }) .collect(); - RecordBatch::try_new(Arc::clone(&schema), cols) + 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. Returns the output and the most memory the stream held + /// between batches. + async fn run( + num_keys: usize, + batches: Vec, + pool: Arc, + ) -> Result<(Vec, usize)> { + let schema = batches[0].schema(); let keys = (0..num_keys) .map(|i| Ok((col(&format!("k{i}"), &schema)?, format!("k{i}")))) .collect::>>()?; @@ -551,35 +495,79 @@ mod tests { let context = Arc::new( TaskContext::default().with_runtime( RuntimeEnvBuilder::new() - .with_memory_pool(pool) + .with_memory_pool(Arc::clone(&pool)) .build_arc()?, ), ); let mut 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) + Ok((output, held)) + } + + /// Merges the partial `SUM` states by key, as the final stage does. + fn merge(output: &[RecordBatch]) -> BTreeMap, i64> { + let mut merged = BTreeMap::new(); + for batch in output { + let columns: Vec<_> = batch + .columns() + .iter() + .map(|column| column.as_primitive::()) + .collect(); + let (sum, keys) = columns.split_last().unwrap(); + for row in 0..batch.num_rows() { + let key = keys.iter().map(|k| k.value(row)).collect(); + *merged.entry(key).or_default() += sum.value(row); + } + } + merged + } + + fn num_rows(output: &[RecordBatch]) -> usize { + output.iter().map(|batch| batch.num_rows()).sum() + } + + /// The limit fits the table with only the active group, but not the table + /// before the completed groups are emitted. Key 256 spans both batches, so + /// emitting its state early splits it into two state rows, which the final + /// stage merges. + #[tokio::test] + async fn completed_groups_under_memory_pressure() -> Result<()> { + let input = || batches(1, &[0..257, 256..513]); + let (expected, held) = + run(1, input()?, Arc::new(UnboundedMemoryPool::default())).await?; + let (actual, _) = run(1, input()?, Arc::new(GreedyMemoryPool::new(held))).await?; + assert_eq!(num_rows(&actual), num_rows(&expected) + 1); + assert_eq!(merge(&actual), merge(&expected)); + Ok(()) } - /// After the completed groups are emitted on memory pressure only one group - /// remains, but the table keeps the capacity it grew for the whole batch: - /// the single-key hash table (~272 KB for 8192 keys), the clustered - /// `new_groups` scratch and `group_indices` (64 KB each). Re-reserving that - /// retained size fails, although the active group fits easily. + /// Every row is a distinct group, so a batch grows the table for 8192 + /// groups: the single-key hash table (~272 KB), the clustered `new_groups` + /// scratch and `group_indices` (64 KB each). Emitting only the completed + /// groups would keep that capacity reserved and fail this limit, although + /// the active group needs a few KB. #[tokio::test] - async fn completed_groups_under_memory_pressure_retained_capacity() -> Result<()> { + async fn memory_pressure_releases_retained_capacity() -> Result<()> { for num_keys in [1, 2] { - let expected = - run_distinct_keys(num_keys, Arc::new(UnboundedMemoryPool::default())) - .await?; - let actual = - run_distinct_keys(num_keys, Arc::new(GreedyMemoryPool::new(100 * 1024))) - .await - .map_err(|e| e.context(format!("num_keys = {num_keys}")))?; - assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); + let input = || batches(num_keys, &[0..8192, 8192..16384, 16384..24576]); + let (expected, _) = + run(num_keys, input()?, Arc::new(UnboundedMemoryPool::default())).await?; + let (actual, _) = run( + num_keys, + input()?, + Arc::new(GreedyMemoryPool::new(100 * 1024)), + ) + .await + .map_err(|e| e.context(format!("num_keys = {num_keys}")))?; + assert_eq!(merge(&actual), merge(&expected)); } Ok(()) } From 175d521580511ea7a1e8121184b8ee9dbad54361 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 24 Oct 2026 17:26:47 -0700 Subject: [PATCH 4/6] test: update the full clustering refusal snapshot for the partial stream Under full group clustering, the partial stream now emits the state of the one incomplete group on a refusal and finishes, instead of returning the refusal. Co-authored-by: Isaac --- .../cases/aggregate_memory_call_sequence.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs index a9eb4e3ec1f37..71de01462e247 100644 --- a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs +++ b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs @@ -802,11 +802,13 @@ async fn clustered_spills_then_replays() -> Result<()> { Ok(()) } -/// Full clustering: a limit too small for even one group returns the refusal -/// in every clustered stream, because none of them spills under full -/// clustering and the one group never completes. ClusteredPartial emits the -/// completed groups on a refusal when there are any (#26576); ClusteredFinal -/// and ClusteredSingle return the refusal even then, which #26628 reports. +/// Full clustering: a limit too small for even one group. None of the +/// clustered streams spills under full clustering, and the one group never +/// completes. ClusteredPartial emits every state on a refusal, including that +/// group, and resets its table, so it finishes (#26725). ClusteredFinal and +/// ClusteredSingle cannot emit an incomplete group, so they return the +/// refusal; they return it even when some groups are complete, which #26628 +/// reports. #[tokio::test] async fn full_clustering_refusal_is_an_error() -> Result<()> { let _bound = hide_byte_counts().bind_to_scope(); @@ -822,8 +824,10 @@ async fn full_clustering_refusal_is_an_error() -> Result<()> { ClusteredPartialAggregateStream[0] try_grow [bytes] ClusteredPartialAggregateStream[0] try_grow [bytes] -> refused ClusteredPartialAggregateStream[0] shrink [bytes] + ClusteredPartialAggregateStream[0] shrink [bytes] + ClusteredPartialAggregateStream[0] shrink [bytes] ClusteredPartialAggregateStream[0] unregister - => Err(Resources exhausted: Failed to allocate additional [bytes] for ClusteredPartialAggregateStream[0] with [bytes] already allocated for this reservation - [bytes] remain available for the total memory pool: greedy(used: [bytes], pool_size: [bytes])) + => Ok(1 rows in 1 batches) "); let run = one_group(AggregateMode::Final).run().await?; From e7403778efc51e3c5ca359168f93052f52283ced Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 26 Oct 2026 11:58:36 -0700 Subject: [PATCH 5/6] fix: release retained group values capacity instead of resetting the table Resetting the table under full group clustering split the active group into two state rows within one partition. The aggregate still declares the group keys unique and ordered, so a sort on the keys and another column above it could pass rows through out of order. Emit only the completed groups again, as before, and then release the capacity the table kept for them: a new GroupValues::shrink_to_fit, implemented for GroupValuesPrimitive, GroupValuesClustered, GroupValuesColumn and GroupValuesRows, and the group_indices scratch. Merge the two unit tests into one that overlaps batch boundaries and covers all four implementations, restore the refusal snapshot, and drop the SQL test, which duplicated the unit test. Co-authored-by: Isaac --- .../memory_limit/clustered_partial_oom.rs | 171 ----------------- datafusion/core/tests/memory_limit/mod.rs | 1 - .../aggregate_hash_table/common_clustered.rs | 11 ++ .../aggregates/clustered_partial_stream.rs | 172 +++++++++--------- .../src/aggregates/group_values/mod.rs | 10 + .../group_values/multi_group_by/clustered.rs | 6 + .../group_values/multi_group_by/mod.rs | 13 ++ .../src/aggregates/group_values/row.rs | 10 + .../group_values/single_group_by/primitive.rs | 5 + .../cases/aggregate_memory_call_sequence.rs | 16 +- 10 files changed, 142 insertions(+), 273 deletions(-) delete mode 100644 datafusion/core/tests/memory_limit/clustered_partial_oom.rs diff --git a/datafusion/core/tests/memory_limit/clustered_partial_oom.rs b/datafusion/core/tests/memory_limit/clustered_partial_oom.rs deleted file mode 100644 index 4cb0e76746ddb..0000000000000 --- a/datafusion/core/tests/memory_limit/clustered_partial_oom.rs +++ /dev/null @@ -1,171 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! A partial aggregate over input sorted on every grouping key must finish -//! when its table needs only a few KB between batches. -//! -//! On a refused reservation the clustered partial stream emits its states and -//! resets the table. Emitting only the completed groups would leave one active -//! group but keep the capacity the table grew for the whole batch, so -//! re-reserving it would fail although the active group alone fits easily. -//! -//! The downstream merge and final aggregate cannot spill and need at least -//! a batch of memory themselves, so a shared limit cannot single out the -//! partial stream. The pool here limits only the partial streams. - -use std::fmt::{self, Display}; -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; - -use arrow::array::{ArrayRef, Int64Array, RecordBatch}; -use arrow::datatypes::{DataType, Field, Schema}; -use datafusion::datasource::MemTable; -use datafusion::physical_plan::displayable; -use datafusion::prelude::*; -use datafusion_common::test_util::batches_to_sort_string; -use datafusion_common::{Result, resources_datafusion_err}; -use datafusion_execution::memory_pool::{ - MemoryPool, MemoryReservation, UnboundedMemoryPool, -}; -use datafusion_execution::runtime_env::RuntimeEnvBuilder; - -const ROWS: usize = 8192; -const PARTITIONS: usize = 4; -const PARTIAL_STREAM: &str = "ClusteredPartialAggregateStream"; - -/// Unbounded pool except that each partial aggregate stream may reserve at -/// most `limit` bytes. -#[derive(Debug)] -struct PartialStreamLimitPool { - limit: usize, - used: AtomicUsize, -} - -impl Display for PartialStreamLimitPool { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}(limit: {})", self.name(), self.limit) - } -} - -impl MemoryPool for PartialStreamLimitPool { - fn name(&self) -> &str { - "partial_stream_limit" - } - - fn grow(&self, _reservation: &MemoryReservation, additional: usize) { - self.used.fetch_add(additional, Ordering::Relaxed); - } - - fn shrink(&self, _reservation: &MemoryReservation, shrink: usize) { - self.used.fetch_sub(shrink, Ordering::Relaxed); - } - - fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { - if reservation.consumer().name().starts_with(PARTIAL_STREAM) - && reservation.size() + additional > self.limit - { - return Err(resources_datafusion_err!( - "Failed to allocate additional {additional} bytes for {} with {} bytes \ - already allocated, limit {}", - reservation.consumer().name(), - reservation.size(), - self.limit - )); - } - self.grow(reservation, additional); - Ok(()) - } - - fn reserved(&self) -> usize { - self.used.load(Ordering::Relaxed) - } -} - -/// Each partition holds three batches sorted on every key, and every row is -/// a distinct group. -fn table(num_keys: usize) -> Result { - let mut fields: Vec<_> = (0..num_keys) - .map(|i| Field::new(format!("k{i}"), DataType::Int64, false)) - .collect(); - fields.push(Field::new("v", DataType::Int64, false)); - let schema = Arc::new(Schema::new(fields)); - let partitions = (0..PARTITIONS) - .map(|p| { - (0..3) - .map(|b| { - let start = ((p * 3 + b) * ROWS) as i64; - let ids = start..start + ROWS as i64; - let columns = (0..=num_keys) - .map(|_| { - Arc::new(Int64Array::from_iter_values(ids.clone())) - as ArrayRef - }) - .collect(); - Ok(RecordBatch::try_new(Arc::clone(&schema), columns)?) - }) - .collect::>>() - }) - .collect::>>()?; - let order = (0..num_keys) - .map(|i| col(format!("k{i}")).sort(true, false)) - .collect(); - Ok(MemTable::try_new(schema, partitions)?.with_sort_order(vec![order])) -} - -async fn run(num_keys: usize, pool: Arc) -> Result> { - let runtime = RuntimeEnvBuilder::new() - .with_memory_pool(pool) - .build_arc()?; - let config = SessionConfig::new().with_target_partitions(PARTITIONS); - let ctx = SessionContext::new_with_config_rt(config, runtime); - ctx.register_table("t", Arc::new(table(num_keys)?))?; - let keys = (0..num_keys) - .map(|i| format!("k{i}")) - .collect::>() - .join(", "); - let df = ctx - .sql(&format!("select {keys}, sum(v) from t group by {keys}")) - .await?; - let plan = df.clone().create_physical_plan().await?; - let plan = displayable(plan.as_ref()).indent(true).to_string(); - assert!( - plan.contains("mode=Partial") && plan.contains("group_clustering_mode=Full"), - "{plan}" - ); - df.collect().await -} - -/// With 8192 distinct keys per batch, the capacity a table grows for a batch -/// is about 344 KB for a single key and 131 KB for two keys, above the limit. -#[tokio::test] -async fn sorted_partial_aggregate_completes_under_memory_pressure() -> Result<()> { - for num_keys in [1, 2] { - let pool = Arc::new(PartialStreamLimitPool { - limit: 100 * 1024, - used: AtomicUsize::new(0), - }); - let actual = run(num_keys, pool) - .await - .map_err(|e| e.context(format!("num_keys = {num_keys}")))?; - let expected = run(num_keys, Arc::new(UnboundedMemoryPool::default())).await?; - assert_eq!( - batches_to_sort_string(&actual), - batches_to_sort_string(&expected) - ); - } - Ok(()) -} diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index bd0587217b109..7239bcf850255 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -20,7 +20,6 @@ use std::num::NonZeroUsize; use std::sync::{Arc, LazyLock}; -mod clustered_partial_oom; mod nlj_spill_unmatched; mod repartition_mem_limit; mod union_nullable_spill; 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_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs index 36aab486a9cb2..3f00225c0adb4 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs @@ -86,10 +86,10 @@ use crate::{SendableRecordBatchStream, metrics}; /// memory usage does not grow with the total number of groups. /// /// If the reservation after an input batch fails, `ClusteredPartialAggregateStream` -/// emits all current intermediate states, including the active group, and resets -/// the table, as with partial group clustering below. Emitting only the completed -/// groups would leave the capacity the table grew for the batch reserved. The -/// final stage merges the two state rows of the split active group. +/// still emits the completed groups, and returns the error only when no group is +/// 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 /// @@ -264,8 +264,10 @@ 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, materialize all current states, including - /// incomplete groups, and reset the table. + /// - 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. /// 4. Input was exhausted with no remaining groups, directly end. /// 5. Yield one slice without materializing the table again. Keep the shared @@ -321,10 +323,19 @@ impl Aggregating { let output = match reservation.try_resize(self.table.memory_size()) { Ok(()) => self.table.take_completed_state_batch()?, Err(oom @ DataFusionError::ResourcesExhausted(_)) => { - // Emit every state, including the active groups, and reset the - // table to release the capacity it grew for this batch. The - // final stage merges the states of a group split across outputs. - Some(self.table.take_state_batch()?.ok_or(oom)?) + if matches!(self.table.group_clustering(), GroupClustering::Full(_)) { + // 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)?) + } } Err(e) => return Err(e), }; @@ -420,14 +431,14 @@ impl Outputting { #[cfg(test)] mod tests { - use std::collections::BTreeMap; use std::ops::Range; use super::*; use crate::aggregates::PhysicalGroupBy; use crate::test::TestMemoryExec; - use arrow::array::{ArrayRef, AsArray, Int64Array}; + 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, }; @@ -438,37 +449,53 @@ mod tests { use datafusion_physical_expr::expressions::col; use datafusion_physical_expr_common::sort_expr::LexOrdering; - /// One batch per range, with `num_keys` key columns and a value column - /// that all hold the range values. - fn batches(num_keys: usize, ranges: &[Range]) -> Result> { - let mut fields: Vec<_> = (0..num_keys) - .map(|i| Field::new(format!("k{i}"), DataType::Int64, false)) + /// 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 columns = (0..=num_keys) - .map(|_| { - Arc::new(Int64Array::from_iter_values(range.clone())) as ArrayRef + 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(); + .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. Returns the output and the most memory the stream held - /// between batches. + /// on every key. async fn run( - num_keys: usize, batches: Vec, pool: Arc, - ) -> Result<(Vec, usize)> { + ) -> Result> { let schema = batches[0].schema(); - let keys = (0..num_keys) + let keys = (0..schema.fields().len() - 1) .map(|i| Ok((col(&format!("k{i}"), &schema)?, format!("k{i}")))) .collect::>>()?; let ordering = LexOrdering::new( @@ -495,79 +522,42 @@ mod tests { 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)) - } - - /// Merges the partial `SUM` states by key, as the final stage does. - fn merge(output: &[RecordBatch]) -> BTreeMap, i64> { - let mut merged = BTreeMap::new(); - for batch in output { - let columns: Vec<_> = batch - .columns() - .iter() - .map(|column| column.as_primitive::()) - .collect(); - let (sum, keys) = columns.split_last().unwrap(); - for row in 0..batch.num_rows() { - let key = keys.iter().map(|k| k.value(row)).collect(); - *merged.entry(key).or_default() += sum.value(row); - } - } - merged + crate::common::collect(stream).await } - fn num_rows(output: &[RecordBatch]) -> usize { - output.iter().map(|batch| batch.num_rows()).sum() - } - - /// The limit fits the table with only the active group, but not the table - /// before the completed groups are emitted. Key 256 spans both batches, so - /// emitting its state early splits it into two state rows, which the final - /// stage merges. + /// 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 input = || batches(1, &[0..257, 256..513]); - let (expected, held) = - run(1, input()?, Arc::new(UnboundedMemoryPool::default())).await?; - let (actual, _) = run(1, input()?, Arc::new(GreedyMemoryPool::new(held))).await?; - assert_eq!(num_rows(&actual), num_rows(&expected) + 1); - assert_eq!(merge(&actual), merge(&expected)); - Ok(()) - } - - /// Every row is a distinct group, so a batch grows the table for 8192 - /// groups: the single-key hash table (~272 KB), the clustered `new_groups` - /// scratch and `group_indices` (64 KB each). Emitting only the completed - /// groups would keep that capacity reserved and fail this limit, although - /// the active group needs a few KB. - #[tokio::test] - async fn memory_pressure_releases_retained_capacity() -> Result<()> { - for num_keys in [1, 2] { - let input = || batches(num_keys, &[0..8192, 8192..16384, 16384..24576]); - let (expected, _) = - run(num_keys, input()?, Arc::new(UnboundedMemoryPool::default())).await?; - let (actual, _) = run( - num_keys, - input()?, - Arc::new(GreedyMemoryPool::new(100 * 1024)), - ) - .await - .map_err(|e| e.context(format!("num_keys = {num_keys}")))?; - assert_eq!(merge(&actual), merge(&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(()) } 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); diff --git a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs index 71de01462e247..a9eb4e3ec1f37 100644 --- a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs +++ b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs @@ -802,13 +802,11 @@ async fn clustered_spills_then_replays() -> Result<()> { Ok(()) } -/// Full clustering: a limit too small for even one group. None of the -/// clustered streams spills under full clustering, and the one group never -/// completes. ClusteredPartial emits every state on a refusal, including that -/// group, and resets its table, so it finishes (#26725). ClusteredFinal and -/// ClusteredSingle cannot emit an incomplete group, so they return the -/// refusal; they return it even when some groups are complete, which #26628 -/// reports. +/// Full clustering: a limit too small for even one group returns the refusal +/// in every clustered stream, because none of them spills under full +/// clustering and the one group never completes. ClusteredPartial emits the +/// completed groups on a refusal when there are any (#26576); ClusteredFinal +/// and ClusteredSingle return the refusal even then, which #26628 reports. #[tokio::test] async fn full_clustering_refusal_is_an_error() -> Result<()> { let _bound = hide_byte_counts().bind_to_scope(); @@ -824,10 +822,8 @@ async fn full_clustering_refusal_is_an_error() -> Result<()> { ClusteredPartialAggregateStream[0] try_grow [bytes] ClusteredPartialAggregateStream[0] try_grow [bytes] -> refused ClusteredPartialAggregateStream[0] shrink [bytes] - ClusteredPartialAggregateStream[0] shrink [bytes] - ClusteredPartialAggregateStream[0] shrink [bytes] ClusteredPartialAggregateStream[0] unregister - => Ok(1 rows in 1 batches) + => Err(Resources exhausted: Failed to allocate additional [bytes] for ClusteredPartialAggregateStream[0] with [bytes] already allocated for this reservation - [bytes] remain available for the total memory pool: greedy(used: [bytes], pool_size: [bytes])) "); let run = one_group(AggregateMode::Final).run().await?; From 521feaa58a2069b2e3d102a95334211c9a698f65 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 27 Oct 2026 10:53:02 -0700 Subject: [PATCH 6/6] fix: release retained table capacity in clustered final and single streams After a refused reservation, ClusteredFinalAggregateStream and ClusteredSingleAggregateStream emit their completed groups, but the table keeps the capacity it grew for the batch, so reserving it can still fail. Call ClusteredAggregateTable::shrink_to_fit after that emit, as the partial stream does. Co-authored-by: Isaac --- .../src/aggregates/clustered_final_stream.rs | 22 +++++++++++++--- .../src/aggregates/clustered_single_stream.rs | 24 +++++++++++++++--- .../src/aggregates/clustered_test_util.rs | 25 +++++++++++++++++++ 3 files changed, 64 insertions(+), 7 deletions(-) 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_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(