fix: preserve aggregate result identity during exchange reuse

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

@@ -30,7 +30,7 @@ import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSet, Expression, ExpressionSet, Generator, NamedExpression, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSeq, AttributeSet, Expression, ExpressionSet, Generator, NamedExpression, SortOrder}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Final, First, Last, Partial, PartialMerge, Percentile}
 import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
 import org.apache.spark.sql.catalyst.plans._
@@ -1993,6 +1993,7 @@ object CometHashAggregateExec
       op.output,
       op.groupingExpressions,
       op.aggregateExpressions,
+      op.aggregateAttributes,
       op.resultExpressions,
       op.child.output,
       op.child,
@@ -2043,6 +2044,7 @@ object CometObjectHashAggregateExec
       adjustOutputForNativeState(op),
       op.groupingExpressions,
       op.aggregateExpressions,
+      op.aggregateAttributes,
       op.resultExpressions,
       op.child.output,
       op.child,
@@ -2056,6 +2058,7 @@ case class CometHashAggregateExec(
     override val output: Seq[Attribute],
     groupingExpressions: Seq[NamedExpression],
     aggregateExpressions: Seq[AggregateExpression],
+    aggregateAttributes: Seq[Attribute],
     resultExpressions: Seq[NamedExpression],
     input: Seq[Attribute],
     child: SparkPlan,
@@ -2068,7 +2071,15 @@ case class CometHashAggregateExec(
   // modes is empty too.
   val modes: Seq[AggregateMode] = aggregateExpressions.map(_.mode).distinct
 
-  override def producedAttributes: AttributeSet = outputSet ++ AttributeSet(resultExpressions)
+  // Match Spark's aggregate canonicalization, including the original result attributes that
+  // rewritten DISTINCT aggregate expressions do not necessarily retain in their resultIds.
+  override lazy val allAttributes: AttributeSeq =
+    child.output ++ aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes) ++
+      aggregateAttributes ++
+      aggregateExpressions.flatMap(_.aggregateFunction.inputAggBufferAttributes)
+
+  override def producedAttributes: AttributeSet =
+    outputSet ++ AttributeSet(resultExpressions) ++ AttributeSet(aggregateAttributes)
 
   override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
     this.copy(child = newChild)
@@ -2091,6 +2102,8 @@ case class CometHashAggregateExec(
         this.output == other.output &&
         this.groupingExpressions == other.groupingExpressions &&
         this.aggregateExpressions == other.aggregateExpressions &&
+        this.aggregateAttributes == other.aggregateAttributes &&
+        this.resultExpressions == other.resultExpressions &&
         this.input == other.input &&
         this.modes == other.modes &&
         this.child == other.child &&
@@ -2101,7 +2114,15 @@ case class CometHashAggregateExec(
   }
 
   override def hashCode(): Int =
-    Objects.hashCode(output, groupingExpressions, aggregateExpressions, input, modes, child)
+    Objects.hashCode(
+      output,
+      groupingExpressions,
+      aggregateExpressions,
+      aggregateAttributes,
+      resultExpressions,
+      input,
+      modes,
+      child)
 
   override lazy val metrics: Map[String, SQLMetric] = {
     val baseline = CometMetricNode.baselineMetrics(sparkContext)

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

@@ -35,6 +35,7 @@ import org.apache.spark.sql.comet.CometHashAggregateExec
 import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec
 import org.apache.spark.sql.execution.SQLExecution
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
 import org.apache.spark.sql.functions.{avg, col, count_distinct, sum}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
@@ -1301,6 +1302,61 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  Seq(
+    ("COUNT(*)", 2L),
+    ("COUNT(DISTINCT _2)", 2L),
+    ("COUNT(DISTINCT _2) + SUM(_2)", 7L),
+    ("CAST(SIZE(COLLECT_SET(_2)) AS BIGINT)", 2L)).foreach { case (function, expected) =>
+    test(
+      s"aggregate canonicalization preserves result expressions and equivalent reuse: $function") {
+      withSQLConf(
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+        SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true",
+        SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+        CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+        CometConf.COMET_SHUFFLE_MODE.key -> "native") {
+        withParquetTable(Seq((0, 2), (0, 3)), "tbl") {
+          def aggregate(result: String, alias: String = "c"): DataFrame =
+            sql(s"SELECT $result AS $alias, _1 FROM tbl GROUP BY _1")
+              .repartition(2, col(alias), col("_1"))
+
+          def finalAggregate(df: DataFrame): CometHashAggregateExec =
+            df.queryExecution.executedPlan
+              .collectFirst {
+                case agg: CometHashAggregateExec if agg.modes.contains(Final) => agg
+              }
+              .getOrElse(fail("Expected a native final aggregate"))
+
+          val plus = aggregate(s"($function) + 1")
+          val minus = aggregate(s"($function) - 1")
+          // The shuffles above the final aggregates must not reuse each other: doing so
+          // would return the first projection twice, even without an existence join.
+          checkSparkAnswerAndOperator(plus.unionAll(minus), classOf[ReusedExchangeExec])
+          checkAnswer(plus.unionAll(minus), Seq(Row(expected + 1L, 0), Row(expected - 1L, 0)))
+          assert(!finalAggregate(plus).sameResult(finalAggregate(minus)))
+
+          // Comparing result expressions must still normalize aggregate-result attributes.
+          // Fresh expression IDs and a different output alias do not change the computation.
+          val same = aggregate(s"($function) + 1", "renamed")
+          assert(finalAggregate(plus).sameResult(finalAggregate(same)))
+          assert(finalAggregate(plus).semanticHash() == finalAggregate(same).semanticHash())
+          val (_, reusedPlan) =
+            checkSparkAnswerAndOperator(plus.unionAll(same), classOf[ReusedExchangeExec])
+          val reusedFinalAggregates = reusedPlan.collect {
+            case reused: ReusedExchangeExec if reused.child.exists {
+                  case agg: CometHashAggregateExec => agg.modes.contains(Final)
+                  case _ => false
+                } =>
+              reused
+          }
+          assert(
+            reusedFinalAggregates.nonEmpty,
+            s"Expected equivalent aggregate reuse:\n$reusedPlan")
+        }
+      }
+    }
+  }
+
   test("test final sum") {
     withSQLConf(
       CometConf.COMET_SHUFFLE_ENABLED.key -> "true",