diff --git a/datafusion/physical-expr/src/equivalence/mod.rs b/datafusion/physical-expr/src/equivalence/mod.rs index ff57146efcce9..48e676bcf707f 100644 --- a/datafusion/physical-expr/src/equivalence/mod.rs +++ b/datafusion/physical-expr/src/equivalence/mod.rs @@ -267,6 +267,50 @@ mod tests { Ok(()) } + #[test] + fn clear_orderings_clears_groupings() -> Result<()> { + let schema = create_test_schema()?; + let a = col("a", &schema)?; + let b = col("b", &schema)?; + let ordering = [PhysicalSortExpr::new_default(Arc::clone(&a))]; + let mut properties = + EquivalenceProperties::new_with_orderings(schema, [ordering.clone()]); + properties.add_grouping([Arc::clone(&b)]); + assert!(properties.grouping_satisfy([Arc::clone(&a)])?); + assert!(properties.grouping_satisfy([Arc::clone(&b)])?); + + properties.clear_orderings(); + + assert!(properties.oeq_class().is_empty()); + assert!(properties.geq_class().is_empty()); + assert!(properties.normalized_geq_class().is_empty()); + assert!(!properties.ordering_satisfy(ordering)?); + assert!(!properties.grouping_satisfy([a])?); + assert!(!properties.grouping_satisfy([b])?); + Ok(()) + } + + #[test] + fn clear_groupings_preserves_orderings() -> Result<()> { + let schema = create_test_schema()?; + let a = col("a", &schema)?; + let b = col("b", &schema)?; + let ordering = [PhysicalSortExpr::new_default(Arc::clone(&a))]; + let mut properties = + EquivalenceProperties::new_with_orderings(schema, [ordering.clone()]); + properties.add_grouping([Arc::clone(&b)]); + assert!(properties.grouping_satisfy([Arc::clone(&b)])?); + + properties.clear_groupings(); + + assert!(properties.geq_class().is_empty()); + assert!(properties.normalized_geq_class().is_empty()); + assert!(properties.ordering_satisfy(ordering)?); + assert!(properties.grouping_satisfy([a])?); + assert!(!properties.grouping_satisfy([b])?); + Ok(()) + } + #[test] fn reorder_clears_explicit_grouping() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 1206315bbb274..3b10fc80a58db 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -350,15 +350,18 @@ impl EquivalenceProperties { Ok(self) } - /// Clears (empties) the ordering equivalence class within this object. - /// Call this method when existing orderings are invalidated. + /// Clears ordering and explicit grouping information. + /// + /// Call this method when existing orderings are invalidated. Grouping + /// assertions are also cleared because the operation may split groups. pub fn clear_orderings(&mut self) { self.oeq_class.clear(); self.oeq_cache.clear(); + // Row reordering can also split groups, invalidating their contiguity. + self.clear_groupings(); } - /// Clears grouping information invalidated by an operation that changes - /// row sequence or combines input partitions. + /// Clears explicit grouping information without changing orderings. pub fn clear_groupings(&mut self) { self.geq_class.clear(); self.geq_cache.clear(); diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 369cf77ae5199..7983a9a7e223d 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -5365,6 +5365,62 @@ mod tests { Ok(()) } + #[rstest::rstest] + fn clustered_aggregate_clears_explicit_grouping( + #[values(AggregateMode::Partial, AggregateMode::Single)] mode: AggregateMode, + ) -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + Field::new("value", DataType::Int64, true), + ])); + let a = col("a", &schema)?; + let c = col("c", &schema)?; + let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ + [PhysicalSortExpr::new_default(Arc::clone(&a))].into(), + ])? + .try_with_grouping_information(vec![vec![Arc::clone(&c)]])?; + assert!( + input + .properties() + .equivalence_properties() + .grouping_satisfy([Arc::clone(&c)])? + ); + let aggregate = AggregateExec::try_new( + mode, + PhysicalGroupBy::new_single(vec![ + (a, "group_a".to_string()), + (c, "group_c".to_string()), + ]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(value)") + .build()?, + )], + vec![None], + Arc::new(input), + schema, + )?; + + assert_eq!( + aggregate.group_clustering_mode(), + &GroupClusteringMode::Partial(vec![0]) + ); + let properties = aggregate.properties().equivalence_properties(); + // Preserving the ordering on a does not preserve the grouping on c. + assert!( + properties.ordering_satisfy([PhysicalSortExpr::new_default(col( + "group_a", + &aggregate.schema(), + )?)])? + ); + assert!(properties.geq_class().is_empty()); + assert!(!properties.grouping_satisfy([col("group_c", &aggregate.schema())?])?); + Ok(()) + } + #[rstest::rstest] #[case::full(true)] #[case::partial(false)] diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index e0d8ae173e670..a2f5324e92cbc 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -84,9 +84,6 @@ impl CoalescePartitionsExec { // Coalescing partitions loses existing orderings: let mut eq_properties = input.equivalence_properties().clone(); eq_properties.clear_orderings(); - if input_partitions > 1 { - eq_properties.clear_groupings(); - } eq_properties.clear_per_partition_constants(); PlanProperties::new( eq_properties, // Equivalence Properties @@ -475,9 +472,11 @@ mod tests { use futures::FutureExt; - #[test] - fn grouping_is_cleared_only_when_partitions_are_merged() -> Result<()> { - let input = test::mem_exec(2); + #[rstest::rstest] + #[case::single_partition(1)] + #[case::multiple_partitions(2)] + fn coalesce_clears_explicit_grouping(#[case] partitions: usize) -> Result<()> { + let input = test::mem_exec(partitions); let grouping = col("i", &input.schema())?; let input = Arc::new( input.try_with_grouping_information(vec![vec![Arc::clone(&grouping)]])?, @@ -492,16 +491,6 @@ mod tests { let properties = merge.properties().equivalence_properties(); assert!(properties.geq_class().is_empty()); assert!(!properties.grouping_satisfy([grouping])?); - - let input = test::mem_exec(1); - let grouping = col("i", &input.schema())?; - let input = Arc::new( - input.try_with_grouping_information(vec![vec![Arc::clone(&grouping)]])?, - ); - let merge = CoalescePartitionsExec::new(input); - let properties = merge.properties().equivalence_properties(); - assert_eq!(properties.geq_class().len(), 1); - assert!(properties.grouping_satisfy([grouping])?); Ok(()) } diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 641e932c280c1..2592892a06bdf 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -164,7 +164,6 @@ impl SortPreservingMergeExec { // Other input orderings hold only within each partition: if input_partitions > 1 { eq_properties.clear_orderings(); - eq_properties.clear_groupings(); } eq_properties.clear_per_partition_constants(); eq_properties.add_ordering(ordering);