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 216f9f1e4b..b7a1d7976c 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, 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) 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 143248f551..19de9f35df 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, 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",