comphead opened a new issue, #26628: URL: https://github.com/apache/datafusion/issues/26628
### Describe the bug `Final` and `Single` aggregates over input sorted on all grouping keys (`ClusteredFinalAggregateStream`, `ClusteredSingleAggregateStream`) fail with `ResourcesExhausted` instead of emitting the completed groups. This is the bug from #26575 in the two sibling streams. After each input batch, both streams reserve memory for the whole table before they emit the completed groups. Under full group clustering neither stream can spill, so a refused reservation returns the error: https://github.com/apache/datafusion/blob/5026ac452492048a3a9fcbe2e653ecfa921cc709/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs#L342-L344 https://github.com/apache/datafusion/blob/5026ac452492048a3a9fcbe2e653ecfa921cc709/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs#L335-L339 The memory is already allocated, and emitting the completed groups would shrink the table to the one active group. The legacy `GroupedHashAggregateStream`, still used with `enable_migration_aggregate = false`, emits completed groups before it updates the reservation, so it does not fail: https://github.com/apache/datafusion/blob/5026ac452492048a3a9fcbe2e653ecfa921cc709/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs#L704-L716 ### To Reproduce `SUM(k) GROUP BY k` over two batches sorted on `k`, with keys `0..=256` and `256..=512`, so key 256 spans both batches. For `Final` the input is partial states `(k, sum)` with the same keys. Each test runs the stream with an unbounded pool and records the most memory it holds between output batches, 12,992 bytes, which is the table with only the active group. It then reruns with a `GreedyMemoryPool` of that size. On `main` at 5026ac452 both reruns fail: ``` ResourcesExhausted("Failed to allocate additional 13.3 KB for ClusteredFinalAggregateStream[0] with 5.4 KB already allocated for this reservation - 7.3 KB remain available for the total memory pool: greedy(used: 5.4 KB, pool_size: 12.7 KB)") ResourcesExhausted("Failed to allocate additional 13.3 KB for ClusteredSingleAggregateStream[0] with 5.4 KB already allocated for this reservation - 7.3 KB remain available for the total memory pool: greedy(used: 5.4 KB, pool_size: 12.7 KB)") ``` Under the same pool, the legacy stream completes both plans and returns the same rows as the unbounded run. So does each new stream once it emits the completed groups on a refused reservation. <details> <summary>Reproducer tests</summary> In the `tests` module of `clustered_single_stream.rs` (also needs `batches_to_string` and `UnboundedMemoryPool`): ```rust /// Runs `SELECT k, 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_fully_clustered( 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(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 = 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)) } /// With full group clustering, 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_clustered_completed_groups_under_memory_pressure() -> Result<()> { let (expected, held) = run_fully_clustered(Arc::new(UnboundedMemoryPool::default())).await?; let (actual, _) = run_fully_clustered(Arc::new(GreedyMemoryPool::new(held))).await?; assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); Ok(()) } ``` In the `tests` module of `clustered_final_stream.rs` (also needs `batches_to_string` and `UnboundedMemoryPool`): ```rust /// Runs a final `SUM(k) GROUP BY k` over two batches of partial states sorted /// on `k` that share a boundary key. Returns the output and the most memory /// the stream held between batches. async fn run_fully_clustered( 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_schema = AggregateExec::try_new( AggregateMode::Partial, groups.clone(), expressions.clone(), vec![None], empty, Arc::clone(&schema), )? .schema(); // Partial states `(k, sum)` with `sum = k`. 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(input), schema, )?; 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)) } /// With full group clustering, 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_clustered_completed_groups_under_memory_pressure() -> Result<()> { let (expected, held) = run_fully_clustered(Arc::new(UnboundedMemoryPool::default())).await?; let (actual, _) = run_fully_clustered(Arc::new(GreedyMemoryPool::new(held))).await?; assert_eq!(batches_to_string(&actual), batches_to_string(&expected)); Ok(()) } ``` </details> ### Expected behavior On a refused reservation without a spill context, emit `take_completed_result_batch()` and keep only the active group, as #26576 does for the partial stream. Return the error only when no group is complete. Both streams already reserve just the retained table when they cannot hold the output batch, and the emitted rows are the ones the success path emits, so results and order stay the same. ### Additional context - The same bug in the partial stream is #26575, with a fix in #26576. - `branch-55` has the same early return in `OrderedFinalAggregateStream`. Sorted `Single` aggregates there still use `GroupedHashAggregateStream`, so only `Final` is affected on 55. This is from reading the code, not tested on 55. https://github.com/apache/datafusion/blob/6fceca2806dfc8f5ec73f17fc315ab2972efa428/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs#L459-L465 - With partial group clustering and temporary files disabled, `spill_context` is also `None`, so those queries reach the same early return. Not tested. -- 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]
