You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/23 11:05:39 UTC
(datafusion) 01/01: fix: emit completed groups when a fully sorted partial aggregate runs out of memory (#26576)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-26576-6a280f15786f59b7cca9fcef1a0ddb3c3e931a07
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit cd17d8ea053af1d55536d640324ad9dc0a5c862c
Author: Oleks V <co...@users.noreply.github.com>
AuthorDate: Fri Oct 23 11:05:15 2026 +0000
fix: emit completed groups when a fully sorted partial aggregate runs out of memory (#26576)
## Which issue does this PR close?
- Closes #26575.
## Rationale for this change
A `Partial` aggregate over input sorted on all grouping keys fails with
`ResourcesExhausted` in a shared memory pool, although emitting the
completed groups would leave only the active group in memory. This is a
regression from 54, see the issue for details.
## What changes are included in this PR?
When the reservation after an input batch is refused under full group
clustering, `ClusteredPartialAggregateStream` now emits the completed
groups, the same rows the success path emits. It returns the error only
when no group is complete.
## What is the testing strategy for this PR?
New unit test `completed_groups_under_memory_pressure` in
`clustered_partial_stream.rs`. It runs the stream with a limit that fits
the table between batches but not after aggregating a batch, and
compares the output with an unlimited run. It fails on `main` with
`ResourcesExhausted`. It is a unit test because a SQL memory limit
cannot target one stream's table size.
## Are there any user-facing changes?
These queries now complete instead of failing. No API changes.
---
.../src/aggregates/clustered_partial_stream.rs | 92 ++++++++++++++++++++--
1 file changed, 85 insertions(+), 7 deletions(-)
diff --git a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs
index f7f7e2d903..f60f7ed30c 100644
--- a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs
@@ -85,8 +85,9 @@ use crate::{SendableRecordBatchStream, metrics};
/// only the current group remains active after completed groups are emitted, and
/// memory usage does not grow with the total number of groups.
///
-/// If a memory reservation nevertheless fails, the stream returns the error
-/// directly, indicating an unexpected behavior.
+/// If the reservation after an input batch fails, `ClusteredPartialAggregateStream`
+/// still emits the completed groups, and returns the error only when no group is
+/// complete.
///
/// ## Partial group clustering
///
@@ -321,12 +322,10 @@ impl Aggregating {
// 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(_)) {
- return Err(oom);
+ Some(self.table.take_completed_state_batch()?.ok_or(oom)?)
+ } else {
+ Some(self.table.take_state_batch()?.ok_or(oom)?)
}
- let Some(batch) = self.table.take_state_batch()? else {
- return Err(oom);
- };
- Some(batch)
}
Err(e) => return Err(e),
};
@@ -419,3 +418,82 @@ impl Outputting {
Ok(next_stage)
}
}
+
+#[cfg(test)]
+mod tests {
+ 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 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;
+
+ /// 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<dyn MemoryPool>) -> Result<(Vec<RecordBatch>, 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::<Result<Vec<_>, _>>()?;
+ 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(())
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org