Skip to content

fix: emit completed groups when fully sorted final and single aggrega… #26637

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yashaggarwal85d wants to merge 1 commit into
base: main
Choose a base branch
from
+190 −8
Open

fix: emit completed groups when fully sorted final and single aggrega… #26637

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 98 additions & 3 deletions datafusion/physical-plan/src/aggregates/clustered_final_stream.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,7 +261,8 @@ impl ClusteredFinalAggregateStream {
/// 1. Start.
/// 2. Merge one input batch:
/// - If memory fits and no groups are complete, continue reading input.
/// - If OOM, spill.
/// - If OOM, spill. Without spilling, emit completed groups through [3],
/// and fail only when no group is complete.
/// 3. Prepare output:
/// - Before any spill, a prefix of groups is complete: materialize the
/// entire prefix once, retaining the input and active groups to resume
Expand Down Expand Up @@ -340,7 +341,15 @@ impl Aggregating {
Ok(()) => {}
Err(oom @ DataFusionError::ResourcesExhausted(_)) => {
let Some(spill_context) = self.spill_context.as_mut() else {
return Err(oom);
// Without spilling, emitting completed groups is the only

@comphead comphead Oct 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Design question for both streams. The fix uses a refused try_resize as the trigger and then emits. The issue points at the legacy GroupedHashAggregateStream, which emits completed groups before it updates the reservation. For the non-spilling case, taking the completed groups first and then sizing the reservation once would not need an error path, and a failure would then mean the active group itself does not fit. Is there a reason to prefer the error-driven fallback, for example keeping the success path untouched? If so, a short comment saying why would help. Note that neither variant lowers the real peak, since aggregate_batch has already allocated the new groups. It only makes the accounting consistent.

// way to release memory. Fail only when none is complete.
let batch =
self.table.take_completed_result_batch()?.ok_or(oom)?;
timer.done();
return Ok(Some(ExecutionStage::Outputting(Outputting {
batch,
resume: Some(self),
})));
};
if self.table.is_empty() {
return Err(oom);
Expand Down Expand Up @@ -464,8 +473,11 @@ mod tests {
use crate::test::TestMemoryExec;
use arrow::array::{Int64Array, StringViewArray};
use arrow::datatypes::{DataType, Field, Schema};
use datafusion_common::test_util::batches_to_string;
use datafusion_execution::config::SessionConfig;
use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool};
use datafusion_execution::memory_pool::{
GreedyMemoryPool, MemoryPool, UnboundedMemoryPool,
};
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
use datafusion_functions_aggregate::{min_max::min_udaf, sum::sum_udaf};
use datafusion_physical_expr::PhysicalSortExpr;
Expand Down Expand Up @@ -732,4 +744,87 @@ mod tests {
assert_eq!(pool.reserved(), 0);
Ok(())
}

/// Runs a final `SUM(k) GROUP BY k` over two partial-state batches sorted on
/// `k` that share a boundary key. Returns the output and the most memory the
/// stream held while handing off output.
async fn run_fully_sorted(
pool: Arc<dyn MemoryPool>,
) -> Result<(Vec<RecordBatch>, usize)> {
let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)]));
let groups = PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".into())]);
let expressions = vec![Arc::new(
AggregateExprBuilder::new(sum_udaf(), vec![col("k", &schema)?])
.schema(Arc::clone(&schema))
.alias("sum")
.build()?,
)];
let empty = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?;
let partial = AggregateExec::try_new(
AggregateMode::Partial,
groups.clone(),
expressions.clone(),
vec![None],
empty,
Arc::clone(&schema),
)?;
let partial_schema = partial.schema();
let batches = [0..=256, 256..=512]
.into_iter()
.map(|keys| {
let k = Int64Array::from_iter_values(keys);
RecordBatch::try_new(
Arc::clone(&partial_schema),
vec![Arc::new(k.clone()), Arc::new(k)],
)
})
.collect::<Result<Vec<_>, _>>()?;
let ordering =
LexOrdering::new([PhysicalSortExpr::new_default(col("k", &partial_schema)?)])
.unwrap();
let input =
TestMemoryExec::try_new(&[batches], Arc::clone(&partial_schema), None)?
.try_with_sort_information(vec![ordering])?;
let aggregate = AggregateExec::try_new(
AggregateMode::Final,
groups.as_final(),
expressions,
vec![None],
Arc::new(TestMemoryExec::update_cache(&Arc::new(input))),
schema,
)?;
assert!(matches!(
aggregate.group_clustering_mode(),
GroupClusteringMode::Full

@comphead comphead Oct 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both new tests assert GroupClusteringMode::Full. The changed branch is spill_context == None, which is also true for Partial clustering with temporary files disabled (DiskManager without tmp files). That case reaches the new code with an active key range defined by the sorted prefix, and it is not covered. A case with a two-column key sorted on the first column and tmp files disabled would cover it, for the final and the single stream.

));
let context = Arc::new(
TaskContext::default().with_runtime(
RuntimeEnvBuilder::new()
.with_memory_pool(Arc::clone(&pool))
.build_arc()?,
),
);
let mut stream =
ClusteredFinalAggregateStream::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<()> {

@comphead comphead Oct 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The test only checks that a tight limit now succeeds. Nothing checks the other half of the contract in the docs: the error is still returned when no group can be released. A rerun with a limit just below held (or a single-key input) that expects ResourcesExhausted would show the fix does not swallow genuine exhaustion and exercises the ok_or(oom) and Ok(None) arms.

let (expected, held) =
run_fully_sorted(Arc::new(UnboundedMemoryPool::default())).await?;
let (actual, _) = run_fully_sorted(Arc::new(GreedyMemoryPool::new(held))).await?;
assert_eq!(batches_to_string(&actual), batches_to_string(&expected));
Ok(())
}
}
97 changes: 92 additions & 5 deletions datafusion/physical-plan/src/aggregates/clustered_single_stream.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,9 @@ use crate::{RecordBatchStream, SendableRecordBatchStream};
/// 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, `ClusteredSingleAggregateStream`
/// still emits the completed groups, and returns the error only when no group is
/// complete.
///
/// ## Partial group clustering
///
Expand Down Expand Up @@ -335,7 +336,27 @@ impl ClusteredSingleAggregateStream {
let Some(spill_context) = spill_context else {

@comphead comphead Oct 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new None branch repeats the tail of the success path below: take_completed_result_batch, reservation_size_for_table, start_outputting(.., ReadingInput { table, spill_context }). It also adds two match arms that both call break_with_err, with the second e shadowing the OOM error.

With spill_context == None the has_spills check below is always false, so the success path already does the right thing. Keeping the refused error in an Option and falling through would remove the duplicate block. The only extra logic is in the Ok(None) arm, which returns the saved error instead of continuing. The same shape applies to ClusteredFinalAggregateStream::Aggregating::handle_stage, where it saves a third copy of the Outputting { batch, resume: Some(self) } construction.

// `None` means spilling is not supported, see comments
// at `ClusteredSingleAggregateState` for details.
return Self::break_with_err(e);
// Emitting completed groups is then the only way to
// release memory. Fail only when none is complete.
let timer = elapsed_compute.timer();
let result = table.take_completed_result_batch();
timer.done();
return match result {
Ok(Some(batch)) => {
let table_memory =
Self::reservation_size_for_table(&table, None);
self.start_outputting(
batch,
table_memory,
ClusteredSingleAggregateState::ReadingInput {
table,
spill_context: None,
},
)
}
Ok(None) => Self::break_with_err(e),
Err(e) => Self::break_with_err(e),
};
};
if table.is_empty() {
return Self::break_with_internal_err(
Expand Down Expand Up @@ -766,10 +787,10 @@ mod tests {
use crate::test::TestMemoryExec;
use arrow::array::Int64Array;
use arrow::datatypes::{DataType, Field, Schema};
use datafusion_common::test_util::batches_to_sort_string;
use datafusion_common::test_util::{batches_to_sort_string, batches_to_string};
use datafusion_execution::config::SessionConfig;
use datafusion_execution::memory_pool::{
GreedyMemoryPool, MemoryPool, PeakRecordingPool,
GreedyMemoryPool, MemoryPool, PeakRecordingPool, UnboundedMemoryPool,
};
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
use datafusion_functions_aggregate::sum::sum_udaf;
Expand Down Expand Up @@ -851,4 +872,70 @@ mod tests {
);
Ok(())
}

/// Runs a single `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 while handing off output.
async fn run_fully_sorted(

@comphead comphead Oct 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This driver (context with RuntimeEnvBuilder, the held sampling loop, unlimited run then GreedyMemoryPool::new(held) run, batches_to_string comparison) is now the third copy, after clustered_partial_stream.rs on main and run_fully_sorted in clustered_final_stream.rs. Only the plan construction differs, and the final one needs partial-state input. AggregateExec::execute already dispatches to the clustered stream by mode, since enable_migration_aggregate defaults to true. One shared helper that takes Arc<dyn ExecutionPlan>, runs it on both pools and returns the output and held would leave each module with only its plan builder. The helper could live next to the existing aggregate test utilities in aggregates/mod.rs.

The unit-test form makes sense here since a SQL memory limit cannot target one stream's table size. If a SQL form is feasible, aggregate_memory_spill.slt already sets a 1M limit and notes that sorted generate_series takes the streaming path. I have not tried it, so it may be too fragile.

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::Single,
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(TestMemoryExec::update_cache(&Arc::new(input))),
Arc::clone(&schema),
)?;
assert!(matches!(
aggregate.group_clustering_mode,
GroupClusteringMode::Full
));
let context = Arc::new(
TaskContext::default().with_runtime(
RuntimeEnvBuilder::new()
.with_memory_pool(Arc::clone(&pool))
.build_arc()?,
),
);
let mut stream = ClusteredSingleAggregateStream::new(&aggregate, &context, 0)?;
// 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 fully_sorted_completed_groups_under_memory_pressure() -> Result<()> {
let (expected, held) =
run_fully_sorted(Arc::new(UnboundedMemoryPool::default())).await?;
let (actual, _) = run_fully_sorted(Arc::new(GreedyMemoryPool::new(held))).await?;
assert_eq!(batches_to_string(&actual), batches_to_string(&expected));
Ok(())
}
}
Loading