feat: support PartialMerge

Author: comphead

native/core/src/execution/mod.rs

@@ -19,6 +19,7 @@
 pub mod columnar_to_row;
 pub mod expressions;
 pub mod jni_api;
+pub(crate) mod merge_as_partial;
 pub(crate) mod metrics;
 pub mod operators;
 pub(crate) mod planner;

native/core/src/execution/planner.rs

@@ -968,24 +968,87 @@ impl PhysicalPlanner {
 
                 let mode = match agg.mode {
                     0 => DFAggregateMode::Partial,
-                    // Both Final and PartialMerge use merge semantics in DataFusion.
-                    // The output difference (final values vs intermediate buffers) is
-                    // handled by the presence/absence of result_exprs.
-                    1 | 2 => DFAggregateMode::Final,
+                    1 => DFAggregateMode::Final,
+                    2 => DFAggregateMode::Partial, // PartialMerge uses Partial + MergeAsPartial
                     other => {
                         return Err(ExecutionError::GeneralError(format!(
                             "Unsupported aggregate mode: {other}"
                         )))
                     }
                 };
 
+                // Determine per-expression modes. PartialMerge (2) expressions use
+                // MergeAsPartial wrapper so they run merge semantics in Partial mode.
+                let per_expr_modes: Vec<i32> = if !agg.expr_modes.is_empty() {
+                    agg.expr_modes.clone()
+                } else {
+                    vec![agg.mode; agg.agg_exprs.len()]
+                };
+
+                let has_partial_merge = per_expr_modes.contains(&2);
+
                 let agg_exprs: PhyAggResult = agg
                     .agg_exprs
                     .iter()
                     .map(|expr| self.create_agg_expr(expr, Arc::clone(&schema)))
                     .collect();
 
-                let aggr_expr = agg_exprs?.into_iter().map(Arc::new).collect();
+                let aggr_expr: Vec<Arc<AggregateFunctionExpr>> = if has_partial_merge {
+                    // Wrap PartialMerge expressions with MergeAsPartial.
+                    // State fields in the child's output start at initial_input_buffer_offset.
+                    let mut state_offset = agg.initial_input_buffer_offset as usize;
+                    let child_schema = child.schema();
+
+                    agg_exprs?
+                        .into_iter()
+                        .enumerate()
+                        .map(|(idx, expr)| {
+                            let expr_mode = per_expr_modes[idx];
+                            if expr_mode == 2 {
+                                // PartialMerge: wrap with MergeAsPartial
+                                let state_fields = expr.state_fields().map_err(|e| {
+                                    ExecutionError::GeneralError(e.to_string())
+                                })?;
+                                let num_state_fields = state_fields.len();
+
+                                // Create Column refs pointing to state field positions
+                                let state_cols: Vec<Arc<dyn PhysicalExpr>> = (0..num_state_fields)
+                                    .map(|i| {
+                                        let col_idx = state_offset + i;
+                                        let field = child_schema.field(col_idx);
+                                        Arc::new(Column::new(field.name(), col_idx))
+                                            as Arc<dyn PhysicalExpr>
+                                    })
+                                    .collect();
+                                state_offset += num_state_fields;
+
+                                let merge_udf = crate::execution::merge_as_partial::MergeAsPartialUDF::new(&expr)
+                                    .map_err(|e| ExecutionError::DataFusionError(e.to_string()))?;
+                                let merge_udf_arc = Arc::new(
+                                    datafusion::logical_expr::AggregateUDF::new_from_impl(merge_udf),
+                                );
+
+                                let merge_expr = AggregateExprBuilder::new(
+                                    merge_udf_arc,
+                                    state_cols,
+                                )
+                                .schema(Arc::clone(&child_schema))
+                                .alias(format!("col_{idx}"))
+                                .with_ignore_nulls(false)
+                                .with_distinct(false)
+                                .build()
+                                .map_err(|e| ExecutionError::DataFusionError(e.to_string()))?;
+
+                                Ok(Arc::new(merge_expr))
+                            } else {
+                                // Partial: use as-is
+                                Ok(Arc::new(expr))
+                            }
+                        })
+                        .collect::<Result<Vec<_>, ExecutionError>>()?
+                } else {
+                    agg_exprs?.into_iter().map(Arc::new).collect()
+                };
 
                 // Build per-aggregate filter expressions from the FILTER (WHERE ...) clause.
                 // Filters are only present in Partial mode; Final/PartialMerge always get None.

native/proto/src/proto/operator.proto

@@ -270,6 +270,13 @@ message HashAggregate {
   repeated spark.spark_expression.AggExpr agg_exprs = 2;
   repeated spark.spark_expression.Expr result_exprs = 3;
   AggregateMode mode = 5;
+  // Per-expression modes for mixed-mode aggregates (e.g., PartialMerge + Partial).
+  // When set, each entry corresponds to agg_exprs at the same index.
+  // When empty, all expressions use the `mode` field.
+  repeated AggregateMode expr_modes = 6;
+  // Offset in the child's output where aggregate buffer attributes start.
+  // Used by PartialMerge to locate state fields in the input.
+  int32 initial_input_buffer_offset = 7;
 }
 
 message Limit {

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

@@ -1356,8 +1356,9 @@ trait CometBaseAggregate {
       childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = {
 
     val modes = aggregate.aggregateExpressions.map(_.mode).distinct
-    // In distinct aggregates there can be a combination of modes
-    val multiMode = modes.size > 1
+    // In distinct aggregates there can be a combination of modes.
+    // We support {Partial, PartialMerge} mix; other combinations are rejected.
+    val multiMode = modes.size > 1 && modes.toSet != Set(Partial, PartialMerge)
     // For a final mode HashAggregate, we only need to transform the HashAggregate
     // if there is Comet partial aggregation.
     val sparkFinalMode = modes.contains(Final) && findCometPartialAgg(aggregate.child).isEmpty
@@ -1430,34 +1431,42 @@ trait CometBaseAggregate {
       Some(builder.setHashAgg(hashAggBuilder).build())
     } else {
       val modes = aggregateExpressions.map(_.mode).distinct
-
-      if (modes.size != 1) {
-        // This shouldn't happen as all aggregation expressions should share the same mode.
-        // Fallback to Spark nevertheless here.
-        withInfo(aggregate, "All aggregate expressions do not have the same mode")
+      val modeSet = modes.toSet
+
+      // Validate mode combinations. We support:
+      // - All Partial
+      // - All Final
+      // - All PartialMerge
+      // - Mixed {Partial, PartialMerge} (for distinct aggregate plans)
+      val isMixedPartialMerge = modeSet == Set(Partial, PartialMerge)
+      if (modes.size > 1 && !isMixedPartialMerge) {
+        withInfo(aggregate, s"Unsupported mixed aggregation modes: ${modes.mkString(", ")}")
         return None
       }
 
-      val mode = modes.head match {
-        case Partial => CometAggregateMode.Partial
-        case Final => CometAggregateMode.Final
-        case PartialMerge => CometAggregateMode.PartialMerge
-        case _ =>
-          withInfo(aggregate, s"Unsupported aggregation mode ${modes.head}")
-          return None
+      // Determine the proto mode. For uniform modes, use that mode directly.
+      // For mixed {Partial, PartialMerge}, use Partial as the base mode since
+      // PartialMerge expressions are wrapped with MergeAsPartial on the native side.
+      val mode = if (isMixedPartialMerge) {
+        CometAggregateMode.Partial
+      } else {
+        modes.head match {
+          case Partial => CometAggregateMode.Partial
+          case Final => CometAggregateMode.Final
+          case PartialMerge => CometAggregateMode.PartialMerge
+          case _ =>
+            withInfo(aggregate, s"Unsupported aggregation mode ${modes.head}")
+            return None
+        }
       }
 
-      // In final/partial-merge mode, the aggregate expressions are bound to the
-      // output of the child and partial aggregate expressions buffer attributes
-      // produced by partial aggregation. This is done in Spark `HashAggregateExec`
-      // internally. In Comet, we don't have to do this because we don't use the
-      // merging expression.
-      val binding = mode == CometAggregateMode.Partial
-      // `output` is only used when `binding` is true (i.e., non-Final)
+      // Per-expression binding: Partial expressions bind to child output,
+      // PartialMerge/Final expressions do not (native planner handles their input).
       val output = child.output
-
-      val aggExprs =
-        aggregateExpressions.map(aggExprToProto(_, output, binding, aggregate.conf))
+      val aggExprs = aggregateExpressions.map { a =>
+        val exprBinding = a.mode != PartialMerge && a.mode != Final
+        aggExprToProto(a, output, exprBinding, aggregate.conf)
+      }
 
       if (aggExprs.exists(_.isEmpty)) {
         withInfo(
@@ -1485,6 +1494,24 @@ trait CometBaseAggregate {
           hashAggBuilder.addAllResultExprs(resultExprs.map(_.get).asJava)
         }
         hashAggBuilder.setModeValue(mode.getNumber)
+
+        // Send per-expression modes and buffer offset for PartialMerge handling
+        val hasPartialMerge = aggregateExpressions.exists(_.mode == PartialMerge)
+        if (hasPartialMerge) {
+          val exprModes = aggregateExpressions.map { a =>
+            a.mode match {
+              case Partial => CometAggregateMode.Partial
+              case PartialMerge => CometAggregateMode.PartialMerge
+              case Final => CometAggregateMode.Final
+              case other =>
+                withInfo(aggregate, s"Unsupported aggregation mode $other")
+                return None
+            }
+          }
+          hashAggBuilder.addAllExprModes(exprModes.asJava)
+          hashAggBuilder.setInitialInputBufferOffset(aggregate.initialInputBufferOffset)
+        }
+
         Some(builder.setHashAgg(hashAggBuilder).build())
       } else {
         val allChildren: Seq[Expression] =

spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala

@@ -638,10 +638,9 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
               dictionaryEnabled) {
               withView("v") {
                 sql("CREATE TEMP VIEW v AS SELECT _1, _2 FROM tbl ORDER BY _1")
-                checkSparkAnswerAndFallbackReason(
+                checkSparkAnswerAndOperator(
                   "SELECT _2, SUM(_1), SUM(DISTINCT _1), MIN(_1), MAX(_1), COUNT(_1)," +
-                    " COUNT(DISTINCT _1), AVG(_1), FIRST(_1), LAST(_1) FROM v GROUP BY _2",
-                  "All aggregate expressions do not have the same mode")
+                    " COUNT(DISTINCT _1), AVG(_1), FIRST(_1), LAST(_1) FROM v GROUP BY _2")
               }
             }
           }