From 717d4239629c96c5fcec878c591442fbeadee242 Mon Sep 17 00:00:00 2001 From: 0lai0 Date: Mon, 7 Sep 2026 01:52:17 +0800 Subject: [PATCH] fix: explain ObjectHashAggregate fallback when Comet shuffle is disabled --- docs/source/user-guide/latest/operators.md | 10 +- .../apache/comet/rules/CometExecRule.scala | 6 - .../apache/spark/sql/comet/operators.scala | 24 ++-- .../comet/rules/CometExecRuleSuite.scala | 124 +++++++++++++++++- 4 files changed, 143 insertions(+), 21 deletions(-) diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 3a18d86606d..59d704c5e7e 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -76,11 +76,11 @@ omitted from the tables below and may be reconsidered based on demand: ## Aggregation -| Operator | Status | Notes | -| ------------------------- | ------ | ----------------------------------------------------------------- | -| `HashAggregateExec` | ✅ | | -| `ObjectHashAggregateExec` | ✅ | Supports a limited set of aggregates, such as `bloom_filter_agg`. | -| `SortAggregateExec` | 🔜 | Falls back today; Comet currently accelerates hash aggregates. | +| Operator | Status | Notes | +| ------------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `HashAggregateExec` | ✅ | | +| `ObjectHashAggregateExec` | ✅ | Supports a limited set of aggregates, such as `bloom_filter_agg`. Falls back when Comet shuffle is disabled, which would otherwise split the aggregate across Comet and Spark. See the [Tuning Guide](tuning.md). | +| `SortAggregateExec` | 🔜 | Falls back today; Comet currently accelerates hash aggregates. | ## Joins diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index db5373e1093..aeafa822c45 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -1178,12 +1178,6 @@ case class CometExecRule(session: SparkSession) val serde = handler.get.asInstanceOf[CometOperatorSerde[SparkPlan]] if (!isOperatorEnabled(serde, agg.asInstanceOf[SparkPlan])) return false - // ObjectHashAggregate has an extra shuffle-enabled guard in its convert method - agg match { - case _: ObjectHashAggregateExec if !isCometShuffleEnabled(agg.conf) => return false - case _ => - } - val aggregateExpressions = agg.aggregateExpressions val groupingExpressions = agg.groupingExpressions 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 3700e97642b..53a43a4c59c 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 @@ -2019,6 +2019,22 @@ object CometObjectHashAggregateExec op.aggregateExpressions.exists(_.mode == Final)) { return Unsupported(Some("Final aggregates disabled via test config")) } + // When Comet shuffle is disabled we do not want to transform the ObjectHashAggregate to + // CometHashAggregate, because we would probably get partial Comet aggregation and final + // Spark aggregation. The reason text stays mode-neutral because `getSupportLevel` runs for + // every stage, so the Final node receives it too and must not be told it is the partial. + // Declining here rather than in `convert` is what lets CometExecRule + // record the reason centrally, so the fallback is explained rather than silent - see + // https://github.com/apache/datafusion-comet/issues/5500. The message deliberately does not + // name a config key: `isCometShuffleEnabled` is a conjunction of the shuffle config, the + // configured shuffle manager and the Celeborn compatibility check, so naming one of them + // would misdirect when another is the cause. + if (!isCometShuffleEnabled(op.conf)) { + return Unsupported( + Some( + "Comet shuffle is not enabled, so converting ObjectHashAggregate would split the " + + "aggregate across Comet and Spark")) + } Compatible() } @@ -2026,14 +2042,6 @@ object CometObjectHashAggregateExec aggregate: ObjectHashAggregateExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { - - if (!isCometShuffleEnabled(aggregate.conf)) { - // When Comet shuffle is disabled, we don't want to transform the HashAggregate - // to CometHashAggregate. Otherwise, we probably get partial Comet aggregation - // and final Spark aggregation. - return None - } - doConvert(aggregate, builder, childOp: _*) } diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 5444a89fa36..79b668444d8 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -24,17 +24,19 @@ import scala.util.Random import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.FunctionIdentifier import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo} -import org.apache.spark.sql.catalyst.expressions.aggregate.BloomFilterAggregate +import org.apache.spark.sql.catalyst.expressions.aggregate.{BloomFilterAggregate, Partial} import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec} import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataTypes, StructField, StructType} -import org.apache.comet.{CometConf, CometExplainInfo} +import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus} +import org.apache.comet.serde.{Compatible, Unsupported} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} /** @@ -77,6 +79,30 @@ class CometExecRuleSuite extends CometTestBase { }.sum } + /** + * Build a Spark plan containing an `ObjectHashAggregateExec` and hand it to `f`. `collect_list` + * is a `TypedImperativeAggregate`, so Spark plans it as `ObjectHashAggregateExec` rather than + * `HashAggregateExec`. Each call builds a fresh plan, which matters because fallback reasons + * accumulate on plan-node tags. + */ + private def withObjectHashAggPlan(f: SparkPlan => Unit): Unit = { + withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + val sparkPlan = + createSparkPlan(spark, "SELECT id, collect_list(name) FROM test_data GROUP BY id") + assert(countOperators(sparkPlan, classOf[ObjectHashAggregateExec]) > 0) + f(sparkPlan) + } + } + } + + /** The partial-mode `ObjectHashAggregateExec` in `plan`. */ + private def partialObjectHashAgg(plan: SparkPlan): ObjectHashAggregateExec = + stripAQEPlan(plan).collectFirst { + case a: ObjectHashAggregateExec if a.aggregateExpressions.forall(_.mode == Partial) => a + }.get + test("expression-level fallback reasons are rolled up onto the operator that falls back") { // Extended explain only walks plan nodes, so a reason recorded on a sub-expression is // invisible unless CometExecRule lifts it onto the enclosing operator. Disabling a single @@ -175,6 +201,100 @@ class CometExecRuleSuite extends CometTestBase { } } + test("ObjectHashAggregate records a reason when it declines because Comet shuffle is off") { + // CometObjectHashAggregateExec deliberately declines when Comet shuffle is disabled, because + // converting it would leave a Comet partial aggregate feeding a Spark final aggregate. That + // decline used to return None from `convert` without recording why, so strict mode saw an + // unexplained fallback and the generic " is not supported" message hid the real + // cause. See https://github.com/apache/datafusion-comet/issues/5500. + // + // Both strict settings matter and are covered here: strict mode turns a missing reason into a + // hard failure, while the lenient production default is where the generic message used to + // stand in for the real cause, so it is the half a user actually sees. + Seq(true, false).foreach { strictFallbackReasons => + // A fresh plan per iteration: fallback reasons accumulate on plan-node tags, so reusing + // one plan would let the first iteration's reasons satisfy the second. + withObjectHashAggPlan { sparkPlan => + withSQLConf( + CometConf.COMET_SHUFFLE_ENABLED.key -> "false", + CometConf.COMET_STRICT_FALLBACK_REASONS.key -> strictFallbackReasons.toString, + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val transformedPlan = applyCometExecRule(sparkPlan) + // Assert on the partial stage: its child is the Comet local table scan, so it is the + // node whose children are all native and which therefore reaches the strict + // unexplained-fallback check in CometExecRule. + val partialAgg = partialObjectHashAgg(transformedPlan) + + val reasons = partialAgg + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .getOrElse(Set.empty[String]) + assert( + reasons.exists(_.contains("Comet shuffle is not enabled")), + "expected the shuffle-disabled reason on the ObjectHashAggregateExec, " + + s"got: $reasons") + // The generic catch-all must not appear: a real reason was available. + assert( + !reasons.contains(s"${partialAgg.nodeName} is not supported"), + s"a real reason was available but the generic message was used too: $reasons") + + // The reason must survive into the rendered extended-explain output, which is what + // the user actually reads. + val explained = new ExtendedExplainInfo().getFallbackReasons(transformedPlan) + // Match the tail, not the "Comet shuffle is not enabled" prefix: with shuffle off, + // CometShuffleExchangeExec.shuffleSupported tags the ShuffleExchangeExec with its own + // "Comet shuffle is not enabled: ..." reason, and getFallbackReasons flattens every + // node's tags into one unattributed set. The prefix alone would match that instead. + assert( + explained.exists(_.contains("would split the aggregate across Comet and Spark")), + s"expected the aggregate's own reason in extended explain output, got: $explained") + } + } + } + } + + test("CometObjectHashAggregateExec reports the shuffle-disabled decline in getSupportLevel") { + // The decline belongs in getSupportLevel, not convert: that is where CometExecRule attaches + // the fallback reason centrally, and it matches CometCollectLimitExec and + // CometTakeOrderedAndProjectExec, which gate on the same shuffle predicate. + // See https://github.com/apache/datafusion-comet/issues/5500. + withObjectHashAggPlan { sparkPlan => + val agg = partialObjectHashAgg(sparkPlan) + + withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "false") { + CometObjectHashAggregateExec.getSupportLevel(agg) match { + case Unsupported(Some(notes)) => + // Assert the aggregate-specific tail, not the "Comet shuffle is not enabled" + // prefix: that prefix is shared with CometCollectLimitExec and + // CometTakeOrderedAndProjectExec, so matching it alone would not notice this + // message losing the half that explains the split. + assert( + notes.contains("would split the aggregate across Comet and Spark"), + s"expected the aggregate-specific shuffle-disabled reason, got: $notes") + case other => + fail(s"expected Unsupported with a shuffle-disabled reason, got: $other") + } + } + + // Enabling shuffle must leave eligibility untouched. + withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { + assert(CometObjectHashAggregateExec.getSupportLevel(agg).isInstanceOf[Compatible]) + } + } + } + + test("ObjectHashAggregate still converts when Comet shuffle is enabled") { + // Guards the other half of https://github.com/apache/datafusion-comet/issues/5500: adding the + // fallback reason must not change which aggregates Comet accepts. + withObjectHashAggPlan { sparkPlan => + withSQLConf( + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val transformedPlan = applyCometExecRule(sparkPlan) + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) > 0) + } + } + } + test( "CometExecRule should apply basic operator transformations, but only when Comet is enabled") { withTempView("test_data") {