You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/30 10:41:29 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-5213-d8c01472bf99fc72dcbf6db7dc06f09d4efc69b3 created (now 490138d88e)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-5213-d8c01472bf99fc72dcbf6db7dc06f09d4efc69b3
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 490138d88e perf: memoize Comet-subtree check in EliminateRedundantTransitions (#5213)
This branch includes the following new commits:
new 490138d88e perf: memoize Comet-subtree check in EliminateRedundantTransitions (#5213)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: perf: memoize Comet-subtree check in EliminateRedundantTransitions (#5213)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5213-d8c01472bf99fc72dcbf6db7dc06f09d4efc69b3
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 490138d88eee256332703cb0bfaee7603c240142
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Sep 30 09:30:14 2026 +0000
perf: memoize Comet-subtree check in EliminateRedundantTransitions (#5213)
* perf: memoize Comet-subtree check in EliminateRedundantTransitions
hasCometNativeChild scans the subtree below every ColumnarToRowExec the
rule visits, so stacked transitions rescan the same nodes and the
combined cost is quadratic in plan size. The scan stops at the first
Comet operator it finds, which keeps the common case cheap, but a stack
of transitions over a Comet-free subtree walks all of it every time.
Memoize the scan for the duration of a single rule invocation. On a
6403-node chain of transitions over a vanilla Parquet scan this cuts the
nodes visited from 10.2M to 6402 and the rule from 264 ms to 0.9 ms.
The memo is keyed on identity because SparkPlan equality and hash are
themselves subtree walks, and it is a local of _apply because the rule
instance lives for the whole session and must not retain plans.
Part of #5199.
* test: register EliminateRedundantTransitionsSuite in CI workflow matrices
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
.../rules/EliminateRedundantTransitions.scala | 46 ++++++++++++--
.../rules/EliminateRedundantTransitionsSuite.scala | 73 ++++++++++++++++++++++
4 files changed, 116 insertions(+), 5 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index ff12e42d49..ee8a947aa4 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -589,6 +589,7 @@ jobs:
org.apache.comet.rules.CometExecRuleSuite
org.apache.spark.sql.comet.CometPlanEqualitySuite
org.apache.comet.rules.CometEmptyRelationExecRuleSuite
+ org.apache.comet.rules.EliminateRedundantTransitionsSuite
org.apache.comet.rules.RevertNativeForTransitionHeavyStagesSuite
org.apache.spark.sql.CometTPCDSQuerySuite
org.apache.spark.sql.CometTPCDSQueryTestSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index d3579dc137..371e4f6531 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -246,6 +246,7 @@ jobs:
org.apache.comet.rules.CometExecRuleSuite
org.apache.spark.sql.comet.CometPlanEqualitySuite
org.apache.comet.rules.CometEmptyRelationExecRuleSuite
+ org.apache.comet.rules.EliminateRedundantTransitionsSuite
org.apache.comet.rules.RevertNativeForTransitionHeavyStagesSuite
org.apache.spark.sql.CometTPCDSQuerySuite
org.apache.spark.sql.CometTPCDSQueryTestSuite
diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
index d076c14f74..b576961e46 100644
--- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
@@ -19,6 +19,8 @@
package org.apache.comet.rules
+import java.util.IdentityHashMap
+
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.util.sideBySide
@@ -75,6 +77,19 @@ case class EliminateRedundantTransitions(session: SparkSession)
}
private def _apply(plan: SparkPlan): SparkPlan = {
+ // `hasCometNativeChild` scans the subtree below every `ColumnarToRowExec` it is asked about,
+ // so stacked transitions rescan the same nodes and the combined cost is quadratic in the plan
+ // size. The scan stops at the first Comet operator it finds, which keeps the common case
+ // cheap, but a stack of transitions over a Comet-free subtree walks all of it every time. The
+ // memo below makes that linear.
+ //
+ // It is keyed on identity rather than equality because `SparkPlan` equality and hashing are
+ // themselves subtree walks, and it is scoped to a single rule invocation because the rule
+ // instance lives for the whole session and must not retain plans. `transformUp` reuses the
+ // identity of subtrees it does not rewrite, so a rebuilt node still hits the memo one level
+ // down.
+ val containsCometPlanMemo = new IdentityHashMap[SparkPlan, java.lang.Boolean]()
+
val eliminatedPlan = stripIcebergWriteInputTransition(plan) transformUp {
case ColumnarToRowExec(shuffleExchangeExec: CometShuffleExchangeExec)
if plan.conf.adaptiveExecutionEnabled =>
@@ -92,7 +107,7 @@ case class EliminateRedundantTransitions(session: SparkSession)
// Write should be final operation in the plan
case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) =>
nativeWrite
- case c @ ColumnarToRowExec(child) if hasCometNativeChild(child) =>
+ case c @ ColumnarToRowExec(child) if hasCometNativeChild(child, containsCometPlanMemo) =>
val op = createColumnarToRowExec(child)
if (c.logicalLink.isEmpty) {
op.unsetTagValue(SparkPlan.LOGICAL_PLAN_TAG)
@@ -162,11 +177,32 @@ case class EliminateRedundantTransitions(session: SparkSession)
}
}
- private def hasCometNativeChild(op: SparkPlan): Boolean = {
+ /**
+ * True if the subtree rooted at `op` contains a Comet operator. `QueryStageExec` and
+ * `ReusedExchangeExec` are leaves for tree traversal, so the plan they wrap is unwrapped
+ * explicitly, and only at the root of the checked subtree.
+ */
+ private def hasCometNativeChild(
+ op: SparkPlan,
+ memo: IdentityHashMap[SparkPlan, java.lang.Boolean]): Boolean = {
op match {
- case c: QueryStageExec => hasCometNativeChild(c.plan)
- case c: ReusedExchangeExec => hasCometNativeChild(c.child)
- case _ => op.exists(_.isInstanceOf[CometPlan])
+ case c: QueryStageExec => hasCometNativeChild(c.plan, memo)
+ case c: ReusedExchangeExec => hasCometNativeChild(c.child, memo)
+ case _ => containsCometPlan(op, memo)
+ }
+ }
+
+ /** Memoized equivalent of `op.exists(_.isInstanceOf[CometPlan])`. */
+ private def containsCometPlan(
+ op: SparkPlan,
+ memo: IdentityHashMap[SparkPlan, java.lang.Boolean]): Boolean = {
+ val cached = memo.get(op)
+ if (cached != null) {
+ cached.booleanValue()
+ } else {
+ val result = op.isInstanceOf[CometPlan] || op.children.exists(containsCometPlan(_, memo))
+ memo.put(op, result)
+ result
}
}
diff --git a/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala b/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala
new file mode 100644
index 0000000000..33293c3d34
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/rules/EliminateRedundantTransitionsSuite.scala
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.rules
+
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometNativeColumnarToRowExec, CometPlan}
+import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnionExec}
+
+import org.apache.comet.CometConf
+
+class EliminateRedundantTransitionsSuite extends CometTestBase {
+
+ private def sparkPlanFor(sql: String): SparkPlan = {
+ var plan: SparkPlan = null
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ plan = spark.sql(sql).queryExecution.executedPlan
+ }
+ stripAQEPlan(plan)
+ }
+
+ private def isCometColumnarToRow(plan: SparkPlan): Boolean = plan match {
+ case _: CometColumnarToRowExec => true
+ case _: CometNativeColumnarToRowExec => true
+ case _ => false
+ }
+
+ // The rule decides, per `ColumnarToRowExec`, whether the subtree below it is Comet columnar.
+ // That check is memoized across the whole plan, so a Comet branch must not make a sibling
+ // Spark-only branch look Comet, or the other way around.
+ test("sibling branches are classified independently") {
+ withParquetTable((0 until 10).map(i => (i, i.toString)), "tbl") {
+ val sparkPlan = sparkPlanFor("SELECT _1 FROM tbl")
+ val sparkScan = sparkPlan.collectFirst { case p if p.supportsColumnar => p }
+ assume(sparkScan.isDefined, "test requires a columnar Spark scan")
+ assert(!sparkScan.get.exists(_.isInstanceOf[CometPlan]))
+
+ val cometPlan = CometExecRule(spark).apply(CometScanRule(spark).apply(sparkPlan))
+ val cometScan = cometPlan.collectFirst { case p: CometPlan if p.supportsColumnar => p }
+ assume(cometScan.isDefined, "test requires a columnar Comet scan")
+
+ // Two columnar branches producing the same output, only one of which is Comet.
+ val sparkTransition = ColumnarToRowExec(sparkScan.get)
+ val union = UnionExec(Seq(ColumnarToRowExec(cometScan.get), sparkTransition))
+
+ val result = EliminateRedundantTransitions(spark).apply(union)
+
+ val Seq(cometBranch, sparkBranch) = result.children
+ assert(
+ isCometColumnarToRow(cometBranch),
+ s"Comet branch should use a Comet columnar-to-row transition:\n${result.treeString}")
+ assert(
+ sparkBranch == sparkTransition,
+ s"Spark-only branch should be left untouched:\n${result.treeString}")
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org