You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "yashaggarwal85d (via GitHub)" <gi...@apache.org> on 2026/10/23 04:43:51 UTC
[PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
yashaggarwal85d opened a new pull request, #26637:
URL: https://github.com/apache/datafusion/pull/26637
## Which issue does this PR close?
- Closes #26628.
## Rationale for this change
A `Final` or `Single` aggregate over input sorted on all grouping keys fails with `ResourcesExhausted` under a memory limit that the legacy `GroupedHashAggregateStream` handles. After each input batch, `ClusteredFinalAggregateStream` and `ClusteredSingleAggregateStream` reserve memory for the whole table before emitting completed groups. If that reservation is refused and the stream cannot spill (full group clustering, or temporary files disabled), it returns the error, even though emitting the completed groups would leave only the active group in memory.
This is the same problem as #26575 for the partial stream, fixed in #26576.
## What changes are included in this PR?
- `ClusteredFinalAggregateStream`: when the reservation is refused and there is no spill context, take the completed groups and move to the output stage, then resume aggregation. The error is returned only when no group is complete.
- `ClusteredSingleAggregateStream`: the same behavior, emitting through the existing `start_outputting` path and resuming `ReadingInput`.
- Updated doc comments to describe the new behavior.
The spilling path (partial group clustering with temporary files enabled) is unchanged.
## What is the testing strategy for this PR?
New unit tests, modeled on the one in #26576:
- `clustered_final_stream::tests::completed_groups_under_memory_pressure`
- `clustered_single_stream::tests::fully_sorted_completed_groups_under_memory_pressure`
Each runs `SUM(k) GROUP BY k` over two batches sorted on `k` that share a boundary key. It first runs with an unbounded pool to measure the most memory the stream holds when it hands off output, then reruns with a `GreedyMemoryPool` of that size and checks that the results match. Both tests fail on `main` with `ResourcesExhausted` and pass with this change. All `aggregates` unit tests in `datafusion-physical-plan` pass, and `cargo fmt` and `cargo clippy --all-targets --all-features -D warnings` on the crate are clean.
## Are there any user-facing changes?
No API changes. Queries that used to fail with `ResourcesExhausted` in this situation now complete.
---
_This solution was developed with AI assistance; I reviewed and tested the changes._
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4379460725
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -732,4 +749,34 @@ mod tests {
assert_eq!(pool.reserved(), 0);
Ok(())
}
+
+ #[tokio::test]
+ async fn fully_clustered_completed_groups_under_memory_pressure() -> Result<()> {
+ let (plan, mode) =
+ sum_over_sorted_input(AggregateMode::Final, &["k"], full_batches())?;
+ assert_eq!(mode, GroupClusteringMode::Full);
+ assert_completed_groups_emitted_under_limit(&plan).await
+ }
+
+ /// A single key never completes a group before the input ends. The limit
+ /// fits only the initial table.
+ #[tokio::test]
+ async fn memory_pressure_without_completed_groups() -> Result<()> {
+ let (plan, _) =
+ sum_over_sorted_input(AggregateMode::Final, &["k"], single_key_batches())?;
+ let limit = held_memory(&plan).await?;
+ assert_resources_exhausted_under_limit(&plan, limit).await
+ }
+
+ /// Without temporary files, partial group clustering takes the same path as
+ /// full clustering. Its group values keep their capacity after emitting, so
+ /// one byte below the unlimited footprint still fails.
+ #[tokio::test]
+ async fn partially_clustered_memory_pressure_without_spilling() -> Result<()> {
Review Comment:
FYI, here is the issue for this - #26777
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead merged PR #26637:
URL: https://github.com/apache/datafusion/pull/26637
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4379126743
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -732,4 +749,34 @@ mod tests {
assert_eq!(pool.reserved(), 0);
Ok(())
}
+
+ #[tokio::test]
+ async fn fully_clustered_completed_groups_under_memory_pressure() -> Result<()> {
+ let (plan, mode) =
+ sum_over_sorted_input(AggregateMode::Final, &["k"], full_batches())?;
+ assert_eq!(mode, GroupClusteringMode::Full);
+ assert_completed_groups_emitted_under_limit(&plan).await
+ }
+
+ /// A single key never completes a group before the input ends. The limit
+ /// fits only the initial table.
+ #[tokio::test]
+ async fn memory_pressure_without_completed_groups() -> Result<()> {
+ let (plan, _) =
+ sum_over_sorted_input(AggregateMode::Final, &["k"], single_key_batches())?;
+ let limit = held_memory(&plan).await?;
+ assert_resources_exhausted_under_limit(&plan, limit).await
+ }
+
+ /// Without temporary files, partial group clustering takes the same path as
+ /// full clustering. Its group values keep their capacity after emitting, so
+ /// one byte below the unlimited footprint still fails.
+ #[tokio::test]
+ async fn partially_clustered_memory_pressure_without_spilling() -> Result<()> {
Review Comment:
Thanks for digging into this and finding the retained capacity. Given that finding, I'd drop this test and its copy in `clustered_single_stream.rs`, along with `partial_batches`.
- From reading the code, I expect it to pass on `main` too. No group completes before batch 3, so the first refusal at `held - 1` comes from the post-batch `try_resize`, and `main` returns it right away. So it does not cover the fix.
- What it asserts (`ResourcesExhausted` is not swallowed) is already covered by `memory_pressure_without_completed_groups`.
- It fixes a limitation in place. If group values ever release capacity after `EmitTo::First`, this test starts failing even though the behavior got better.
The capacity observation could go in a follow-up issue instead.
##########
datafusion/physical-plan/src/aggregates/clustered_single_stream.rs:
##########
@@ -329,30 +330,35 @@ impl ClusteredSingleAggregateStream {
spill_context.as_deref(),
));
timer.done();
- match resize_result {
- Ok(()) => {}
+ // Check the reservation before taking completed groups, as the partial
+ // stream does, so spilling is unchanged. Without spilling, the refusal
+ // is returned below only if no group is complete.
+ let oom = match resize_result {
+ Ok(()) => None,
Err(e @ DataFusionError::ResourcesExhausted(_)) => {
- let Some(spill_context) = spill_context else {
- // `None` means spilling is not supported, see comments
- // at `ClusteredSingleAggregateState` for details.
- return Self::break_with_err(e);
- };
- if table.is_empty() {
- return Self::break_with_internal_err(
- "Clustered single aggregate ran out of memory with no aggregated groups",
- );
+ match spill_context.take() {
Review Comment:
The `Some` arm always returns, so matching on `spill_context` by value compiles without `take()`. That also lets the `mut` on `spill_context` at line 298 go away.
```rust
Err(e @ DataFusionError::ResourcesExhausted(_)) => match spill_context {
Some(spill_context) => { /* unchanged, returns Spilling */ }
// Spilling is not supported, see comments at
// `ClusteredSingleAggregateState` for details.
None => Some(e),
},
```
##########
datafusion/physical-plan/src/aggregates/clustered_test_util.rs:
##########
@@ -0,0 +1,210 @@
+// 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.
+
+//! Test helpers for the clustered aggregate streams under memory pressure.
+
+use std::sync::Arc;
+
+use arrow::array::Int64Array;
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::test_util::batches_to_string;
+use datafusion_common::{DataFusionError, Result};
+use datafusion_execution::TaskContext;
+use datafusion_execution::config::SessionConfig;
+use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
+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;
+use futures::StreamExt;
+
+use super::{AggregateExec, AggregateMode, GroupClusteringMode, PhysicalGroupBy};
+use crate::ExecutionPlan;
+use crate::test::TestMemoryExec;
+
+/// Key columns of two batches sorted on `k` that share the boundary key 256, so
+/// groups complete within each batch under full group clustering.
+pub(super) fn full_batches() -> Vec<Vec<Vec<i64>>> {
+ vec![vec![(0..=256).collect()], vec![(256..=512).collect()]]
+}
+
+/// Key columns `(a, b)` of batches sorted on `a` only, with 256 groups for each
+/// `a` spread over two batches. The groups of one `a` complete together when the
+/// next `a` starts.
+pub(super) fn partial_batches() -> Vec<Vec<Vec<i64>>> {
+ [0, 0, 1, 1, 2, 2]
+ .into_iter()
+ .map(|a| vec![vec![a; 256], (0..256).collect()])
+ .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<Vec<Vec<i64>>> {
+ vec![vec![vec![0; 512]]; 2]
+}
+
+/// Builds `SUM(v) GROUP BY <keys>` in `mode` over `batches` of `(keys..., v)`
+/// columns, where `v` equals the last key. The input is declared sorted on the
+/// first key only, so one key gives full group clustering and more keys give
+/// partial group clustering.
+///
+/// For `Final` mode, the batches are the partial states, since the state of
+/// `SUM(v)` is a single `Int64` column.
+pub(super) fn sum_over_sorted_input(
Review Comment:
`ExecutionPlan::execute` takes `&self`, so this could return `AggregateExec` and the helpers could take `&AggregateExec`. Callers would then check `plan.group_clustering_mode()` directly. That removes the tuple, the `group_clustering_mode().clone()`, and the `(plan, _)` discards.
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -336,12 +338,19 @@ impl Aggregating {
let timer = elapsed_compute.timer();
self.table.aggregate_batch(&batch)?;
- match reservation.try_resize(self.reservation_size()) {
- Ok(()) => {}
- Err(oom @ DataFusionError::ResourcesExhausted(_)) => {
- let Some(spill_context) = self.spill_context.as_mut() else {
- return Err(oom);
- };
+ // Check the reservation before taking completed groups, as the partial
+ // stream does, so spilling is unchanged. Without spilling, the refusal
+ // is returned below only if no group is complete.
+ let oom = match (
Review Comment:
The tuple match repeats the `Err(oom @ DataFusionError::ResourcesExhausted(_))` pattern and calls `as_mut()` on the success path too. A nested match keeps the original shape and reads the same as the single stream:
```rust
let oom = match reservation.try_resize(self.reservation_size()) {
Ok(()) => None,
Err(oom @ DataFusionError::ResourcesExhausted(_)) => {
match self.spill_context.as_mut() {
None => Some(oom),
Some(spill_context) => { /* unchanged spill, then continue */ }
}
}
Err(e) => return Err(e),
};
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4366065885
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -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
Review Comment:
Added `partially_clustered_memory_pressure_without_spilling` for both streams: a two-column key `(a, b)` sorted on `a`, with the disk manager disabled, asserting `GroupClusteringMode::Partial([0])`. While writing it I found that under partial clustering the group values keep their capacity after `EmitTo::First`. `memory_size()` stays the same after the emit (31192 bytes before and after for 512 → 256 groups in this test), so emitting can't bring the table under a limit it already exceeded. The test therefore pins that the branch still returns `ResourcesExhausted` cleanly one byte below the unlimited footprint, rather than claiming success. I checked that it actually reaches the new branch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#issuecomment-6380136045
> Thanks for the thorough review! I rebased on main, which fixes the `cargo-semver-checks` failure (the old base didn't have `input_keeps_groups_local`), and pushed a commit addressing the inline comments:
>
> - **Fall-through:** both streams now keep the refused reservation and fall through to the normal emit path. The duplicated tail and the shadowed `e` are gone.
> - **Shared helper:** the driver now lives in a `#[cfg(test)] clustered_test_util` module that all three clustered streams use (`clustered_partial_stream.rs` too). The tests have consistent names in each module: `fully_clustered_completed_groups_under_memory_pressure`, `memory_pressure_without_completed_groups`, `partially_clustered_memory_pressure_without_spilling`.
> - **Comments:** I trimmed the stage-list addition. What remains is the type doc and one short inline note per stream.
>
> I didn't do the randomized/fuzz variant in this PR, to keep it focused. I'm happy to send it as a follow-up that extends the aggregate fuzz cases with multi-batch, multi-column and string keys, if that works for you.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#issuecomment-6419126273
Thanks for the suggestion @jayzhan211! I updated the PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#issuecomment-6368392623
Thanks @yashaggarwal85d I'll check this today, also @jayzhan211 FYI as you also reviewed https://github.com/apache/datafusion/pull/26576
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4379368914
##########
datafusion/physical-plan/src/aggregates/clustered_single_stream.rs:
##########
@@ -329,30 +330,35 @@ impl ClusteredSingleAggregateStream {
spill_context.as_deref(),
));
timer.done();
- match resize_result {
- Ok(()) => {}
+ // Check the reservation before taking completed groups, as the partial
+ // stream does, so spilling is unchanged. Without spilling, the refusal
+ // is returned below only if no group is complete.
+ let oom = match resize_result {
+ Ok(()) => None,
Err(e @ DataFusionError::ResourcesExhausted(_)) => {
- let Some(spill_context) = spill_context else {
- // `None` means spilling is not supported, see comments
- // at `ClusteredSingleAggregateState` for details.
- return Self::break_with_err(e);
- };
- if table.is_empty() {
- return Self::break_with_internal_err(
- "Clustered single aggregate ran out of memory with no aggregated groups",
- );
+ match spill_context.take() {
Review Comment:
Nice, that's cleaner. It now matches on `spill_context` by value, and the `mut` is gone (baa8fe5).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#issuecomment-6358317296
<!-- semver-check-comment -->
Thank you for opening this pull request!
Reviewer note: [cargo-semver-checks](https://github.com/obi1kenobi/cargo-semver-checks) reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).
<details>
<summary>Details</summary>
```
Cloning apache/main
Building datafusion-physical-plan v55.1.0 (current)
Built [ 38.653s] (current)
Parsing datafusion-physical-plan v55.1.0 (current)
Parsed [ 0.155s] (current)
Building datafusion-physical-plan v55.1.0 (baseline)
Built [ 26.896s] (baseline)
Parsing datafusion-physical-plan v55.1.0 (baseline)
Parsed [ 0.160s] (baseline)
Checking datafusion-physical-plan v55.1.0 -> v55.1.0 (no change; assume patch)
Checked [ 0.481s] 229 checks: 228 pass, 1 fail, 0 warn, 31 skip
--- failure inherent_method_missing: pub method removed or renamed ---
Description:
A publicly-visible method or associated fn is no longer available under its prior name. It may have been renamed or removed entirely.
ref: https://doc.rust-lang.org/cargo/reference/semver.html#item-remove
impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.51.0/src/lints/inherent_method_missing.ron
Failed in:
AggregateExec::input_keeps_groups_local, previously in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/c91879df6bf841ecef350e933fd580c2a37d4b29/datafusion/physical-plan/src/aggregates/mod.rs:1663
Summary semver requires new major version: 1 major and 0 minor checks failed
Finished [ 68.690s] datafusion-physical-plan
```
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4366063623
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -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
Review Comment:
It's intentional, to keep this PR minimal. When a spill context exists, the same refused reservation decides whether to spill, so checking the reservation first leaves that path and the success path exactly as they are. It also matches what `ClusteredPartialAggregateStream` does after #26576. I added a one-line comment saying this. I agree that "take completed groups first, then size once" is cleaner, and as you say neither lowers the real peak. I'd rather change that ordering across all three streams in a separate PR, if you think it's worth doing.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#issuecomment-6420133028
Thanks @comphead, I'll read through your suggestions and make the necessary changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4366069067
##########
datafusion/physical-plan/src/aggregates/clustered_single_stream.rs:
##########
@@ -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(
Review Comment:
Agreed. I moved it into one shared `#[cfg(test)] mod clustered_test_util` next to the stream modules. It provides `sum_over_sorted_input(mode, keys, batches)`, which builds the plan (for `Final`, it builds the partial-state input), and `collect_with_held_memory(&Arc<dyn ExecutionPlan>, pool)`, which runs it through `execute` with `enable_migration_aggregate` on and temp files off. I put it in its own file instead of inside `aggregates/mod.rs` because that file is already very large; happy to move it if you prefer. `clustered_partial_stream.rs` now uses it too. I didn't add the SLT case: as you noted, it would depend on exact byte sizes and could be fragile.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4379376853
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -732,4 +749,34 @@ mod tests {
assert_eq!(pool.reserved(), 0);
Ok(())
}
+
+ #[tokio::test]
+ async fn fully_clustered_completed_groups_under_memory_pressure() -> Result<()> {
+ let (plan, mode) =
+ sum_over_sorted_input(AggregateMode::Final, &["k"], full_batches())?;
+ assert_eq!(mode, GroupClusteringMode::Full);
+ assert_completed_groups_emitted_under_limit(&plan).await
+ }
+
+ /// A single key never completes a group before the input ends. The limit
+ /// fits only the initial table.
+ #[tokio::test]
+ async fn memory_pressure_without_completed_groups() -> Result<()> {
+ let (plan, _) =
+ sum_over_sorted_input(AggregateMode::Final, &["k"], single_key_batches())?;
+ let limit = held_memory(&plan).await?;
+ assert_resources_exhausted_under_limit(&plan, limit).await
+ }
+
+ /// Without temporary files, partial group clustering takes the same path as
+ /// full clustering. Its group values keep their capacity after emitting, so
+ /// one byte below the unlimited footprint still fails.
+ #[tokio::test]
+ async fn partially_clustered_memory_pressure_without_spilling() -> Result<()> {
Review Comment:
Fair point. It passes on main too, so it doesn't really test the fix. I removed both copies and `partial_batches` (c3ea1ce). I'll open a separate issue for the retained capacity.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4379401480
##########
datafusion/physical-plan/src/aggregates/clustered_test_util.rs:
##########
@@ -0,0 +1,210 @@
+// 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.
+
+//! Test helpers for the clustered aggregate streams under memory pressure.
+
+use std::sync::Arc;
+
+use arrow::array::Int64Array;
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::test_util::batches_to_string;
+use datafusion_common::{DataFusionError, Result};
+use datafusion_execution::TaskContext;
+use datafusion_execution::config::SessionConfig;
+use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
+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;
+use futures::StreamExt;
+
+use super::{AggregateExec, AggregateMode, GroupClusteringMode, PhysicalGroupBy};
+use crate::ExecutionPlan;
+use crate::test::TestMemoryExec;
+
+/// Key columns of two batches sorted on `k` that share the boundary key 256, so
+/// groups complete within each batch under full group clustering.
+pub(super) fn full_batches() -> Vec<Vec<Vec<i64>>> {
+ vec![vec![(0..=256).collect()], vec![(256..=512).collect()]]
+}
+
+/// Key columns `(a, b)` of batches sorted on `a` only, with 256 groups for each
+/// `a` spread over two batches. The groups of one `a` complete together when the
+/// next `a` starts.
+pub(super) fn partial_batches() -> Vec<Vec<Vec<i64>>> {
+ [0, 0, 1, 1, 2, 2]
+ .into_iter()
+ .map(|a| vec![vec![a; 256], (0..256).collect()])
+ .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<Vec<Vec<i64>>> {
+ vec![vec![vec![0; 512]]; 2]
+}
+
+/// Builds `SUM(v) GROUP BY <keys>` in `mode` over `batches` of `(keys..., v)`
+/// columns, where `v` equals the last key. The input is declared sorted on the
+/// first key only, so one key gives full group clustering and more keys give
+/// partial group clustering.
+///
+/// For `Final` mode, the batches are the partial states, since the state of
+/// `SUM(v)` is a single `Int64` column.
+pub(super) fn sum_over_sorted_input(
Review Comment:
Good spot. It now builds the key arrays once and pushes an `Arc::clone` of the last one for `v` (c3ea1ce).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4366061051
##########
datafusion/physical-plan/src/aggregates/clustered_single_stream.rs:
##########
@@ -335,7 +336,27 @@ impl ClusteredSingleAggregateStream {
let Some(spill_context) = spill_context else {
Review Comment:
Done. The refused error is now kept as `oom: Option<DataFusionError>`, and both streams fall through to the existing `take_completed_result_batch` / emit path. `oom` is returned only when nothing is complete, and the second `e` no longer shadows it. `ClusteredFinalAggregateStream::Aggregating::handle_stage` got the same change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4379363880
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -336,12 +338,19 @@ impl Aggregating {
let timer = elapsed_compute.timer();
self.table.aggregate_batch(&batch)?;
- match reservation.try_resize(self.reservation_size()) {
- Ok(()) => {}
- Err(oom @ DataFusionError::ResourcesExhausted(_)) => {
- let Some(spill_context) = self.spill_context.as_mut() else {
- return Err(oom);
- };
+ // Check the reservation before taking completed groups, as the partial
+ // stream does, so spilling is unchanged. Without spilling, the refusal
+ // is returned below only if no group is complete.
+ let oom = match (
Review Comment:
Makes sense. Switched it back to a nested match so the success path stays as it was (baa8fe5).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "yashaggarwal85d (via GitHub)" <gi...@apache.org>.
yashaggarwal85d commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4366068022
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -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
+ ));
+ 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<()> {
Review Comment:
Added `memory_pressure_without_completed_groups` for both streams. A single key never completes a group before EOF, and the limit fits only the initial table, so the refused reservation must come back as `ResourcesExhausted`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: emit completed groups when fully sorted final and single aggrega… [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4361499205
##########
datafusion/physical-plan/src/aggregates/clustered_single_stream.rs:
##########
@@ -335,7 +336,27 @@ impl ClusteredSingleAggregateStream {
let Some(spill_context) = spill_context else {
Review Comment:
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.
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -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
+ ));
+ 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<()> {
Review Comment:
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.
##########
datafusion/physical-plan/src/aggregates/clustered_single_stream.rs:
##########
@@ -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(
Review Comment:
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.
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -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
Review Comment:
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.
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -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
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org