jayzhan211 opened a new issue, #26701: URL: https://github.com/apache/datafusion/issues/26701
### Describe the bug `Final` and `FinalPartitioned` aggregates (`FinalHashAggregateStream`) fail with `ResourcesExhausted` when the pool refuses the reservation for their output, although the output is already materialized and the table it came from is gone. In the same situation, the other new aggregate streams and the legacy `GroupedHashAggregateStream` carry on. When the stream starts producing output, the table materializes all groups into one batch. On the first output batch, the stream resizes its reservation from the table to that batch. If the batch is larger than what the table had reserved, the reservation has to grow, and a refusal is returned to the query: https://github.com/apache/datafusion/blob/836b4dda4fb265c07c7060b1f3d341cd209ea6a4/datafusion/physical-plan/src/aggregates/hash_stream.rs#L838-L841 `string_agg` produces such an output. In the reproducer below, the table reserves 435.0 KB and the materialized output needs 543.8 KB. I also ran `sum`, `count`, `max`, `first_value`, `array_agg` and `approx_distinct` with Int64, Utf8, Utf8View and two-column keys: only `string_agg` grew the reservation at this point (by 16% to 27%), the others shrank it. The new streams are the default since 55.0.0 (#22899), so any `Final` or `FinalPartitioned` hash aggregate whose output is larger than its table can fail this way under memory pressure, unless `datafusion.execution.enable_migration_aggregate` is set to `false`. ### To Reproduce A SQL memory limit cannot target one stream's output, so the reproducer is a test. It runs `SELECT k, string_agg(s, ',') FROM t GROUP BY k` over 20,000 rows and 2,000 groups, first with no limit to measure how much the `Final` stream reserves for its output, then under a `GreedyMemoryPool` one byte smaller, so the table still fits. It checks the legacy stream and `Single` under the same limit first; on `main` at 836b4dda4f only the last step, the new `Final` stream, fails. Save it as `datafusion/physical-plan/tests/final_output_refusal.rs` and run `cargo test -p datafusion-physical-plan --test final_output_refusal`. <details> <summary>Reproducer test</summary> ```rust use std::sync::Arc; use arrow::array::{Int64Array, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_common::Result; use datafusion_common::test_util::batches_to_sort_string; use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::string_agg::string_agg_udaf; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::{col, lit}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; use datafusion_physical_plan::common::collect; use datafusion_physical_plan::test::TestMemoryExec; use futures::StreamExt; fn context(pool: Arc<dyn MemoryPool>, migration: bool) -> Result<Arc<TaskContext>> { let config = SessionConfig::new() .with_batch_size(512) .set_bool("datafusion.execution.enable_migration_aggregate", migration); let runtime = RuntimeEnvBuilder::new() .with_memory_pool(pool) .build_arc()?; Ok(Arc::new( TaskContext::default() .with_session_config(config) .with_runtime(runtime), )) } /// `SELECT k, string_agg(s, ',') FROM t GROUP BY k` in `mode`. fn aggregate( mode: AggregateMode, input: Arc<dyn ExecutionPlan>, raw: &SchemaRef, ) -> Result<Arc<AggregateExec>> { let group_by = PhysicalGroupBy::new_single(vec![(col("k", raw)?, "k".into())]); let group_by = if mode == AggregateMode::Final { group_by.as_final() } else { group_by }; let agg = AggregateExprBuilder::new(string_agg_udaf(), vec![col("s", raw)?, lit(",")]) .schema(Arc::clone(raw)) .alias("agg") .build()?; Ok(Arc::new(AggregateExec::try_new( mode, group_by, vec![Arc::new(agg)], vec![None], input, Arc::clone(raw), )?)) } /// Returns the output and the memory reserved when the first output batch /// arrives, which for the new streams is the whole materialized output. async fn run( plan: Arc<AggregateExec>, pool: Arc<dyn MemoryPool>, migration: bool, ) -> Result<(Vec<RecordBatch>, usize)> { let mut stream = plan.execute(0, context(Arc::clone(&pool), migration)?)?; let mut held = None; let mut output = vec![]; while let Some(batch) = stream.next().await.transpose()? { held.get_or_insert(pool.reserved()); output.push(batch); } Ok((output, held.unwrap())) } #[tokio::test] async fn final_output_refusal() -> Result<()> { // 20,000 rows in 8 batches, 2,000 groups. let raw = Arc::new(Schema::new(vec![ Field::new("k", DataType::Int64, false), Field::new("s", DataType::Utf8, false), ])); let rows = RecordBatch::try_new( Arc::clone(&raw), vec![ Arc::new(Int64Array::from_iter_values((0..20_000).map(|i| i % 2_000))), Arc::new(StringArray::from_iter_values( (0..20_000).map(|i| format!("value-{i:08}")), )), ], )?; let batches: Vec<_> = (0..8).map(|i| rows.slice(i * 2_500, 2_500)).collect(); let raw_input = || TestMemoryExec::try_new_exec(&[batches.clone()], Arc::clone(&raw), None); // The Final input is the Partial output, computed once without a limit. let partial = aggregate(AggregateMode::Partial, raw_input()?, &raw)?; let unbounded = || Arc::new(UnboundedMemoryPool::default()); let states = collect(partial.execute(0, context(unbounded(), true)?)?).await?; let state_input = || TestMemoryExec::try_new_exec(&[states.clone()], partial.schema(), None); // Without a limit, Final reserves its whole materialized output on the // first output batch. Rerun with one byte less: the table still fits. let final_plan = || aggregate(AggregateMode::Final, state_input()?, &raw); let (expected, held) = run(final_plan()?, unbounded(), true).await?; let limited = || Arc::new(GreedyMemoryPool::new(held - 1)); // The legacy stream (`enable_migration_aggregate = false`) completes. let (legacy, _) = run(final_plan()?, limited(), false).await?; assert_eq!( batches_to_sort_string(&legacy), batches_to_sort_string(&expected) ); // Single over the raw rows completes. let single_plan = aggregate(AggregateMode::Single, raw_input()?, &raw)?; let (single, _) = run(single_plan, limited(), true).await?; assert_eq!( batches_to_sort_string(&single), batches_to_sort_string(&expected) ); // Final fails with ResourcesExhausted. let (actual, _) = run(final_plan()?, limited(), true).await?; assert_eq!( batches_to_sort_string(&actual), batches_to_sort_string(&expected) ); Ok(()) } ``` </details> ### Expected behavior | Plan, under the same limit | DataFusion returns | Expected | |---|---|---| | `Final`, new stream (the default) | `ResourcesExhausted("Failed to allocate additional 108.9 KB for FinalHashAggregateStream[0] with 435.0 KB already allocated for this reservation - 108.8 KB remain available for the total memory pool: greedy(used: 435.0 KB, pool_size: 543.8 KB)")` | 2,000 rows, the same as with no limit | | `Final`, legacy stream (`enable_migration_aggregate = false`) | 2,000 rows, the same as with no limit | the same | | `Single` over the raw rows, new stream | 2,000 rows, the same as with no limit | the same | ### Additional context - The cause is the `?` on the reservation in the output loop of `FinalHashAggregateStream` ([`hash_stream.rs:841`](https://github.com/apache/datafusion/blob/836b4dda4fb265c07c7060b1f3d341cd209ea6a4/datafusion/physical-plan/src/aggregates/hash_stream.rs#L841)). Once the table is materialized, its `memory_size()` is the size of the output batch ([`aggregate_hash_table/common.rs:356-357`](https://github.com/apache/datafusion/blob/836b4dda4fb265c07c7060b1f3d341cd209ea6a4/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs#L356-L357)). So the first output batch asks for the whole output, and nothing is left to spill or release. - The other streams don't fail at this point. `SingleHashAggregateStream` hands the output off unreserved ([`single_stream.rs:450-460`](https://github.com/apache/datafusion/blob/836b4dda4fb265c07c7060b1f3d341cd209ea6a4/datafusion/physical-plan/src/aggregates/single_stream.rs#L450-L460)), `PartialHashAggregateStream` ignores the refusal ([`hash_stream.rs:499-502`](https://github.com/apache/datafusion/blob/836b4dda4fb265c07c7060b1f3d341cd209ea6a4/datafusion/physical-plan/src/aggregates/hash_stream.rs#L499-L502)), and so does the legacy stream ([`grouped_hash_stream.rs:1094-1096`](https://github.com/apache/datafusion/blob/836b4dda4fb265c07c7060b1f3d341cd209ea6a4/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs#L1094-L1096)). - `Final` and `FinalPartitioned` both use this stream ([`mod.rs:1789`](https://github.com/apache/datafusion/blob/836b4dda4fb265c07c7060b1f3d341cd209ea6a4/datafusion/physical-plan/src/aggregates/mod.rs#L1789)). - Related: #26628 (ClusteredFinal and ClusteredSingle fail on a refused reservation after the memory is already allocated) and #26575, fixed by #26576 (the same in ClusteredPartial). The open PR #26562 adds a test, `output_materialization_refused`, that pins how each stream reacts to this refusal today. A fix will change its `Final` snapshot. - A fix direction: handle the refusal the way `Single` does, by releasing the reservation and handing the already materialized output off unreserved. With that change, the reproducer passes. ### Willingness to contribute I'm not planning to fix this myself, so anyone is welcome to pick it up. Feel free to ping me with questions about the cause or when you'd like a review. -- 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]
