fix: special case to not remove the needed coalesce

Author: wiedld

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -2174,21 +2174,20 @@ async fn test_preserve_needed_coalesce() -> Result<()> {
         get_plan_string(&optimized),
         vec![
             "SortPreservingMergeExec: [a@0 ASC]",
-            "  SortExec: expr=[a@0 ASC], preserve_partitioning=[true]",
+            "  SortExec: expr=[a@0 ASC], preserve_partitioning=[false]",
             "    AggregateExec: mode=SinglePartitioned, gby=[a@0 as a1], aggr=[]",
-            "      ProjectionExec: expr=[a@0 as a, b@1 as value]",
-            "        UnionExec",
-            "          ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
-            "          ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+            "      CoalescePartitionsExec",
+            "        ProjectionExec: expr=[a@0 as a, b@1 as value]",
+            "          UnionExec",
+            "            ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
+            "            ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e]",
         ],
     );
 
-    // Plan is now invalid.
+    // Plan is valid.
     let checker = SanityCheckPlan::new();
-    let err = checker
-        .optimize(optimized, &Default::default())
-        .unwrap_err();
-    assert!(err.message().contains(" does not satisfy distribution requirements: HashPartitioned[[a@0]]). Child-0 output partitioning: UnknownPartitioning(2)"));
+    let checker = checker.optimize(optimized, &Default::default());
+    assert!(checker.is_ok());
 
     Ok(())
 }

datafusion/physical-optimizer/src/enforce_sorting/mod.rs

@@ -47,8 +47,8 @@ use crate::enforce_sorting::sort_pushdown::{
     assign_initial_requirements, pushdown_sorts, SortPushDown,
 };
 use crate::utils::{
-    add_sort_above, add_sort_above_with_check, is_coalesce_partitions, is_limit,
-    is_repartition, is_sort, is_sort_preserving_merge, is_union, is_window,
+    add_sort_above, add_sort_above_with_check, is_aggregation, is_coalesce_partitions,
+    is_limit, is_repartition, is_sort, is_sort_preserving_merge, is_union, is_window,
 };
 use crate::PhysicalOptimizerRule;
 
@@ -518,7 +518,7 @@ fn remove_bottleneck_in_subplan(
 ) -> Result<PlanWithCorrespondingCoalescePartitions> {
     let plan = &requirements.plan;
     let children = &mut requirements.children;
-    if is_coalesce_partitions(&children[0].plan) {
+    if is_coalesce_partitions(&children[0].plan) && !is_aggregation(plan) {
         // We can safely use the 0th index since we have a `CoalescePartitionsExec`.
         let mut new_child_node = children[0].children.swap_remove(0);
         while new_child_node.plan.output_partitioning() == plan.output_partitioning()

datafusion/physical-optimizer/src/utils.rs

@@ -19,6 +19,7 @@ use std::sync::Arc;
 
 use datafusion_physical_expr::LexRequirement;
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
+use datafusion_physical_plan::aggregates::AggregateExec;
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::repartition::RepartitionExec;
@@ -105,3 +106,8 @@ pub fn is_repartition(plan: &Arc<dyn ExecutionPlan>) -> bool {
 pub fn is_limit(plan: &Arc<dyn ExecutionPlan>) -> bool {
     plan.as_any().is::<GlobalLimitExec>() || plan.as_any().is::<LocalLimitExec>()
 }
+
+// Checks whether the given operator is a [`AggregateExec`].
+pub fn is_aggregation(plan: &Arc<dyn ExecutionPlan>) -> bool {
+    plan.as_any().is::<AggregateExec>()
+}