don't combine if input has multiple partitions - preserve distributed aggregation

Author: okhsunrog

datafusion/physical-optimizer/src/combine_partial_final_agg.rs

@@ -24,6 +24,7 @@ use datafusion_common::error::Result;
 use datafusion_physical_plan::aggregates::{
     AggregateExec, AggregateMode, PhysicalGroupBy,
 };
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion_physical_plan::ExecutionPlan;
 
 use crate::PhysicalOptimizerRule;
@@ -85,7 +86,10 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
                         input_agg_exec.aggr_expr(),
                         input_agg_exec.filter_expr(),
                     ),
-                ) {
+                )
+                // Don't combine if input has multiple partitions - preserve distributed aggregation
+                && !has_multi_partition_coalesce(input_agg_exec.input())
+            {
                 let mode = if agg_exec.mode() == &AggregateMode::Final {
                     AggregateMode::Single
                 } else {
@@ -161,4 +165,28 @@ fn can_combine(final_agg: GroupExprsRef, partial_agg: GroupExprsRef) -> bool {
         )
 }
 
+/// Check if the plan subtree contains a CoalescePartitionsExec with multiple input partitions.
+fn has_multi_partition_coalesce(plan: &Arc<dyn ExecutionPlan>) -> bool {
+    // Check if this node is CoalescePartitionsExec with multiple inputs
+    if plan.as_any().is::<CoalescePartitionsExec>() {
+        let partition_count = plan
+            .children()
+            .first()
+            .map(|child| child.properties().partitioning.partition_count())
+            .unwrap_or(0);
+        if partition_count > 1 {
+            return true;
+        }
+    }
+
+    // Recursively check all children
+    for child in plan.children() {
+        if has_multi_partition_coalesce(child) {
+            return true;
+        }
+    }
+
+    false
+}
+
 // See tests in datafusion/core/tests/physical_optimizer

datafusion/sql/src/parser.rs

@@ -279,15 +279,14 @@ impl<'a> DFParser<'a> {
         sql: &str,
         dialect: &'a dyn Dialect,
     ) -> Result<Self, ParserError> {
-        let tokens = Tokenizer::new(dialect, sql).into_tokens().collect::<Result<_, _>>()?;
+        let tokens = Tokenizer::new(dialect, sql)
+            .into_tokens()
+            .collect::<Result<_, _>>()?;
         Ok(Self::from_dialect_and_tokens(dialect, tokens))
     }
 
     /// Create a new parser from specified dialect and tokens.
-    pub fn from_dialect_and_tokens(
-        dialect: &'a dyn Dialect,
-        tokens: Vec<Token>,
-    ) -> Self {
+    pub fn from_dialect_and_tokens(dialect: &'a dyn Dialect, tokens: Vec<Token>) -> Self {
         let parser = Parser::new(dialect).with_tokens(tokens);
         DFParser { parser }
     }