From 3d6b9523ae559d8ec9e5855d54207755d11f15d3 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 27 Sep 2026 09:57:17 -0600 Subject: [PATCH] fix: preserve aggregate result identity during exchange reuse (#5470) CometHashAggregateExec's equals and hashCode left out resultExpressions, so two final aggregates that differ only in their result expressions, such as COUNT(*) + 1 and COUNT(*) - 1 over the same input, compared equal after canonicalization. Exchange reuse then served both from one shuffle and returned the first projection twice. Include the result expressions and the aggregate attributes in plan identity, and canonicalize the aggregate the way Spark does so equivalent aggregates are still reused. Adapted for branch-1.0: operators.scala keeps its First and Last imports, and CometAggregateSuite imports Final, which the new test uses. (cherry picked from commit c8ee6aef50dcb4d4f8592dec4d264f6a81a4a0c9) --- .../apache/spark/sql/comet/operators.scala | 27 +++++++- .../comet/exec/CometAggregateSuite.scala | 64 ++++++++++++++++++- 2 files changed, 87 insertions(+), 4 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 73d50931c19..c0ffd9a0ff5 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -32,7 +32,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._ @@ -2011,6 +2011,7 @@ object CometHashAggregateExec op.output, op.groupingExpressions, op.aggregateExpressions, + op.aggregateAttributes, op.resultExpressions, op.child.output, op.child, @@ -2061,6 +2062,7 @@ object CometObjectHashAggregateExec adjustOutputForNativeState(op), op.groupingExpressions, op.aggregateExpressions, + op.aggregateAttributes, op.resultExpressions, op.child.output, op.child, @@ -2074,6 +2076,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, @@ -2086,7 +2089,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) @@ -2109,6 +2120,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 && @@ -2119,7 +2132,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 protected def outputExpressions: Seq[NamedExpression] = resultExpressions } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 2aaf3eb8290..dd44000c940 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -25,9 +25,11 @@ import org.apache.hadoop.fs.Path import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.Cast +import org.apache.spark.sql.catalyst.expressions.aggregate.Final import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.comet.CometHashAggregateExec -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, 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} @@ -1165,6 +1167,66 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + Seq( + ("COUNT(*)", 2L, false), + ("COUNT(DISTINCT _2)", 2L, false), + ("COUNT(DISTINCT _2) + SUM(_2)", 7L, false), + ("CAST(SIZE(COLLECT_SET(_2)) AS BIGINT)", 2L, false), + ("COUNT(*)", 2L, true)).foreach { case (function, expected, adaptive) => + test( + "aggregate canonicalization preserves result expressions and equivalent reuse: " + + s"$function, AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + 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") { + // Build independent branches with an exchange above Final and the requested output alias. + 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")) + + // Traverse adaptive/query-stage wrappers and fail if the native Final fell back to Spark. + def finalAggregate(df: DataFrame): CometHashAggregateExec = + collectFirst(df.queryExecution.executedPlan) { + 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]) + if (adaptive) { + assert(reusedPlan.isInstanceOf[AdaptiveSparkPlanExec]) + } + // Adaptive-aware traversal must find reuse above Final, not just a shared Partial stage. + val reusedFinalAggregates = collect(reusedPlan) { + case reused: ReusedExchangeExec if collect(reused.child) { + case agg: CometHashAggregateExec if agg.modes.contains(Final) => agg + }.nonEmpty => + reused + } + assert( + reusedFinalAggregates.nonEmpty, + s"Expected equivalent aggregate reuse:\n$reusedPlan") + } + } + } + } + test("test final sum") { withSQLConf( CometConf.COMET_SHUFFLE_ENABLED.key -> "true",