diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 5db3d23ca1c8c..6457deeda98ae 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -910,7 +910,7 @@ async fn test_aggregate_with_pk2() -> Result<()> { assert_snapshot!( physical_plan_to_string(&df).await, @r" - AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[], ordering_mode=Sorted + AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[] FilterExec: id@0 = 1 AND name@1 = a DataSourceExec: partitions=1, partition_sizes=[1] " @@ -958,7 +958,7 @@ async fn test_aggregate_with_pk3() -> Result<()> { assert_snapshot!( physical_plan_to_string(&df).await, @r" - AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[], ordering_mode=PartiallySorted([0]) + AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[] FilterExec: id@0 = 1 DataSourceExec: partitions=1, partition_sizes=[1] " @@ -1008,7 +1008,7 @@ async fn test_aggregate_with_pk4() -> Result<()> { assert_snapshot!( physical_plan_to_string(&df).await, @r" - AggregateExec: mode=Single, gby=[id@0 as id], aggr=[], ordering_mode=Sorted + AggregateExec: mode=Single, gby=[id@0 as id], aggr=[] FilterExec: id@0 = 1 DataSourceExec: partitions=1, partition_sizes=[1] " diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index 26e6e0c74c49d..4c5c8af984cee 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -740,12 +740,12 @@ fn test_pushdown_through_aggregates_on_grouping_columns() { OptimizationTest: input: - FilterExec: b@1 = bar - - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=PartiallySorted([0]) + - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt] - FilterExec: a@0 = foo - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true output: Ok: - - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=Sorted + - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt] - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=a@0 = foo AND b@1 = bar " ); @@ -804,7 +804,7 @@ fn test_pushdown_through_aggregates_preserves_parent_filter_order() { output: Ok: - FilterExec: cnt@2 = 1 - - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=PartiallySorted([1]) + - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt] - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar " ); @@ -2099,7 +2099,7 @@ fn test_pushdown_filter_on_non_first_grouping_column() { - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true output: Ok: - - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=PartiallySorted([1]) + - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt] - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar " ); @@ -2435,7 +2435,7 @@ fn test_pushdown_through_aggregate_with_reordered_input_columns() { - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true output: Ok: - - AggregateExec: mode=Final, gby=[a@1 as a, b@2 as b], aggr=[cnt], ordering_mode=PartiallySorted([1]) + - AggregateExec: mode=Final, gby=[a@1 as a, b@2 as b], aggr=[cnt] - ProjectionExec: expr=[c@2 as c, a@0 as a, b@1 as b] - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar " diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index 32f02d0f6a2f5..800d459c339ab 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -321,6 +321,27 @@ async fn test_aggregate() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_constant_grouping_expr_on_unbounded_input() -> Result<()> { + let case = QueryCase { + sql: "SELECT c1, c2, MIN(c4) FROM test WHERE c2 = 1 GROUP BY c1, c2".to_string(), + cases: vec![ + Arc::new(UnaryTestCase { + source_type: SourceType::Bounded, + expect_fail: false, + }), + Arc::new(UnaryTestCase { + source_type: SourceType::Unbounded, + expect_fail: true, + }), + ], + error_operator: "operator: AggregateExec".to_string(), + }; + + case.run().await?; + Ok(()) +} + #[tokio::test] async fn test_window_agg_hash_partition() -> Result<()> { let test1 = UnaryTestCase { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 854fbf8484eb9..ab8ab7ffdd105 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1029,6 +1029,8 @@ impl AggregateExec { let required_input_ordering = 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` @@ -1037,9 +1039,18 @@ impl AggregateExec { let indices: Vec = indices .into_iter() .filter(|idx| group_by.groups.iter().all(|group| !group[*idx])) + .filter(|idx| { + input_eq_properties + .is_expr_constant(&groupby_exprs[*idx]) + .is_none() + }) .collect(); - let mut input_order_mode = if indices.len() == groupby_exprs.len() + let num_non_constant_groupby_exprs = groupby_exprs + .iter() + .filter(|expr| input_eq_properties.is_expr_constant(expr).is_none()) + .count(); + let mut input_order_mode = if indices.len() == num_non_constant_groupby_exprs && !indices.is_empty() && group_by.groups.len() == 1 { @@ -3259,7 +3270,7 @@ mod tests { use arrow::array::{ BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array, - Int64Array, NullArray, StructArray, UInt32Array, UInt64Array, + Int64Array, NullArray, StringArray, StructArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::Int32Type; @@ -3270,7 +3281,7 @@ mod tests { use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::{ - Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, + Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, Operator, Signature, Volatility, }; use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; @@ -3284,7 +3295,7 @@ mod tests { use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; - use datafusion_physical_expr::expressions::{Literal, NotExpr}; + use datafusion_physical_expr::expressions::{Literal, NotExpr, binary}; use crate::projection::ProjectionExec; use crate::repartition::RepartitionExec; @@ -4728,6 +4739,87 @@ mod tests { Ok(()) } + #[tokio::test] + async fn constant_grouping_expr_is_not_a_completion_boundary() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("market", DataType::Utf8, true), + Field::new("value", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![None, Some(10), Some(10)])), + Arc::new(StringArray::from(vec![Some("US"), Some("US"), Some("US")])), + Arc::new(Int64Array::from(vec![3, 1, 2])), + ], + )?; + + let build_aggregate = |input: Arc| -> Result { + let predicate = + binary(col("market", &schema)?, Operator::Eq, lit("US"), &schema)?; + let input = Arc::new(FilterExecBuilder::new(predicate, input).build()?); + AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![ + (col("key", &schema)?, "key".to_string()), + (col("market", &schema)?, "market".to_string()), + ]), + vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(value)") + .build()?, + )], + vec![None], + input, + Arc::clone(&schema), + ) + }; + + let unordered_input = TestMemoryExec::try_new_exec( + &[vec![batch.clone()]], + Arc::clone(&schema), + None, + )?; + let aggregate = build_aggregate(unordered_input)?; + assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Linear); + assert_eq!( + aggregate.schema().as_ref(), + &Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("market", DataType::Utf8, true), + Field::new("COUNT(value)", DataType::Int64, false), + ]) + ); + + let output = + collect(aggregate.execute(0, Arc::new(TaskContext::default()))?).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::(), 2); + assert_snapshot!(batches_to_sort_string(&output), @r" ++-----+--------+--------------+ +| key | market | COUNT(value) | ++-----+--------+--------------+ +| | US | 1 | +| 10 | US | 2 | ++-----+--------+--------------+ +"); + + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("key", 0), + ))]) + .unwrap(); + let ordered_input = + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let ordered_input = + Arc::new(TestMemoryExec::update_cache(&Arc::new(ordered_input))); + let aggregate = build_aggregate(ordered_input)?; + assert_eq!(aggregate.input_order_mode(), &InputOrderMode::Sorted); + + Ok(()) + } + fn partial_reduce_test_aggregate() -> Result { partial_reduce_test_aggregate_with_batches(1) } @@ -6615,7 +6707,10 @@ mod tests { "); } } - Err(e) => assert!(matches!(e, DataFusionError::ResourcesExhausted(_))), + Err(e) => assert!( + matches!(e.find_root(), DataFusionError::ResourcesExhausted(_)), + "unexpected error: {e}" + ), } Ok(()) diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index a4cd3b0ccc1f0..464777fa5a5b6 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -8751,9 +8751,9 @@ logical_plan 15)------------EmptyRelation: rows=1 physical_plan 01)ProjectionExec: expr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST]@1 as last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))@2 as sum(DISTINCT Int64(1))] -02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))], ordering_mode=Sorted +02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))] 03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=5 -04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))], ordering_mode=Sorted +04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[last_value(a.foo) ORDER BY [a.foo ASC NULLS LAST], sum(DISTINCT Int64(1))] 05)--------UnionExec 06)----------ProjectionExec: expr=[1 as id, 2 as foo] 07)------------PlaceholderRowExec diff --git a/datafusion/sqllogictest/test_files/array_agg.slt b/datafusion/sqllogictest/test_files/array_agg.slt index b3cd60b9fed8c..451241756d4ec 100644 --- a/datafusion/sqllogictest/test_files/array_agg.slt +++ b/datafusion/sqllogictest/test_files/array_agg.slt @@ -574,9 +574,9 @@ logical_plan 15)------------EmptyRelation: rows=1 physical_plan 01)ProjectionExec: expr=[array_length(array_agg(DISTINCT a.foo)@1) as array_length(array_agg(DISTINCT a.foo)), sum(DISTINCT Int64(1))@2 as sum(DISTINCT Int64(1))] -02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted +02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))] 03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=5 -04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted +04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))] 05)--------UnionExec 06)----------ProjectionExec: expr=[1 as id, CAST(2 AS Int64) as foo] 07)------------PlaceholderRowExec diff --git a/datafusion/sqllogictest/test_files/push_down_filter_regression.slt b/datafusion/sqllogictest/test_files/push_down_filter_regression.slt index 23013d481fa19..c327f8092f871 100644 --- a/datafusion/sqllogictest/test_files/push_down_filter_regression.slt +++ b/datafusion/sqllogictest/test_files/push_down_filter_regression.slt @@ -512,9 +512,9 @@ query TT EXPLAIN SELECT a, count(b) FROM agg_filter_pushdown WHERE a = 'x' GROUP BY a; ---- physical_plan -01)AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(agg_filter_pushdown.b)], ordering_mode=Sorted +01)AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(agg_filter_pushdown.b)] 02)--RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 -03)----AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(agg_filter_pushdown.b)], ordering_mode=Sorted +03)----AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(agg_filter_pushdown.b)] 04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_filter_pushdown.parquet]]}, projection=[a, b], file_type=parquet, predicate=a@0 = x, pruning_predicate=a_null_count@2 != row_count@3 AND a_min@0 <= x AND x <= a_max@1, required_guarantees=[a in (x)] diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index 9d953f5075425..115a010330103 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -783,7 +783,7 @@ logical_plan physical_plan 01)UnionExec 02)--ProjectionExec: expr=[Int64(1)@0 as a] -03)----AggregateExec: mode=SinglePartitioned, gby=[1 as Int64(1)], aggr=[], ordering_mode=Sorted +03)----AggregateExec: mode=SinglePartitioned, gby=[1 as Int64(1)], aggr=[] 04)------PlaceholderRowExec 05)--ProjectionExec: expr=[2 as a] 06)----PlaceholderRowExec @@ -811,7 +811,7 @@ logical_plan physical_plan 01)UnionExec 02)--ProjectionExec: expr=[count(Int64(1))@1 as count, CAST(n@0 AS Int64) as n] -03)----AggregateExec: mode=SinglePartitioned, gby=[n@0 as n], aggr=[count(Int64(1))], ordering_mode=Sorted +03)----AggregateExec: mode=SinglePartitioned, gby=[n@0 as n], aggr=[count(Int64(1))] 04)------ProjectionExec: expr=[5 as n] 05)--------PlaceholderRowExec 06)--ProjectionExec: expr=[1 as count, max(Int64(10))@0 as n] diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 6ffe1b4cc087c..b0eac5528a382 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -276,7 +276,7 @@ physical_plan 03)----SortExec: expr=[b@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[max(d.a)] 05)--------RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[b@1 as b], aggr=[max(d.a)], ordering_mode=Sorted +06)----------AggregateExec: mode=Partial, gby=[b@1 as b], aggr=[max(d.a)] 07)------------UnionExec 08)--------------ProjectionExec: expr=[1 as a, aa as b] 09)----------------PlaceholderRowExec