From c11a7f072f2a0a97138c0f4826b8d45a568be19f Mon Sep 17 00:00:00 2001 From: xavlee <52435309+xavlee@users.noreply.github.com> Date: Fri, 4 Sep 2026 14:15:30 -0400 Subject: [PATCH 1/5] feat: stream aggregates over grouped input --- .../physical-plan/src/aggregates/mod.rs | 104 +++++++++++++++--- 1 file changed, 89 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 39389ad77df9b..86f1d23d156ae 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1044,6 +1044,10 @@ impl AggregateExec { // states without using group boundaries to recognize completed groups. if group_by.has_grouping_set() || mode == AggregateMode::PartialReduce { group_clustering_mode = GroupClusteringMode::None; + } else if num_non_constant_groupby_exprs > 0 + && input_eq_properties.grouping_satisfy(groupby_exprs.iter().cloned())? + { + group_clustering_mode = GroupClusteringMode::Full; } // construct a map from the input expression to the output expression of the Aggregation group by @@ -2457,7 +2461,7 @@ impl ExecutionPlan for AggregateExec { metrics: _, // Derived at construction from the input ordering and `group_by`. required_input_ordering: _, - // Derived at construction from the input ordering and `group_by`. + // Derived at construction from the input properties and `group_by`. group_clustering_mode: _, // Derived at construction by `Self::compute_properties`. cache: _, @@ -3381,7 +3385,7 @@ mod tests { Int64Array, NullArray, StringArray, StructArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; - use arrow::datatypes::{Float64Type, Int32Type, Int64Type, UInt32Type}; + use arrow::datatypes::{Float64Type, Int32Type, Int64Type, TimeUnit, UInt32Type}; use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; use datafusion_common::{DataFusionError, assert_contains, internal_err}; use datafusion_execution::config::SessionConfig; @@ -3394,6 +3398,7 @@ mod tests { Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, Operator, Signature, Volatility, }; + use datafusion_functions::datetime::date_bin; use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; use datafusion_functions_aggregate::array_agg::array_agg_udaf; use datafusion_functions_aggregate::average::avg_udaf; @@ -3402,10 +3407,9 @@ mod tests { use datafusion_functions_aggregate::median::median_udaf; use datafusion_functions_aggregate::min_max::min_udaf; use datafusion_functions_aggregate::sum::sum_udaf; - use datafusion_physical_expr::Partitioning; - use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::{Literal, NotExpr, binary}; + use datafusion_physical_expr::{Partitioning, PhysicalSortExpr, ScalarFunctionExpr}; use crate::projection::ProjectionExec; use crate::repartition::RepartitionExec; @@ -5961,7 +5965,7 @@ mod tests { } #[tokio::test] - async fn unsorted_contiguous_groups_use_final_emission() -> Result<()> { + async fn unsorted_contiguous_groups_use_incremental_emission() -> Result<()> { let schema = Arc::new(Schema::new(vec![ Field::new("key", DataType::Int32, false), Field::new("time_bin", DataType::Int64, false), @@ -5989,9 +5993,11 @@ mod tests { ], )?, ]; + let key = col("key", &schema)?; + let time_bin = col("time_bin", &schema)?; let group_by = PhysicalGroupBy::new_single(vec![ - (col("key", &schema)?, "key".to_string()), - (col("time_bin", &schema)?, "time_bin".to_string()), + (Arc::clone(&key), "key".to_string()), + (Arc::clone(&time_bin), "time_bin".to_string()), ]); let aggr_expr = Arc::new( AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) @@ -5999,8 +6005,9 @@ mod tests { .alias("SUM(value)") .build()?, ); - let input: Arc = - TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; + let input = TestMemoryExec::try_new(&[input_batches], Arc::clone(&schema), None)? + .try_with_grouping_information(vec![vec![key, time_bin]])?; + let input: Arc = Arc::new(input); assert_eq!(input.output_partitioning().partition_count(), 1); let aggregate = AggregateExec::try_new( @@ -6014,16 +6021,21 @@ mod tests { assert_eq!( aggregate.group_clustering_mode(), - &GroupClusteringMode::None + &GroupClusteringMode::Full + ); + assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental); + assert!( + aggregate + .cache() + .equivalence_properties() + .geq_class() + .is_empty() ); - // This captures the behavior before #24438. When the source can declare - // `(key, time_bin)` group-contiguous, the corresponding case can use - // `EmissionType::Incremental`. - assert_eq!(aggregate.cache().emission_type, EmissionType::Final); + assert!(aggregate.cache().output_ordering().is_none()); let task_ctx = new_migrated_hash_ctx(1024); let stream = aggregate.execute_typed(0, &task_ctx)?; - assert!(matches!(stream, StreamType::SingleHash(_))); + assert!(matches!(stream, StreamType::ClusteredSingleAggregate(_))); let stream: SendableRecordBatchStream = stream.into(); let output = collect(stream).await?; assert_snapshot!(batches_to_sort_string(&output), @r" @@ -6040,6 +6052,68 @@ mod tests { Ok(()) } + #[test] + fn grouped_date_bin_projects_to_aggregate() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("time", DataType::Timestamp(TimeUnit::Second, None), false), + ])); + let time_bin_expr = || -> Result> { + Ok(Arc::new(ScalarFunctionExpr::try_new( + date_bin(), + vec![ + lit(ScalarValue::new_interval_dt(0, 10_000)), + col("time", &schema)?, + ], + &schema, + Arc::new(ConfigOptions::default()), + )?)) + }; + let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_grouping_information(vec![vec![ + col("key", &schema)?, + time_bin_expr()?, + ]])?; + let projection = ProjectionExec::try_new( + [ + ProjectionExpr::new(col("key", &schema)?, "key"), + // Build this expression independently from the source + // assertion so the test exercises semantic expression matching. + ProjectionExpr::new(time_bin_expr()?, "time_bin"), + ], + Arc::new(input), + )?; + + let projected_schema = projection.schema(); + let key = col("key", &projected_schema)?; + let time_bin = col("time_bin", &projected_schema)?; + assert!( + projection + .properties() + .equivalence_properties() + .grouping_satisfy([Arc::clone(&key), Arc::clone(&time_bin)])? + ); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![ + (key, "key".to_string()), + (time_bin, "time_bin".to_string()), + ]), + vec![], + vec![], + Arc::new(projection), + projected_schema, + )?; + + assert!(aggregate.input().output_ordering().is_none()); + assert_eq!( + aggregate.group_clustering_mode(), + &GroupClusteringMode::Full + ); + assert_eq!(aggregate.cache().emission_type, EmissionType::Incremental); + Ok(()) + } + /// Ensures for ordered input, `ClusteredPartialAggregateStream` is used. #[tokio::test] async fn ordered_partial_aggregate_planning() -> Result<()> { From 2d6d6004e2863227a6a281fc944264c3641545ed Mon Sep 17 00:00:00 2001 From: xavlee <52435309+xavlee@users.noreply.github.com> Date: Wed, 23 Sep 2026 12:04:50 -0400 Subject: [PATCH 2/5] fix: guard grouped aggregate completion and test early emission Keep PartialReduce on hash execution and require a non-constant grouping expression for full group completion. Cover both cases and gate the second input batch to verify completed groups are emitted before input finishes. --- .../physical-plan/src/aggregates/mod.rs | 194 +++++++++++++----- datafusion/physical-plan/src/test/exec.rs | 32 +++ 2 files changed, 179 insertions(+), 47 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 86f1d23d156ae..30935a3d7b47d 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3376,7 +3376,8 @@ mod tests { use crate::test::TestMemoryExec; use crate::test::assert_is_pending; use crate::test::exec::{ - BlockingExec, PanicExec, StatisticsExec, assert_strong_count_converges_to_zero, + BarrierExec, BlockingExec, PanicExec, StatisticsExec, + assert_strong_count_converges_to_zero, }; use arrow::array::AsArray; @@ -5569,7 +5570,7 @@ mod tests { } #[test] - fn partial_reduce_does_not_advertise_input_ordering() -> Result<()> { + fn partial_reduce_does_not_use_group_clustering() -> Result<()> { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, false), Field::new("b", DataType::Float64, false), @@ -5578,39 +5579,113 @@ mod tests { Column::new("a", 0), ))]) .unwrap(); - let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? - .try_with_sort_information(vec![ordering])?; - let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); - assert!( - input.properties().output_ordering().is_some(), - "test setup: the input is ordered by the group key" - ); + let ordered_input = + TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let grouped_input = + TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_grouping_information(vec![vec![col("a", &schema)?]])?; - let partial_reduce = AggregateExec::try_new( - AggregateMode::PartialReduce, - PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]), - vec![Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("SUM(b)") - .build()?, - )], - vec![None], - input, + for input in [ordered_input, grouped_input] { + assert!( + input + .properties() + .equivalence_properties() + .grouping_satisfy([col("a", &schema)?])? + ); + let partial_reduce = AggregateExec::try_new( + AggregateMode::PartialReduce, + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )], + vec![None], + Arc::new(input), + Arc::clone(&schema), + )?; + + assert_eq!( + partial_reduce.group_clustering_mode(), + &GroupClusteringMode::None + ); + assert_eq!(partial_reduce.maintains_input_order(), vec![false]); + assert!( + partial_reduce.properties().output_ordering().is_none(), + "partial reduce advertised an ordering it does not maintain: {:?}", + partial_reduce.properties().output_ordering() + ); + + let stream = partial_reduce.execute_typed(0, &new_migrated_hash_ctx(1024))?; + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + } + + Ok(()) + } + + #[tokio::test] + async fn constant_only_grouping_uses_final_emission() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int64Array::from(vec![10, 20])), + ], + )?; + let input = TestMemoryExec::try_new_exec( + &[vec![batch.clone(), batch]], Arc::clone(&schema), + None, )?; + let predicate = binary(col("key", &schema)?, Operator::Eq, lit(1i32), &schema)?; + let input: Arc = + Arc::new(FilterExecBuilder::new(predicate, input).build()?); - assert_eq!( - partial_reduce.group_clustering_mode(), - &GroupClusteringMode::None - ); - assert_eq!(partial_reduce.maintains_input_order(), vec![false]); - assert!( - partial_reduce.properties().output_ordering().is_none(), - "partial reduce advertised an ordering it does not maintain: {:?}", - partial_reduce.properties().output_ordering() - ); + // Cover both a literal and a column known to be constant after filtering. + for key in [lit(1i32), col("key", &schema)?] { + assert!( + input + .equivalence_properties() + .grouping_satisfy([Arc::clone(&key)])? + ); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![(key, "key".to_string())]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + )], + vec![None], + Arc::clone(&input), + Arc::clone(&schema), + )?; + assert_eq!( + aggregate.group_clustering_mode(), + &GroupClusteringMode::None + ); + assert_eq!(aggregate.cache().emission_type, EmissionType::Final); + let stream = aggregate.execute_typed(0, &new_migrated_hash_ctx(1024))?; + assert!(matches!(stream, StreamType::SingleHash(_))); + let output = collect(stream.into()).await?; + allow_duplicates! { + assert_snapshot!(batches_to_string(&output), @r" + +-----+------------+ + | key | SUM(value) | + +-----+------------+ + | 1 | 60 | + +-----+------------+ + "); + } + } Ok(()) } @@ -5971,25 +6046,24 @@ mod tests { Field::new("time_bin", DataType::Int64, false), Field::new("value", DataType::Int64, false), ])); - // Two sorted logical runs are emitted as batches in one DataFusion - // partition. Every distinct grouping tuple occupies one contiguous range, - // but tuple order resets at the batch boundary, so (key, time_bin) is not - // globally sorted. + // Every grouping tuple occupies one contiguous range in a single + // partition, with (2, 20) spanning both batches. Tuple order resets + // within the second batch, so (key, time_bin) is not globally sorted. let input_batches = vec![ RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int32Array::from(vec![1, 1, 2, 2])), - Arc::new(Int64Array::from(vec![20, 20, 20, 20])), - Arc::new(Int64Array::from(vec![10, 20, 30, 40])), + Arc::new(Int32Array::from(vec![1, 1, 2])), + Arc::new(Int64Array::from(vec![20, 20, 20])), + Arc::new(Int64Array::from(vec![10, 20, 30])), ], )?, RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int32Array::from(vec![1, 1, 2, 2])), - Arc::new(Int64Array::from(vec![0, 0, 0, 0])), - Arc::new(Int64Array::from(vec![50, 60, 70, 80])), + Arc::new(Int32Array::from(vec![2, 1, 1, 2, 2])), + Arc::new(Int64Array::from(vec![20, 0, 0, 0, 0])), + Arc::new(Int64Array::from(vec![40, 50, 60, 70, 80])), ], )?, ]; @@ -6005,17 +6079,26 @@ mod tests { .alias("SUM(value)") .build()?, ); - let input = TestMemoryExec::try_new(&[input_batches], Arc::clone(&schema), None)? - .try_with_grouping_information(vec![vec![key, time_bin]])?; - let input: Arc = Arc::new(input); - assert_eq!(input.output_partitioning().partition_count(), 1); + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + eq_properties.add_groupings([vec![key, time_bin]]); + let input = Arc::new( + BarrierExec::new(vec![input_batches], Arc::clone(&schema)) + .without_start_barrier() + .with_batch_barrier() + .with_equivalence_properties(eq_properties) + .with_log(false), + ); + assert_eq!( + input.properties().output_partitioning().partition_count(), + 1 + ); let aggregate = AggregateExec::try_new( AggregateMode::Single, group_by, vec![aggr_expr], vec![None], - input, + Arc::clone(&input) as Arc, schema, )?; @@ -6036,8 +6119,25 @@ mod tests { let task_ctx = new_migrated_hash_ctx(1024); let stream = aggregate.execute_typed(0, &task_ctx)?; assert!(matches!(stream, StreamType::ClusteredSingleAggregate(_))); - let stream: SendableRecordBatchStream = stream.into(); - let output = collect(stream).await?; + let mut stream: SendableRecordBatchStream = stream.into(); + input.wait_batch().await; + // The second batch and EOF remain blocked until the first completed + // group has been emitted. + let first = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("completed group was not emitted before the next input batch") + .expect("stream ended before the next input batch")?; + assert_snapshot!(batches_to_string(std::slice::from_ref(&first)), @r" ++-----+----------+------------+ +| key | time_bin | SUM(value) | ++-----+----------+------------+ +| 1 | 20 | 30 | ++-----+----------+------------+ +"); + input.wait_batch().await; + let mut output = vec![first]; + output.extend(collect(stream).await?); assert_snapshot!(batches_to_sort_string(&output), @r" +-----+----------+------------+ | key | time_bin | SUM(value) | diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index 2f77f903c9ad1..f8648f6cdf0a6 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -318,6 +318,9 @@ pub struct BarrierExec { /// all streams wait on this barrier to produce start_data_barrier: Option>, + /// each stream waits on this barrier before sending each batch + batch_barrier: Option>, + /// the stream wait for this to return Poll::Ready(None) finish_barrier: Option>, @@ -336,6 +339,7 @@ impl BarrierExec { data, schema, start_data_barrier: barrier, + batch_barrier: None, cache: Arc::new(cache), finish_barrier: None, log: true, @@ -352,6 +356,30 @@ impl BarrierExec { self } + /// Wait for [`Self::wait_batch`] before sending each batch. + pub fn with_batch_barrier(mut self) -> Self { + self.batch_barrier = Some(Arc::new(Barrier::new(self.data.len() + 1))); + self + } + + /// Release the next batch in each partition. + pub async fn wait_batch(&self) { + self.batch_barrier + .as_ref() + .expect("Must only be called when having a batch barrier") + .wait() + .await; + } + + /// Declare the equivalence properties of the test data. + pub fn with_equivalence_properties( + mut self, + properties: EquivalenceProperties, + ) -> Self { + self.cache = Arc::new(self.cache.as_ref().clone().with_eq_properties(properties)); + self + } + pub fn with_finish_barrier(mut self) -> Self { let barrier = Arc::new(( // wait for all streams and the input @@ -486,6 +514,7 @@ impl ExecutionPlan for BarrierExec { // task simply sends data in order after barrier is reached let data = self.data[partition].clone(); let start_barrier = self.start_data_barrier.as_ref().map(Arc::clone); + let batch_barrier = self.batch_barrier.as_ref().map(Arc::clone); let finish_barrier = self.finish_barrier.as_ref().map(Arc::clone); let log = self.log; let tx = builder.tx(); @@ -497,6 +526,9 @@ impl ExecutionPlan for BarrierExec { barrier.wait().await; } for batch in data { + if let Some(barrier) = &batch_barrier { + barrier.wait().await; + } if log { println!("Partition {partition} sending batch"); } From 2a529bf2d4ee208e9ae65b987704a8d56cb4d010 Mon Sep 17 00:00:00 2001 From: xavlee <52435309+xavlee@users.noreply.github.com> Date: Fri, 6 Nov 2026 15:53:28 -0500 Subject: [PATCH 3/5] fix: disable group clustering for multiple group masks --- .../physical-plan/src/aggregates/mod.rs | 93 +++++++++++++++---- 1 file changed, 76 insertions(+), 17 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 30935a3d7b47d..1c9278246c681 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1009,7 +1009,6 @@ impl AggregateExec { LexRequirement::new(new_requirements).map(OrderingRequirements::new_soft); // Constant expressions never change, so they cannot mark a completed group. - // Exclude them from both the ordering indices and the group expression count. // If our aggregation has grouping sets then our base grouping exprs will // be expanded based on the flags in `group_by.groups` where for each // group we swap the grouping expr for `null` if the flag is `true` @@ -1025,13 +1024,18 @@ impl AggregateExec { }) .collect(); - let num_non_constant_groupby_exprs = groupby_exprs + let has_non_constant_groupby_exprs = groupby_exprs .iter() - .filter(|expr| input_eq_properties.is_expr_constant(expr).is_none()) - .count(); - let mut group_clustering_mode = if indices.len() == num_non_constant_groupby_exprs - && !indices.is_empty() - && group_by.groups.len() == 1 + .any(|expr| input_eq_properties.is_expr_constant(expr).is_none()); + // Grouping sets can change group keys. PartialReduce combines intermediate + // states without using group boundaries to recognize completed groups. + let group_clustering_mode = if group_by.has_grouping_set() + || group_by.groups.len() != 1 + || mode == AggregateMode::PartialReduce + { + GroupClusteringMode::None + } else if has_non_constant_groupby_exprs + && input_eq_properties.grouping_satisfy(groupby_exprs.iter().cloned())? { GroupClusteringMode::Full } else if !indices.is_empty() { @@ -1040,16 +1044,6 @@ impl AggregateExec { GroupClusteringMode::None }; - // Grouping sets can change group keys. PartialReduce combines intermediate - // states without using group boundaries to recognize completed groups. - if group_by.has_grouping_set() || mode == AggregateMode::PartialReduce { - group_clustering_mode = GroupClusteringMode::None; - } else if num_non_constant_groupby_exprs > 0 - && input_eq_properties.grouping_satisfy(groupby_exprs.iter().cloned())? - { - group_clustering_mode = GroupClusteringMode::Full; - } - // construct a map from the input expression to the output expression of the Aggregation group by let group_expr_mapping = ProjectionMapping::try_new(group_by.expr.clone(), &input.schema())?; @@ -5625,6 +5619,71 @@ mod tests { Ok(()) } + #[tokio::test] + async fn multiple_group_masks_do_not_use_group_clustering() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let a = col("a", &schema)?; + let b = col("b", &schema)?; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 2, 1])), + Arc::new(Int32Array::from(vec![20, 20, 10])), + Arc::new(Int64Array::from(vec![10, 20, 40])), + ], + )?; + let input = TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)?; + let grouped_input = input + .clone() + .try_with_grouping_information(vec![vec![Arc::clone(&a), Arc::clone(&b)]])?; + // Repeating a mask repeats the keys, so input grouping does not prove + // contiguity after the masks are expanded. + let group_by = PhysicalGroupBy::new( + vec![(a, "a".to_string()), (b, "b".to_string())], + vec![], + vec![vec![false, false]; 2], + false, + ); + let aggr_expr = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + ); + + for input in [input, grouped_input] { + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by.clone(), + vec![Arc::clone(&aggr_expr)], + vec![None], + Arc::new(input), + Arc::clone(&schema), + )?; + let output = + collect(aggregate.execute(0, new_migrated_hash_ctx(1024))?).await?; + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&output), @r" + +---+----+------------+ + | a | b | SUM(value) | + +---+----+------------+ + | 1 | 10 | 80 | + | 2 | 20 | 60 | + +---+----+------------+ + "); + } + assert_eq!( + aggregate.group_clustering_mode(), + &GroupClusteringMode::None + ); + } + Ok(()) + } + #[tokio::test] async fn constant_only_grouping_uses_final_emission() -> Result<()> { let schema = Arc::new(Schema::new(vec![ From 42a47346467c3e93dfef9d019c6e3c20bc87d03a Mon Sep 17 00:00:00 2001 From: xavlee <52435309+xavlee@users.noreply.github.com> Date: Mon, 9 Nov 2026 13:02:42 -0500 Subject: [PATCH 4/5] fix: validate grouped masks and cover partitioned aggregation --- .../physical-plan/src/aggregates/mod.rs | 170 +++++++++++++++++- .../library-user-guide/upgrading/57.0.0.md | 22 +++ .../library-user-guide/upgrading/index.rst | 1 + 3 files changed, 190 insertions(+), 3 deletions(-) create mode 100644 docs/source/library-user-guide/upgrading/57.0.0.md diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 1c9278246c681..3b892af4432de 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -37,8 +37,9 @@ //! //! ### Group clustering optimization //! -//! When the input is ordered by group keys, rows are clustered by those keys. -//! The clustered paths use this guarantee to emit completed groups early. +//! When the input is ordered by group keys, or its equivalence properties declare +//! them grouped, rows are clustered by those keys. The clustered paths use this +//! guarantee to emit completed groups early. //! //! ```text //! AggregateExec (final, clustered) @@ -1030,7 +1031,7 @@ impl AggregateExec { // Grouping sets can change group keys. PartialReduce combines intermediate // states without using group boundaries to recognize completed groups. let group_clustering_mode = if group_by.has_grouping_set() - || group_by.groups.len() != 1 + || group_by.groups != [vec![false; group_by.expr.len()]] || mode == AggregateMode::PartialReduce { GroupClusteringMode::None @@ -5684,6 +5685,79 @@ mod tests { Ok(()) } + #[tokio::test] + async fn nulling_group_mask_does_not_use_group_clustering() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let a = col("a", &schema)?; + let b = col("b", &schema)?; + let batches = [(1, 1, 10), (2, 2, 20), (3, 1, 30)] + .into_iter() + .map(|(a, b, value)| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![a])), + Arc::new(Int32Array::from(vec![b])), + Arc::new(Int64Array::from(vec![value])), + ], + ) + }) + .collect::, _>>()?; + let input = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)?; + let grouped_input = input + .clone() + .try_with_grouping_information(vec![vec![Arc::clone(&a), Arc::clone(&b)]])?; + // Nulling a turns distinct input tuples into (NULL, 1), (NULL, 2), + // (NULL, 1), which are no longer contiguous. + let group_by = PhysicalGroupBy::new( + vec![(a, "a".to_string()), (b, "b".to_string())], + vec![ + (lit(ScalarValue::Int32(None)), "a".to_string()), + (lit(ScalarValue::Int32(None)), "b".to_string()), + ], + vec![vec![true, false]], + false, + ); + let aggr_expr = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + ); + + for input in [input, grouped_input] { + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + group_by.clone(), + vec![Arc::clone(&aggr_expr)], + vec![None], + Arc::new(input), + Arc::clone(&schema), + )?; + let output = + collect(aggregate.execute(0, new_migrated_hash_ctx(1024))?).await?; + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&output), @r" + +---+---+------------+ + | a | b | SUM(value) | + +---+---+------------+ + | | 1 | 40 | + | | 2 | 20 | + +---+---+------------+ + "); + } + assert_eq!( + aggregate.group_clustering_mode(), + &GroupClusteringMode::None + ); + } + Ok(()) + } + #[tokio::test] async fn constant_only_grouping_uses_final_emission() -> Result<()> { let schema = Arc::new(Schema::new(vec![ @@ -6211,6 +6285,96 @@ mod tests { Ok(()) } + #[tokio::test] + async fn grouped_input_uses_clustered_partial_and_single_partitioned() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let batch = |keys: Vec, values: Vec| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int64Array::from(values)), + ], + ) + }; + // Groups are contiguous within each partition but not sorted. Key 1 spans + // a batch boundary, and keys 2 and 3 occur in both partitions. + let partitions = vec![ + vec![ + batch(vec![2, 1], vec![1, 2])?, + batch(vec![1, 3], vec![3, 4])?, + ], + vec![batch(vec![3, 3, 2], vec![5, 6, 7])?], + ]; + let key = col("key", &schema)?; + let input: Arc = Arc::new( + TestMemoryExec::try_new(&partitions, Arc::clone(&schema), None)? + .try_with_grouping_information(vec![vec![Arc::clone(&key)]])?, + ); + let group_by = PhysicalGroupBy::new_single(vec![(key, "key".to_string())]); + let aggr_expr = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + ); + let task_ctx = new_migrated_hash_ctx(1024); + + let single_partitioned = AggregateExec::try_new( + AggregateMode::SinglePartitioned, + group_by.clone(), + vec![Arc::clone(&aggr_expr)], + vec![None], + Arc::clone(&input), + Arc::clone(&schema), + )?; + assert_eq!( + single_partitioned.group_clustering_mode(), + &GroupClusteringMode::Full + ); + assert!(matches!( + single_partitioned.execute_typed(0, &task_ctx)?, + StreamType::ClusteredSingleAggregate(_) + )); + + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![Arc::clone(&aggr_expr)], + vec![None], + input, + Arc::clone(&schema), + )?); + assert_eq!(partial.group_clustering_mode(), &GroupClusteringMode::Full); + assert!(matches!( + partial.execute_typed(0, &task_ctx)?, + StreamType::ClusteredPartialAggregate(_) + )); + + let final_aggregate = AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![aggr_expr], + vec![None], + Arc::new(CoalescePartitionsExec::new(partial)), + schema, + )?; + let output = collect(final_aggregate.execute(0, task_ctx)?).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" + +-----+------------+ + | key | SUM(value) | + +-----+------------+ + | 1 | 5 | + | 2 | 8 | + | 3 | 15 | + +-----+------------+ + "); + Ok(()) + } + #[test] fn grouped_date_bin_projects_to_aggregate() -> Result<()> { let schema = Arc::new(Schema::new(vec![ diff --git a/docs/source/library-user-guide/upgrading/57.0.0.md b/docs/source/library-user-guide/upgrading/57.0.0.md new file mode 100644 index 0000000000000..422912f701fde --- /dev/null +++ b/docs/source/library-user-guide/upgrading/57.0.0.md @@ -0,0 +1,22 @@ + + +# Upgrade Guides + +## DataFusion 57.0.0 + +**Note:** DataFusion `57.0.0` has not been released yet. The information provided +in this section pertains to features and changes that have already been merged +to the main branch and are awaiting release in this version. + +### `GroupClusteringMode::Full` no longer implies sorted input + +`AggregateExec` can now select `GroupClusteringMode::Full` when the input's equivalence properties declare the complete `GROUP BY` tuple as grouped (contiguous within each partition), even if the input is not sorted. + +`Full` describes group contiguity only. Code that previously used +`group_clustering_mode() == GroupClusteringMode::Full` to infer sorted input or +output should use `properties().output_ordering()` on the input or aggregate +plan, respectively. diff --git a/docs/source/library-user-guide/upgrading/index.rst b/docs/source/library-user-guide/upgrading/index.rst index f64662f5865d2..e3cdacc104f39 100644 --- a/docs/source/library-user-guide/upgrading/index.rst +++ b/docs/source/library-user-guide/upgrading/index.rst @@ -8,6 +8,7 @@ Upgrade Guides .. toctree:: :maxdepth: 1 + DataFusion 57.0.0 <57.0.0> DataFusion 56.0.0 <56.0.0> DataFusion 55.0.0 <55.0.0> DataFusion 54.0.0 <54.0.0> From cff750b0011a2d9475209f43e983301b204c657c Mon Sep 17 00:00:00 2001 From: xavlee <52435309+xavlee@users.noreply.github.com> Date: Mon, 9 Nov 2026 23:52:17 -0500 Subject: [PATCH 5/5] docs: wrap grouped aggregation upgrade note --- docs/source/library-user-guide/upgrading/57.0.0.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/source/library-user-guide/upgrading/57.0.0.md b/docs/source/library-user-guide/upgrading/57.0.0.md index 422912f701fde..25f88deff8bfd 100644 --- a/docs/source/library-user-guide/upgrading/57.0.0.md +++ b/docs/source/library-user-guide/upgrading/57.0.0.md @@ -14,7 +14,9 @@ to the main branch and are awaiting release in this version. ### `GroupClusteringMode::Full` no longer implies sorted input -`AggregateExec` can now select `GroupClusteringMode::Full` when the input's equivalence properties declare the complete `GROUP BY` tuple as grouped (contiguous within each partition), even if the input is not sorted. +`AggregateExec` can now select `GroupClusteringMode::Full` when the input's +equivalence properties declare the complete `GROUP BY` tuple as grouped +(contiguous within each partition), even if the input is not sorted. `Full` describes group contiguity only. Code that previously used `group_clustering_mode() == GroupClusteringMode::Full` to infer sorted input or