fix: revert unsafe partial aggregates after final fallback

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -616,7 +616,7 @@ case class CometExecRule(session: SparkSession)
       // during the bottom-up conversion. Tags persist through AQE stage creation.
       tagUnsafePartialAggregates(planWithJoinRewritten)
 
-      var newPlan = transform(planWithJoinRewritten)
+      var newPlan = revertUnsafePartialAggregates(transform(planWithJoinRewritten))
 
       // if the plan cannot be run fully natively then explain why (when appropriate
       // config is enabled)
@@ -974,7 +974,7 @@ case class CometExecRule(session: SparkSession)
         // PartialMerge stages of a distinct-aggregate rewrite. See issues #1389 and #4813.
         val modes = agg.aggregateExpressions.map(_.mode).distinct
         if (modes == Seq(Final) &&
-          !QueryPlanSerde.allAggsSupportMixedExecution(agg.aggregateExpressions) &&
+          !QueryPlanSerde.allAggsSupportNativePartialToSparkFinal(agg.aggregateExpressions) &&
           !canAggregateBeConverted(agg, Final)) {
           findPartialAggInPlan(agg.child).foreach { partial =>
             // Only tag if the Partial would otherwise have been converted. If the Partial itself
@@ -1016,6 +1016,63 @@ case class CometExecRule(session: SparkSession)
     }
   }
 
+  /**
+   * The early tagging pass cannot know whether a Final's child will become native. Check the
+   * actual conversion result as well, before native blocks are serialized or AQE launches stages.
+   * Restore only the feeding aggregate/exchange chain; keep native work below its Partial.
+   */
+  private def revertUnsafePartialAggregates(plan: SparkPlan): SparkPlan = {
+    def revertChain(node: SparkPlan): Option[SparkPlan] = node match {
+      case agg: CometHashAggregateExec if agg.modes == Seq(Partial) =>
+        val partial = agg.originalPlan.withNewChildren(Seq(agg.child))
+        partial.setTagValue(
+          CometExecRule.COMET_UNSAFE_PARTIAL,
+          "Partial aggregate disabled: corresponding final aggregate " +
+            "cannot be converted to Comet and intermediate buffer formats are incompatible")
+        Some(partial)
+
+      case agg: CometHashAggregateExec
+          if agg.modes.forall(m => m == Partial || m == PartialMerge) =>
+        revertChain(agg.child).map(child => agg.originalPlan.withNewChildren(Seq(child)))
+
+      case agg: BaseAggregateExec
+          if agg.aggregateExpressions.nonEmpty &&
+            agg.aggregateExpressions.forall(_.mode == Partial) =>
+        // This producer already emits Spark buffers. Do not reach through it to an unrelated
+        // aggregate below it.
+        None
+
+      case agg: BaseAggregateExec
+          if agg.aggregateExpressions.forall(e => e.mode == Partial || e.mode == PartialMerge) =>
+        revertChain(agg.child).map(child => agg.withNewChildren(Seq(child)))
+
+      case CometSinkPlaceHolder(_, _, shuffle: CometShuffleExchangeExec) =>
+        revertChain(shuffle)
+      case shuffle: CometShuffleExchangeExec =>
+        revertChain(shuffle.child).map(child => shuffle.originalPlan.withNewChildren(Seq(child)))
+      case shuffle: ShuffleExchangeExec =>
+        revertChain(shuffle.child).map(child => shuffle.withNewChildren(Seq(child)))
+
+      case _: ShuffleQueryStageExec | _: ReusedExchangeExec =>
+        // A stage owns (and may already have materialized) its buffers. Never rewrite it here.
+        // The whole-plan QueryStagePrep pass must tag the Partial before stages are created;
+        // that tag keeps it in Spark when the rule is reapplied to the exchange in isolation.
+        None
+      case _ => None
+    }
+
+    plan.transformUp {
+      case agg: BaseAggregateExec
+          if agg.aggregateExpressions.map(_.mode).distinct == Seq(Final) &&
+            !QueryPlanSerde.allAggsSupportNativePartialToSparkFinal(agg.aggregateExpressions) =>
+        revertChain(agg.child)
+          // Rebuild native consumers and shuffles from their original Spark operators. Merely
+          // replacing their children would leave a native protobuf reading the old buffers.
+          .map(child => transform(agg.withNewChildren(Seq(child))))
+          .getOrElse(agg)
+    }
+  }
+
   /**
    * Look for the bottom Partial-mode aggregate that feeds into the given plan (the child of a
    * Final). Walks through exchanges and AQE stages, and continues down through intermediate
@@ -1045,8 +1102,8 @@ case class CometExecRule(session: SparkSession)
   /**
    * Conservative check for whether an aggregate could be converted to Comet. Checks operator
    * enablement, grouping expressions, aggregate expressions, and result expressions.
-   * Intentionally skips the sparkFinalMode / child-native checks since those depend on
-   * transformation state.
+   * Intentionally skips the child-native checks since those depend on transformation state;
+   * [[revertUnsafePartialAggregates]] checks the actual conversion result before execution.
    *
    * WARNING: this intentionally mirrors the predicate checks in `CometBaseAggregate.doConvert`
    * (operators.scala). Any change to the convertibility rules there must be reflected here or

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

@@ -82,16 +82,23 @@ trait CometAggregateExpressionSerde[T <: AggregateFunction] {
   def getSupportLevel(expr: T): SupportLevel = Compatible(None)
 
   /**
-   * Whether this aggregate's intermediate buffer format is compatible between Spark and Comet for
-   * the given function instance, making it safe to run the Partial in one engine and the Final in
-   * the other. Aggregates with simple single-value buffers (MIN, MAX, bitwise) are always safe;
-   * SUM and non-decimal AVG match Spark's buffer and are safe except where noted per instance
-   * (e.g. TRY-mode SUM uses a Comet-internal flag column). COUNT is intentionally excluded
-   * despite a matching buffer: mixed COUNT partial/final regressed AQE's
+   * Whether a Comet aggregate can consume this function's Spark intermediate buffer. This covers
+   * Spark Partial to Comet Final, including intermediate PartialMerge stages. COUNT is excluded
+   * despite a matching buffer: a Comet Final above a Spark Partial regressed AQE's
    * PropagateEmptyRelationAfterAQE pattern (which matches BaseAggregateExec only) and the Spark
    * 4.0 count-bug decorrelation for correlated IN subqueries.
    */
