This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new c8ee6aef50 fix: preserve aggregate result identity during exchange
reuse (#5470)
c8ee6aef50 is described below
commit c8ee6aef50dcb4d4f8592dec4d264f6a81a4a0c9
Author: Chao Sun <[email protected]>
AuthorDate: Wed Sep 9 14:41:43 2026 -0700
fix: preserve aggregate result identity during exchange reuse (#5470)
* fix: preserve aggregate result identity during exchange reuse
* test: cover aggregate exchange reuse with AQE
* style: remove redundant test-name interpolation
---
.../org/apache/spark/sql/comet/operators.scala | 27 ++++++++--
.../apache/comet/exec/CometAggregateSuite.scala | 63 +++++++++++++++++++++-
2 files changed, 86 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 e71603d480..a4a4c9d49e 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 a356b5b78c..3592b3b362 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",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
