fix: clear groupings when clearing orderings

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()?;

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();

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)]

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(())
     }
 

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);