-  def supportsMixedPartialFinal(fn: T): Boolean = false
+  def supportsSparkPartialToNativeFinal(fn: T): Boolean = false
+
+  /**
+   * Whether Spark can consume this function's Comet intermediate buffer. Keep this separate from
+   * the reverse direction: planner restrictions on a Comet Final do not necessarily prohibit a
+   * Comet Partial feeding a Spark Final. Existing bidirectional implementations share the
+   * default; handlers may admit an additional forward direction only after validating it
+   * independently.
+   */
+  def supportsNativePartialToSparkFinal(fn: T): Boolean =
+    supportsSparkPartialToNativeFinal(fn)
 
   /**
    * Convert a Spark expression into a protocol buffer representation that can be passed into

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

@@ -421,31 +421,35 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     classOf[VarianceSamp] -> CometVarianceSamp)
 
   /**
-   * Returns true if all aggregate expressions in the list have intermediate buffer formats that
-   * are compatible between Spark and Comet, making it safe to run Partial in one engine and Final
-   * in the other.
+   * Returns true if Spark can consume all the intermediate buffers produced by Comet. Used when a
+   * Spark Final would otherwise consume a native Partial, including after shuffle fallback.
    */
-  def allAggsSupportMixedExecution(aggExprs: Seq[AggregateExpression]): Boolean = {
-    aggExprs.forall(aggExpr => supportsMixedExecution(aggExpr.aggregateFunction))
+  def allAggsSupportNativePartialToSparkFinal(aggExprs: Seq[AggregateExpression]): Boolean = {
+    aggExprs.forall { aggExpr =>
+      val fn = aggExpr.aggregateFunction
+      aggrSerdeMap.get(fn.getClass).exists { handler =>
+        handler
+          .asInstanceOf[CometAggregateExpressionSerde[AggregateFunction]]
+          .supportsNativePartialToSparkFinal(fn)
+      }
+    }
   }
 
   /**
-   * Returns the aggregate functions in the list whose intermediate buffer formats are not known
-   * to be compatible between Spark and Comet. These are the functions that prevent a Spark Final
-   * aggregate (without a Comet Partial) from running, since the buffer produced by one engine
-   * cannot be safely consumed by the other.
+   * Returns functions whose Spark intermediate buffers cannot safely be consumed by a Comet Final
+   * or PartialMerge. This is independent of native Partial to Spark Final compatibility.
    */
-  def aggsNotSupportingMixedExecution(
+  def aggsNotSupportingSparkPartialToNativeFinal(
       aggExprs: Seq[AggregateExpression]): Seq[AggregateFunction] = {
-    aggExprs.map(_.aggregateFunction).filterNot(supportsMixedExecution)
+    aggExprs.map(_.aggregateFunction).filterNot(supportsSparkPartialToNativeFinal)
   }
 
-  private def supportsMixedExecution(fn: AggregateFunction): Boolean = {
+  private def supportsSparkPartialToNativeFinal(fn: AggregateFunction): Boolean = {
     aggrSerdeMap.get(fn.getClass) match {
       case Some(handler) =>
         handler
           .asInstanceOf[CometAggregateExpressionSerde[AggregateFunction]]
-          .supportsMixedPartialFinal(fn)
+          .supportsSparkPartialToNativeFinal(fn)
       case None => false
     }
   }

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

@@ -35,7 +35,7 @@ import org.apache.comet.shims.{CometCollectShim, CometEvalModeUtil}
 
 object CometMin extends CometAggregateExpressionSerde[Min] {
 
-  override def supportsMixedPartialFinal(fn: Min): Boolean = true
+  override def supportsSparkPartialToNativeFinal(fn: Min): Boolean = true
 
   override def getSupportLevel(expr: Min): SupportLevel =
     AggSerde.minMaxSupportLevel(expr.dataType)
@@ -71,7 +71,7 @@ object CometMin extends CometAggregateExpressionSerde[Min] {
 
 object CometMax extends CometAggregateExpressionSerde[Max] {
 
-  override def supportsMixedPartialFinal(fn: Max): Boolean = true
+  override def supportsSparkPartialToNativeFinal(fn: Max): Boolean = true
 
   override def getSupportLevel(expr: Max): SupportLevel =
     AggSerde.minMaxSupportLevel(expr.dataType)
@@ -106,6 +106,10 @@ object CometMax extends CometAggregateExpressionSerde[Max] {
 }
 
 object CometCount extends CometAggregateExpressionSerde[Count] {
+  // Both buffers are a single non-null Long. The AQE/count-bug restrictions documented on the
+  // reverse direction concern a Comet Final; retaining Spark's Final preserves those rewrites.
+  override def supportsNativePartialToSparkFinal(fn: Count): Boolean = true
+
   override def convert(
       aggExpr: AggregateExpression,
       expr: Count,
@@ -129,7 +133,12 @@ object CometCount extends CometAggregateExpressionSerde[Count] {
 
 object CometAverage extends CometAggregateExpressionSerde[Average] {
 
-  override def supportsMixedPartialFinal(fn: Average): Boolean =
+  // A never-updated non-decimal native AVG partial emits (null, 0), but Spark's merge needs
+  // (0.0, 0). Keep this direction disabled until #5420 repairs the emitted state, including
+  // aggregate nodes that also contain a buffer-compatible function such as COUNT.
+  override def supportsNativePartialToSparkFinal(fn: Average): Boolean = false
+
+  override def supportsSparkPartialToNativeFinal(fn: Average): Boolean =
     // Non-decimal AVG has a (sum: double, count: long) buffer matching Spark. Decimal AVG is
     // deferred (overflow nulls count differently) and stays unsafe for mixed execution.
     !fn.child.dataType.isInstanceOf[DecimalType]
@@ -189,7 +198,7 @@ object CometAverage extends CometAggregateExpressionSerde[Average] {
 
 object CometSum extends CometAggregateExpressionSerde[Sum] {
 
-  override def supportsMixedPartialFinal(fn: Sum): Boolean =
+  override def supportsSparkPartialToNativeFinal(fn: Sum): Boolean =
     // Decimal SUM is excluded: overflow detection (ANSI throw / Legacy null) does not survive a
     // Spark-partial / Comet-final split, so the required ArithmeticException is never raised.
     // TRY-mode integer SUM carries a Comet-internal has_all_nulls column that Spark cannot read.
@@ -306,7 +315,7 @@ object CometLast extends CometAggregateExpressionSerde[Last] {
 }
 
 object CometBitAndAgg extends CometAggregateExpressionSerde[BitAndAgg] {
-  override def supportsMixedPartialFinal(fn: BitAndAgg): Boolean = true
+  override def supportsSparkPartialToNativeFinal(fn: BitAndAgg): Boolean = true
 
   override def getSupportLevel(expr: BitAndAgg): SupportLevel =
     if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) {
@@ -344,7 +353,7 @@ object CometBitAndAgg extends CometAggregateExpressionSerde[BitAndAgg] {
 }
 
 object CometBitOrAgg extends CometAggregateExpressionSerde[BitOrAgg] {
-  override def supportsMixedPartialFinal(fn: BitOrAgg): Boolean = true
+  override def supportsSparkPartialToNativeFinal(fn: BitOrAgg): Boolean = true
 
   override def getSupportLevel(expr: BitOrAgg): SupportLevel =
     if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) {
@@ -382,7 +391,7 @@ object CometBitOrAgg extends CometAggregateExpressionSerde[BitOrAgg] {
 }
 
 object CometBitXOrAgg extends CometAggregateExpressionSerde[BitXorAgg] {
-  override def supportsMixedPartialFinal(fn: BitXorAgg): Boolean = true
+  override def supportsSparkPartialToNativeFinal(fn: BitXorAgg): Boolean = true
 
   override def getSupportLevel(expr: BitXorAgg): SupportLevel =
     if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) {
@@ -765,7 +774,7 @@ object CometCorr extends CometAggregateExpressionSerde[Corr] {
 
 object CometBloomFilterAggregate extends CometAggregateExpressionSerde[BloomFilterAggregate] {
 
-  override def supportsMixedPartialFinal(fn: BloomFilterAggregate): Boolean = true
+  override def supportsSparkPartialToNativeFinal(fn: BloomFilterAggregate): Boolean = true
 
   override def getSupportLevel(expr: BloomFilterAggregate): SupportLevel =
     expr.child.dataType match {
@@ -934,7 +943,7 @@ object CometApproxCountDistinct extends CometAggregateExpressionSerde[HyperLogLo
   // The register buffer uses Spark's identical packed-`Long` layout (`numWords` `Long` columns),
   // matching Spark's `aggBufferSchema`, so a Comet partial and Spark final (or the reverse) can
   // be mixed in one plan.
-  override def supportsMixedPartialFinal(fn: HyperLogLogPlusPlus): Boolean = true
+  override def supportsSparkPartialToNativeFinal(fn: HyperLogLogPlusPlus): Boolean = true
 
   // Types that Comet's native `xxhash64` hashes identically to Spark's `XxHash64Function`.
   // `StringType` here is the default UTF8_BINARY collation; a collated `StringType(collationId)`

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

@@ -1656,7 +1656,7 @@ trait CometBaseAggregate {
 
     if (missingCometProducer) {
       val incompatibleAggs =
-        QueryPlanSerde.aggsNotSupportingMixedExecution(aggregate.aggregateExpressions)
+        QueryPlanSerde.aggsNotSupportingSparkPartialToNativeFinal(aggregate.aggregateExpressions)
       if (incompatibleAggs.nonEmpty) {
         val names = incompatibleAggs.map(_.prettyName).distinct.sorted.mkString(", ")
         withFallbackReason(