feat: support PartialMerge

Author: comphead

native/core/src/execution/planner.rs

@@ -967,10 +967,17 @@ impl PhysicalPlanner {
                 let group_by = PhysicalGroupBy::new_single(group_exprs?);
                 let schema = child.schema();
 
-                let mode = if agg.mode == 0 {
-                    DFAggregateMode::Partial
-                } else {
-                    DFAggregateMode::Final
+                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,
+                    other => {
+                        return Err(ExecutionError::GeneralError(format!(
+                            "Unsupported aggregate mode: {other}"
+                        )))
+                    }
                 };
 
                 let agg_exprs: PhyAggResult = agg

native/proto/src/proto/operator.proto

@@ -319,6 +319,7 @@ message ParquetWriter {
 enum AggregateMode {
   Partial = 0;
   Final = 1;
+  PartialMerge = 2;
 }
 
 message Expand {

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -461,15 +461,15 @@ object QueryPlanSerde extends Logging with CometExprShim {
       binding: Boolean,
       conf: SQLConf): Option[AggExpr] = {
 
-    // Support Count(distinct single_value)
-    // COUNT(DISTINCT x) - supported
-    // COUNT(DISTINCT x, x) - supported through transition to COUNT(DISTINCT x)
-    // COUNT(DISTINCT x, y) - not supported
+    // Distinct aggregates with a single column are supported (e.g., COUNT(DISTINCT x),
+    // SUM(DISTINCT x), AVG(DISTINCT x)). The multi-stage plan generated by Spark
+    // guarantees distinct semantics through grouping — the native side does not need
+    // to handle deduplication.
+    // Multi-column distinct is only supported for COUNT (e.g., COUNT(DISTINCT x, y)).
     if (aggExpr.isDistinct
-      &&
-      !(aggExpr.aggregateFunction.prettyName == "count" &&
-        aggExpr.aggregateFunction.children.length == 1)) {
-      withInfo(aggExpr, s"Distinct aggregate not supported for: $aggExpr")
+      && aggExpr.aggregateFunction.children.length > 1
+      && aggExpr.aggregateFunction.prettyName != "count") {
+      withInfo(aggExpr, s"Multi-column distinct aggregate not supported for: $aggExpr")
       return None
     }
 

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

@@ -1441,16 +1441,18 @@ trait CometBaseAggregate {
       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
       }
 
-      // In final 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.Final
+      // 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)
       val output = child.output
 
@@ -1496,14 +1498,21 @@ trait CometBaseAggregate {
 
   /**
    * Find the first Comet partial aggregate in the plan. If it reaches a Spark HashAggregate with
-   * partial mode, it will return None.
+   * partial or partial-merge mode, it will return None.
    */
   private def findCometPartialAgg(plan: SparkPlan): Option[CometHashAggregateExec] = {
+    def isPartialOrMerge(mode: AggregateMode): Boolean =
+      mode == Partial || mode == PartialMerge
+
     plan.collectFirst {
-      case agg: CometHashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) =>
+      case agg: CometHashAggregateExec
+          if agg.aggregateExpressions.forall(e => isPartialOrMerge(e.mode)) =>
         Some(agg)
-      case agg: HashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) => None
-      case agg: ObjectHashAggregateExec if agg.aggregateExpressions.forall(_.mode == Partial) =>
+      case agg: HashAggregateExec
+          if agg.aggregateExpressions.forall(e => isPartialOrMerge(e.mode)) =>
+        None
+      case agg: ObjectHashAggregateExec
+          if agg.aggregateExpressions.forall(e => isPartialOrMerge(e.mode)) =>
         None
       case a: AQEShuffleReadExec => findCometPartialAgg(a.child)
       case s: ShuffleQueryStageExec => findCometPartialAgg(s.plan)
@@ -1642,11 +1651,8 @@ case class CometHashAggregateExec(
 
   // The aggExprs could be empty. For example, if the aggregate functions only have
   // distinct aggregate functions or only have group by, the aggExprs is empty and
-  // modes is empty too. If aggExprs is not empty, we need to verify all the
-  // aggregates have the same mode.
+  // modes is empty too.
   val modes: Seq[AggregateMode] = aggregateExpressions.map(_.mode).distinct
-  assert(modes.length == 1 || modes.isEmpty)
-  val mode = modes.headOption
 
   override def producedAttributes: AttributeSet = outputSet ++ AttributeSet(resultExpressions)
 
@@ -1663,7 +1669,7 @@ case class CometHashAggregateExec(
   }
 
   override def stringArgs: Iterator[Any] =
-    Iterator(input, mode, groupingExpressions, aggregateExpressions, child)
+    Iterator(input, modes, groupingExpressions, aggregateExpressions, child)
 
   override def equals(obj: Any): Boolean = {
     obj match {
@@ -1672,7 +1678,7 @@ case class CometHashAggregateExec(
         this.groupingExpressions == other.groupingExpressions &&
         this.aggregateExpressions == other.aggregateExpressions &&
         this.input == other.input &&
-        this.mode == other.mode &&
+        this.modes == other.modes &&
         this.child == other.child &&
         this.serializedPlanOpt == other.serializedPlanOpt
       case _ =>
@@ -1681,7 +1687,7 @@ case class CometHashAggregateExec(
   }
 
   override def hashCode(): Int =
-    Objects.hashCode(output, groupingExpressions, aggregateExpressions, input, mode, child)
+    Objects.hashCode(output, groupingExpressions, aggregateExpressions, input, modes, child)
 
   override protected def outputExpressions: Seq[NamedExpression] = resultExpressions
 }

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

@@ -641,7 +641,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
                 checkSparkAnswerAndFallbackReason(
                   "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",
-                  "Unsupported aggregation mode PartialMerge")
+                  "All aggregate expressions do not have the same mode")
               }
             }
           }
@@ -650,6 +650,56 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("partialMerge - cnt distinct + sum") {
+    withTempDir(dir => {
+      withSQLConf("spark.comet.enabled" -> "false") {
+        sql("""
+              CREATE OR REPLACE TEMP VIEW t (v, v1, i) AS
+              VALUES
+                ('c',  'a',  1),
+                ('c1', 'a1', 1),
+                ('c2', 'a2', 2),
+                ('c3', 'a3', 2),
+                ('c4', 'a4', 2),
+                ('c',  'a',  1),
+                ('c1', 'a1', 1),
+                ('c2', 'a2', 2),
+                ('c3', 'a3', 2),
+                ('c4', 'a4', 2),
+                ('c',  'a',  1),
+                ('c1', 'a1', 1),
+                ('c2', 'a2', 2),
+                ('c3', 'a3', 2),
+                ('c4', 'a4', 2)
+              """)
+        sql("select * from t")
+          .repartition(3)
+          .write
+          .mode("overwrite")
+          .parquet(dir.getAbsolutePath)
+      }
+
+      withSQLConf(
+        CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+        "spark.comet.exec.shuffle.fallbackToColumnar" -> "false",
+        "spark.comet.cast.allowIncompatible" -> "true",
+        "spark.sql.adaptive.enabled" -> "false",
+        "spark.comet.explain.native.enabled" -> "true",
+        "spark.comet.enabled" -> "true",
+        "spark.comet.expression.Cast.allowIncompatible" -> "true",
+        "spark.comet.exec.shuffle.enableFastEncoding" -> "true",
+        "spark.comet.exec.shuffle.enabled" -> "true",
+        "spark.comet.explainFallback.enabled" -> "true",
+        CometConf.COMET_NATIVE_SCAN_IMPL.key -> "native_iceberg_compat",
+        "spark.shuffle.manager" ->
+          "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager",
+        "spark.comet.logFallbackReasons.enabled" -> "true") {
+        spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t2")
+        checkSparkAnswerAndOperator("SELECT i, sum(v1), count(distinct v) FROM t2 group by i")
+      }
+    })
+  }
+
   test("multiple group-by columns + single aggregate column (first/last), with nulls") {
     val numValues = 10000
 

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

@@ -484,9 +484,12 @@ class CometExecSuite extends CometTestBase {
         case s: CometHashAggregateExec => s
       }.get
 
-      assert(agg.mode.isDefined && agg.mode.get.isInstanceOf[AggregateMode])
+      assert(
+        agg.modes.nonEmpty && agg.modes.headOption.get.isInstanceOf[AggregateMode])
       val newAgg = agg.cleanBlock().asInstanceOf[CometHashAggregateExec]
-      assert(newAgg.mode.isDefined && newAgg.mode.get.isInstanceOf[AggregateMode])
+      assert(
+        newAgg.modes.nonEmpty &&
+          newAgg.modes.headOption.get.isInstanceOf[AggregateMode])
     }
   }