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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
