From d7a29d9c67c9ef9ae8f258b1fb1440df11274ea9 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Tue, 28 Jul 2026 18:34:41 +0300 Subject: [PATCH 1/2] test: cover partially ordered aggregate spilling --- .../physical-plan/src/aggregates/mod.rs | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 787cd4f03ff6c..12974aa0aa394 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4401,6 +4401,124 @@ mod tests { Ok(()) } + /// Single-mode grouped aggregation over *partially sorted* input + /// (`GroupOrdering::Partial` + `OutOfMemoryMode::Spill` in + /// `GroupedHashAggregateStream`) must spill and produce the same results + /// as an unlimited-memory run. + #[tokio::test] + async fn partially_sorted_single_aggregate_spills_and_matches() -> Result<()> { + // Grouping by (sort_col, group_col) while the input carries an ordering + // on `sort_col` alone produces `InputOrderMode::PartiallySorted`, which + // drives `GroupOrdering::Partial` inside `GroupedHashAggregateStream`. + // A single constant `sort_col` value means no sort boundary ever + // completes mid-stream, so groups accumulate until the pool is + // exhausted and the stream must spill sorted intermediate state and + // merge it back. + let schema = Arc::new(Schema::new(vec![ + Field::new("sort_col", DataType::Int32, false), + Field::new("group_col", DataType::UInt32, false), + Field::new("value_col", DataType::Float64, false), + ])); + + let batch = |groups: Vec, values: Vec| -> Result { + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![0; groups.len()])), + Arc::new(UInt32Array::from(groups)), + Arc::new(Float64Array::from(values)), + ], + )?) + }; + // Every group appears in both batches with different values, so the + // first batch's states get spilled and the read-back merge must + // combine two partial states per group instead of passing spilled + // singleton states through. + let n_groups: u32 = 400; + let groups: Vec = (0..n_groups).collect(); + let batches = vec![ + batch(groups.clone(), (0..n_groups).map(f64::from).collect())?, + batch(groups, (0..n_groups).map(|g| f64::from(g + 1000)).collect())?, + ]; + + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("sort_col", 0), + ))]) + .unwrap(); + + let build_exec = || -> Result> { + let input = TestMemoryExec::try_new( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + )? + .try_with_sort_information(vec![ordering.clone()])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + // MIN keeps a single intermediate value; AVG keeps two (sum + count), + // so both single- and multi-state accumulators are spilled and merged. + let aggregates: Vec> = vec![ + Arc::new( + AggregateExprBuilder::new( + min_udaf(), + vec![col("value_col", &schema)?], + ) + .schema(Arc::clone(&schema)) + .alias("MIN(value_col)") + .build()?, + ), + Arc::new( + AggregateExprBuilder::new( + avg_udaf(), + vec![col("value_col", &schema)?], + ) + .schema(Arc::clone(&schema)) + .alias("AVG(value_col)") + .build()?, + ), + ]; + Ok(Arc::new(AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![ + (col("sort_col", &schema)?, "sort_col".to_string()), + (col("group_col", &schema)?, "group_col".to_string()), + ]), + aggregates, + vec![None, None], + input, + Arc::clone(&schema), + )?)) + }; + + // Sanity: partial ordering must be detected, otherwise the test would + // exercise `GroupOrdering::None` instead of the intended path. + assert!(matches!( + build_exec()?.input_order_mode(), + InputOrderMode::PartiallySorted(_) + )); + + // Unlimited-memory reference result (no spill). + let reference_agg = build_exec()?; + let reference = + collect(reference_agg.execute(0, new_spill_ctx(8, 1024 * 1024))?).await?; + assert_spill_count_metric(false, reference_agg); + + // A small `FairSpillPool` (disk enabled by default) forces the + // partially-ordered aggregate to spill sorted intermediate state and + // merge it back. The pool is large enough to still reserve the sort + // headroom required while spilling, but too small to hold every group. + let spill_agg = build_exec()?; + let spilled = collect(spill_agg.execute(0, new_spill_ctx(8, 16_000))?).await?; + assert_spill_count_metric(true, spill_agg); + + // Spilled output must match the unlimited-memory reference. + assert_eq!( + batches_to_sort_string(&reference), + batches_to_sort_string(&spilled), + ); + + Ok(()) + } + #[tokio::test] async fn ordered_partial_aggregate_partially_sorted_no_emit_panic() -> Result<()> { // Reproducer for #20445: emitting from PartiallySorted input must not From 5f5470e44f1d8e7b4fabe42dd71a836721e1bf4d Mon Sep 17 00:00:00 2001 From: buraksenn Date: Wed, 29 Jul 2026 09:21:44 +0300 Subject: [PATCH 2/2] refactor to slt test instead of rust --- .../physical-plan/src/aggregates/mod.rs | 118 ------------------ .../test_files/ordered_aggregate_spill.slt | 52 ++++++++ 2 files changed, 52 insertions(+), 118 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 12974aa0aa394..787cd4f03ff6c 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4401,124 +4401,6 @@ mod tests { Ok(()) } - /// Single-mode grouped aggregation over *partially sorted* input - /// (`GroupOrdering::Partial` + `OutOfMemoryMode::Spill` in - /// `GroupedHashAggregateStream`) must spill and produce the same results - /// as an unlimited-memory run. - #[tokio::test] - async fn partially_sorted_single_aggregate_spills_and_matches() -> Result<()> { - // Grouping by (sort_col, group_col) while the input carries an ordering - // on `sort_col` alone produces `InputOrderMode::PartiallySorted`, which - // drives `GroupOrdering::Partial` inside `GroupedHashAggregateStream`. - // A single constant `sort_col` value means no sort boundary ever - // completes mid-stream, so groups accumulate until the pool is - // exhausted and the stream must spill sorted intermediate state and - // merge it back. - let schema = Arc::new(Schema::new(vec![ - Field::new("sort_col", DataType::Int32, false), - Field::new("group_col", DataType::UInt32, false), - Field::new("value_col", DataType::Float64, false), - ])); - - let batch = |groups: Vec, values: Vec| -> Result { - Ok(RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(vec![0; groups.len()])), - Arc::new(UInt32Array::from(groups)), - Arc::new(Float64Array::from(values)), - ], - )?) - }; - // Every group appears in both batches with different values, so the - // first batch's states get spilled and the read-back merge must - // combine two partial states per group instead of passing spilled - // singleton states through. - let n_groups: u32 = 400; - let groups: Vec = (0..n_groups).collect(); - let batches = vec![ - batch(groups.clone(), (0..n_groups).map(f64::from).collect())?, - batch(groups, (0..n_groups).map(|g| f64::from(g + 1000)).collect())?, - ]; - - let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( - Column::new("sort_col", 0), - ))]) - .unwrap(); - - let build_exec = || -> Result> { - let input = TestMemoryExec::try_new( - std::slice::from_ref(&batches), - Arc::clone(&schema), - None, - )? - .try_with_sort_information(vec![ordering.clone()])?; - let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); - // MIN keeps a single intermediate value; AVG keeps two (sum + count), - // so both single- and multi-state accumulators are spilled and merged. - let aggregates: Vec> = vec![ - Arc::new( - AggregateExprBuilder::new( - min_udaf(), - vec![col("value_col", &schema)?], - ) - .schema(Arc::clone(&schema)) - .alias("MIN(value_col)") - .build()?, - ), - Arc::new( - AggregateExprBuilder::new( - avg_udaf(), - vec![col("value_col", &schema)?], - ) - .schema(Arc::clone(&schema)) - .alias("AVG(value_col)") - .build()?, - ), - ]; - Ok(Arc::new(AggregateExec::try_new( - AggregateMode::Single, - PhysicalGroupBy::new_single(vec![ - (col("sort_col", &schema)?, "sort_col".to_string()), - (col("group_col", &schema)?, "group_col".to_string()), - ]), - aggregates, - vec![None, None], - input, - Arc::clone(&schema), - )?)) - }; - - // Sanity: partial ordering must be detected, otherwise the test would - // exercise `GroupOrdering::None` instead of the intended path. - assert!(matches!( - build_exec()?.input_order_mode(), - InputOrderMode::PartiallySorted(_) - )); - - // Unlimited-memory reference result (no spill). - let reference_agg = build_exec()?; - let reference = - collect(reference_agg.execute(0, new_spill_ctx(8, 1024 * 1024))?).await?; - assert_spill_count_metric(false, reference_agg); - - // A small `FairSpillPool` (disk enabled by default) forces the - // partially-ordered aggregate to spill sorted intermediate state and - // merge it back. The pool is large enough to still reserve the sort - // headroom required while spilling, but too small to hold every group. - let spill_agg = build_exec()?; - let spilled = collect(spill_agg.execute(0, new_spill_ctx(8, 16_000))?).await?; - assert_spill_count_metric(true, spill_agg); - - // Spilled output must match the unlimited-memory reference. - assert_eq!( - batches_to_sort_string(&reference), - batches_to_sort_string(&spilled), - ); - - Ok(()) - } - #[tokio::test] async fn ordered_partial_aggregate_partially_sorted_no_emit_panic() -> Result<()> { // Reproducer for #20445: emitting from PartiallySorted input must not diff --git a/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt index a4d492ab82e12..2c53c94144eb3 100644 --- a/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt +++ b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt @@ -176,6 +176,58 @@ Plan with Metrics 03)----AggregateExec: mode=Partial,aggr=[sum(t1.v1 * Int64(2)), min(t1.v1 % Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] +# ================================================================================== +# Single mode: with one partition the whole aggregation runs in a `Single` mode +# AggregateExec. min() keeps one intermediate state and avg() keeps two (sum + +# count), so both single- and multi-state accumulators are spilled and merged. +# ================================================================================== + +statement ok +SET datafusion.execution.target_partitions = 1 + +# Reference round: enough memory to aggregate without spilling. +statement ok +SET datafusion.runtime.memory_limit = '10M' + +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, min(v1 * 2), avg(v1) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=Single,aggr=[min(t1.v1 * Int64(2)), avg(t1.v1)], ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +query IIIR rowsort +SELECT round(v1, -4), v1 % 5000, min(v1 * 2), avg(v1) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +60000 values hashing to 872df6cefd51f81820fc5c6e5d7480df + +# Spilling round: the same query under a 600 KB limit must spill. +statement ok +SET datafusion.runtime.memory_limit = '600K' + +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, min(v1 * 2), avg(v1) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=Single,aggr=[min(t1.v1 * Int64(2)), avg(t1.v1)], ordering_mode=PartiallySorted([0]), metrics=[spilled_bytes= KB,] + + +# Same result hash as the no-spill round above +query IIIR rowsort +SELECT round(v1, -4), v1 % 5000, min(v1 * 2), avg(v1) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +60000 values hashing to 872df6cefd51f81820fc5c6e5d7480df + statement ok RESET datafusion.runtime.memory_limit