fix: no longer remove needed coalesce

Author: wiedld

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -18,9 +18,7 @@
 use std::sync::Arc;
 
 use crate::physical_optimizer::enforce_distribution::projection_exec_with_alias;
-use crate::physical_optimizer::sanity_checker::{
-    assert_sanity_check, assert_sanity_check_err,
-};
+use crate::physical_optimizer::sanity_checker::assert_sanity_check;
 use crate::physical_optimizer::test_utils::{
     aggregate_exec, bounded_window_exec, bounded_window_exec_non_set_monotonic,
     bounded_window_exec_with_partition, check_integrity, coalesce_batches_exec,
@@ -2307,40 +2305,27 @@ async fn test_preserve_needed_coalesce() -> Result<()> {
     let plan = sort_exec(sort_key, plan);
 
     // Starting plan: as in our test case.
-    assert_eq!(
-        get_plan_string(&plan),
-        vec![
-            "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]",
-            "  AggregateExec: mode=SinglePartitioned, gby=[a@0 as a1], aggr=[]",
-            "    CoalescePartitionsExec",
-            "      ProjectionExec: expr=[a@0 as a, b@1 as value]",
-            "        UnionExec",
-            "          DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
-            "          DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
-        ],
-    );
+    let starting_plan = vec![
+        "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]",
+        "  AggregateExec: mode=SinglePartitioned, gby=[a@0 as a1], aggr=[]",
+        "    CoalescePartitionsExec",
+        "      ProjectionExec: expr=[a@0 as a, b@1 as value]",
+        "        UnionExec",
+        "          DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
+        "          DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
+    ];
+    assert_eq!(get_plan_string(&plan), starting_plan);
+
     // Test: plan is valid.
     assert_sanity_check(&plan, true);
 
-    // EnforceSorting will remove the coalesce, and add an SPM further up (above the aggregate).
+    // EnforceSorting will not remove the coalesce, as it's required.
     let optimizer = EnforceSorting::new();
     let optimized = optimizer.optimize(plan, &Default::default())?;
-    assert_eq!(
-        get_plan_string(&optimized),
-        vec![
-            "SortPreservingMergeExec: [a@0 ASC]",
-            "  SortExec: expr=[a@0 ASC], preserve_partitioning=[true]",
-            "    AggregateExec: mode=SinglePartitioned, gby=[a@0 as a1], aggr=[]",
-            "      ProjectionExec: expr=[a@0 as a, b@1 as value]",
-            "        UnionExec",
-            "          DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
-            "          DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
-        ],
-    );
+    assert_eq!(get_plan_string(&optimized), starting_plan);
 
-    // Bug: Plan is now invalid.
-    let err = "does not satisfy distribution requirements: HashPartitioned[[a@0]]). Child-0 output partitioning: UnknownPartitioning(2)";
-    assert_sanity_check_err(&optimized, err);
+    // Test: plan is valid.
+    assert_sanity_check(&optimized, true);
 
     Ok(())
 }

datafusion/core/tests/physical_optimizer/sanity_checker.rs

@@ -397,14 +397,6 @@ pub(crate) fn assert_sanity_check(plan: &Arc<dyn ExecutionPlan>, is_sane: bool)
     );
 }
 
-/// Assert reason for sanity check failure.
-pub(crate) fn assert_sanity_check_err(plan: &Arc<dyn ExecutionPlan>, err: &str) {
-    let sanity_checker = SanityCheckPlan::new();
-    let opts = ConfigOptions::default();
-    let error = sanity_checker.optimize(plan.clone(), &opts).unwrap_err();
-    assert!(error.message().contains(err));
-}
-
 /// Check if the plan we created is as expected by comparing the plan
 /// formatted as a string.
 fn assert_plan(plan: &dyn ExecutionPlan, expected_lines: Vec<&str>) {

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_aggregate, is_coalesce_partitions,
+    is_limit, is_repartition, is_sort, is_sort_preserving_merge, is_union, is_window,
 };
 use crate::PhysicalOptimizerRule;
 
@@ -138,27 +138,25 @@ fn is_coalesce_to_remove(
     node: &Arc<dyn ExecutionPlan>,
     parent: &Arc<dyn ExecutionPlan>,
 ) -> bool {
-    node.as_any()
-        .downcast_ref::<CoalescePartitionsExec>()
+    node.as_any().downcast_ref::<CoalescePartitionsExec>()
         .map(|_coalesce| {
             // TODO(wiedld): find a more generalized approach that does not rely on
             // pattern matching the structure of the DAG
             // Note that the `Partitioning::satisfy()` (parent vs. coalesce.child) cannot be used for cases of:
             //      * Repartition -> Coalesce -> Repartition
+            //      * Coalesce -> AggregateExec(input=hash-partitioned)
 
-            let parent_req_single_partition = matches!(
-                parent.required_input_distribution()[0],
-                Distribution::SinglePartition
-            );
+            let parent_req_single_partition = matches!(parent.required_input_distribution()[0], Distribution::SinglePartition)
+                // handle aggregates with input=hashPartitioning with a single output partition
+                || (is_aggregate(parent) && parent.properties().output_partitioning().partition_count() <= 1);
 
             // node above does not require single distribution
             !parent_req_single_partition
             // it doesn't immediately repartition
             || is_repartition(parent)
             // any adjacent Coalesce->Sort can be replaced
             || is_sort(parent)
-        })
-        .unwrap_or(false)
+        }).unwrap_or(false)
 }
 
 fn update_coalesce_ctx_children(

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;
@@ -100,6 +101,11 @@ pub fn is_repartition(plan: &Arc<dyn ExecutionPlan>) -> bool {
     plan.as_any().is::<RepartitionExec>()
 }
 
+/// Checks whether the given operator is an [`AggregateExec`].
+pub fn is_aggregate(plan: &Arc<dyn ExecutionPlan>) -> bool {
+    plan.as_any().is::<AggregateExec>()
+}
+
 /// Checks whether the given operator is a limit;
 /// i.e. either a [`LocalLimitExec`] or a [`GlobalLimitExec`].
 pub fn is_limit(plan: &Arc<dyn ExecutionPlan>) -> bool {