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..3592b3b362f 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,8 @@ 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 import org.apache.spark.sql.types.{DataTypes, StructField, StructType} @@ -1305,6 +1306,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",