From 91b45ef964c5d9d6317e791d1a99ae1e52884eff Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 26 Aug 2026 05:16:09 +0000 Subject: [PATCH 1/3] fix: preserve aggregate result identity during exchange reuse --- .../apache/spark/sql/comet/operators.scala | 27 ++++++++- .../comet/exec/CometAggregateSuite.scala | 56 +++++++++++++++++++ 2 files changed, 80 insertions(+), 3 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 e71603d4809..a4a4c9d49e7 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 @@ -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, Partial, PartialMerge, Percentile} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans._ @@ -1974,6 +1974,7 @@ object CometHashAggregateExec op.output, op.groupingExpressions, op.aggregateExpressions, + op.aggregateAttributes, op.resultExpressions, op.child.output, op.child, @@ -2024,6 +2025,7 @@ object CometObjectHashAggregateExec adjustOutputForNativeState(op), op.groupingExpressions, op.aggregateExpressions, + op.aggregateAttributes, op.resultExpressions, op.child.output, op.child, @@ -2037,6 +2039,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, @@ -2049,7 +2052,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) @@ -2072,6 +2083,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 && @@ -2082,7 +2095,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) 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 a356b5b78c3..bcda3c8a5cb 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/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, expr, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} @@ -1305,6 +1306,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", From ceed3a9f268b4e5c03ca6535597156ff240cc20f Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 9 Sep 2026 02:18:12 +0000 Subject: [PATCH 2/3] test: cover aggregate exchange reuse with AQE --- .../comet/exec/CometAggregateSuite.scala | 39 +++++++++++-------- 1 file changed, 22 insertions(+), 17 deletions(-) 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 bcda3c8a5cb..7e69fd2f25f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning 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.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper} import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.functions.{avg, col, count_distinct, expr, sum} import org.apache.spark.sql.internal.SQLConf @@ -1307,29 +1307,31 @@ 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) => + ("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( - s"aggregate canonicalization preserves result expressions and equivalent reuse: $function") { + s"aggregate canonicalization preserves result expressions and equivalent reuse: " + + s"$function, AQE=$adaptive") { withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + 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 = - df.queryExecution.executedPlan - .collectFirst { - case agg: CometHashAggregateExec if agg.modes.contains(Final) => agg - } - .getOrElse(fail("Expected a native final aggregate")) + 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") @@ -1346,11 +1348,14 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { 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 - } => + 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( From 5c67f3e4d06308ba8fdb048a2f6aa454b6d08bd2 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 9 Sep 2026 03:17:51 +0000 Subject: [PATCH 3/3] style: remove redundant test-name interpolation --- .../test/scala/org/apache/comet/exec/CometAggregateSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 7e69fd2f25f..3592b3b362f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -1313,7 +1313,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { ("CAST(SIZE(COLLECT_SET(_2)) AS BIGINT)", 2L, false), ("COUNT(*)", 2L, true)).foreach { case (function, expected, adaptive) => test( - s"aggregate canonicalization preserves result expressions and equivalent reuse: " + + "aggregate canonicalization preserves result expressions and equivalent reuse: " + s"$function, AQE=$adaptive") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,