From 81736558d43e00f55499192761c6cd9303c31444 Mon Sep 17 00:00:00 2001 From: comphead Date: Sat, 3 Oct 2026 11:20:34 -0700 Subject: [PATCH 1/2] fix: draw the cached plan below CometInMemoryTableScan in EXPLAIN Spark's InMemoryTableScanExec lists its InMemoryRelation, and through it the cached plan, as inner children, so EXPLAIN draws the plan that built the cache below the scan. Do the same for CometInMemoryTableScanExec. ExtendedExplainInfo walks inner children, so leave the scan's out of Comet's coverage and fallback reporting, as it already does for CometEmptyRelationExec. The cached plan runs when the relation is materialized, not as part of every query that reads it. Closes #6572. --- .../apache/comet/ExtendedExplainInfo.scala | 6 +- .../comet/CometInMemoryTableScanExec.scala | 7 ++ .../comet/exec/CometInMemoryCacheSuite.scala | 71 ++++++++++++++++++- 3 files changed, 81 insertions(+), 3 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala index fd3b9c484d5..07a7f57ccd8 100644 --- a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala +++ b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala @@ -26,7 +26,7 @@ import scala.collection.mutable import org.apache.spark.sql.ExtendedExplainGenerator import org.apache.spark.sql.catalyst.expressions.{Attribute, BoundReference, Expression, Literal, ScalaUDF} import org.apache.spark.sql.catalyst.trees.{TreeNode, TreeNodeTag} -import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometEmptyRelationExec, CometNativeColumnarToRowExec, CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometColumnarToRowExec, CometEmptyRelationExec, CometInMemoryTableScanExec, CometNativeColumnarToRowExec, CometPlan, CometSparkToColumnarExec} import org.apache.spark.sql.execution.{ColumnarToRowExec, InputAdapter, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} @@ -109,6 +109,10 @@ class ExtendedExplainInfo extends ExtendedExplainGenerator { // Spark's treeString displays the eliminated plan, but it does not execute and must not // contribute Spark operators, fallback reasons, or expressions to Comet's reporting. case _: CometEmptyRelationExec => Seq.empty + // Spark's treeString displays the cached relation and the plan that built it. That plan runs + // when the relation is materialized, not as part of every query that reads it, so it must not + // count toward the coverage or the fallback reasons of each of those queries. + case _: CometInMemoryTableScanExec => Seq.empty case _ => node.innerChildren } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala index d3f1ac03d38..3e8b45eb565 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala @@ -23,6 +23,7 @@ import scala.collection.JavaConverters._ import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.logical.Statistics import org.apache.spark.sql.columnar.{CachedBatch, CachedBatchSerializer} import org.apache.spark.sql.comet.shims.ShimCometInMemoryTableScanExec @@ -80,6 +81,12 @@ case class CometInMemoryTableScanExec( // newlines, which breaks the tree of every plan that reads the cache. override def stringArgs: Iterator[Any] = Iterator(originalPlan) + // Spark's own scan lists its InMemoryRelation as an inner child, and the relation lists the + // cached plan, so EXPLAIN draws the plan that built the cache below the scan. Do the same. + // ExtendedExplainInfo leaves them out of Comet's own reporting: the cached plan runs when the + // relation is materialized, not as part of every query that reads it. + override def innerChildren: Seq[QueryPlan[_]] = Seq(originalPlan.relation) + // `originalPlan` is a plan-typed field rather than a child, so QueryPlan's canonicalization // walks straight past it: its attributes and predicates keep the expression IDs of whichever // occurrence of the cached relation produced them. Two scans of one cache then compare unequal, diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index b2c3764fd6d..190a8ad97de 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -40,7 +40,7 @@ import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch} import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.execution.SortExec +import org.apache.spark.sql.execution.{FormattedMode, SortExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike} @@ -352,9 +352,13 @@ class CometInMemoryCacheSuite extends CometTestBase { line.startsWith("CometInMemoryTableScan Scan In-memory table explain_cache ["), line) assert(line.contains("= 3)"), s"the pruning predicates should be shown: $line") + // EXPLAIN FORMATTED also details the InMemoryRelation drawn below the scan. Before Spark + // 4.0 that detail prints the relation's CachedRDDBuilder below Spark's own scan as well + // (SPARK-51861), so there only the scan's own detail is checked. + val formatted = df.queryExecution.explainString(FormattedMode) Seq( plan.treeString, - df.queryExecution.explainString(org.apache.spark.sql.execution.FormattedMode)) + if (isSpark40Plus) formatted else scans.head.verboseStringWithOperatorId()) .foreach { text => assert(!text.contains("CachedRDDBuilder"), text) assert(!text.contains(classOf[ArrowCachedBatchSerializer].getName), text) @@ -365,6 +369,69 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + test("EXPLAIN draws the cached plan below CometInMemoryTableScan") { + // Spark's own cache scan draws its InMemoryRelation, and below that the plan that built the + // relation. See https://github.com/apache/datafusion-comet/issues/6572. + def nodeLines(tree: String): Seq[String] = + tree.linesIterator.map(_.replaceAll("^[ :+-]*", "")).toSeq + + Seq("false", "true").foreach { aqe => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + spark.catalog.clearCache() + spark + .range(1000) + .selectExpr("id AS key", "id % 8 AS value") + .createOrReplaceTempView("explain_cached_plan") + // The cached plan is planned here. Keeping its Range on Spark gives it a fallback reason + // that the reporting for a query reading the cache must leave out. + withSQLConf( + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false") { + spark.catalog.cacheTable("explain_cached_plan") + } + try { + val df = spark.sql("SELECT value, count(*) FROM explain_cached_plan GROUP BY value") + df.collect() + val plan = df.queryExecution.executedPlan + val scans = collect(plan) { case s: CometInMemoryTableScanExec => s } + assert(scans.size == 1, s"AQE $aqe: $plan") + val relation = scans.head.originalPlan.relation + + // The relation's line, then the cached plan's lines, follow the scan's line. + val lines = nodeLines(plan.treeString) + val scanAt = lines.indexWhere(_.startsWith("CometInMemoryTableScan ")) + assert(scanAt >= 0, s"AQE $aqe: ${plan.treeString}") + val below = nodeLines(relation.treeString) + assert(below.size > 1, relation.treeString) + assert( + lines.slice(scanAt + 1, scanAt + 1 + below.size) == below, + s"AQE $aqe: ${plan.treeString}") + + // EXPLAIN FORMATTED numbers the relation and draws it below the scan too. + val formatted = nodeLines(df.queryExecution.explainString(FormattedMode)) + val formattedScanAt = formatted.indexWhere(_.startsWith("CometInMemoryTableScan (")) + assert(formattedScanAt >= 0, s"AQE $aqe: ${formatted.mkString("\n")}") + assert( + formatted(formattedScanAt + 1).startsWith("InMemoryRelation ("), + s"AQE $aqe: ${formatted.mkString("\n")}") + + // Comet's own reporting leaves the cached plan out. + val info = new ExtendedExplainInfo() + val cachedReasons = info.getFallbackReasons(relation.cachedPlan) + assert(cachedReasons.nonEmpty, relation.cachedPlan.treeString) + val reasons = info.getFallbackReasons(plan) + assert(cachedReasons.forall(r => !reasons.contains(r)), s"AQE $aqe: $reasons") + val verbose = info.generateVerboseInfo(plan) + assert(!verbose.contains("InMemoryRelation"), s"AQE $aqe: $verbose") + } finally { + spark.catalog.clearCache() + } + } + } + } + test("Comet in-memory cache disabled keeps SparkToColumnar fallback path") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", From 53c7a83e744b0bafe350da6a9ed7ebe9df2cf12a Mon Sep 17 00:00:00 2001 From: comphead Date: Sat, 3 Oct 2026 13:39:45 -0700 Subject: [PATCH 2/2] Simplify the cached plan EXPLAIN test and its comment --- .../comet/CometInMemoryTableScanExec.scala | 3 +- .../comet/exec/CometInMemoryCacheSuite.scala | 101 +++++++++--------- 2 files changed, 50 insertions(+), 54 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala index 3e8b45eb565..ac1ca3854b7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala @@ -83,8 +83,7 @@ case class CometInMemoryTableScanExec( // Spark's own scan lists its InMemoryRelation as an inner child, and the relation lists the // cached plan, so EXPLAIN draws the plan that built the cache below the scan. Do the same. - // ExtendedExplainInfo leaves them out of Comet's own reporting: the cached plan runs when the - // relation is materialized, not as part of every query that reads it. + // ExtendedExplainInfo.executionInnerChildren leaves them out of Comet's own reporting. override def innerChildren: Seq[QueryPlan[_]] = Seq(originalPlan.relation) // `originalPlan` is a plan-typed field rather than a child, so QueryPlan's canonicalization diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 190a8ad97de..f7553745282 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -355,10 +355,10 @@ class CometInMemoryCacheSuite extends CometTestBase { // EXPLAIN FORMATTED also details the InMemoryRelation drawn below the scan. Before Spark // 4.0 that detail prints the relation's CachedRDDBuilder below Spark's own scan as well // (SPARK-51861), so there only the scan's own detail is checked. - val formatted = df.queryExecution.explainString(FormattedMode) Seq( plan.treeString, - if (isSpark40Plus) formatted else scans.head.verboseStringWithOperatorId()) + if (isSpark40Plus) df.queryExecution.explainString(FormattedMode) + else scans.head.verboseStringWithOperatorId()) .foreach { text => assert(!text.contains("CachedRDDBuilder"), text) assert(!text.contains(classOf[ArrowCachedBatchSerializer].getName), text) @@ -376,57 +376,54 @@ class CometInMemoryCacheSuite extends CometTestBase { tree.linesIterator.map(_.replaceAll("^[ :+-]*", "")).toSeq Seq("false", "true").foreach { aqe => - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, - CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { - spark.catalog.clearCache() - spark - .range(1000) - .selectExpr("id AS key", "id % 8 AS value") - .createOrReplaceTempView("explain_cached_plan") - // The cached plan is planned here. Keeping its Range on Spark gives it a fallback reason - // that the reporting for a query reading the cache must leave out. + withClue(s"AQE $aqe: ") { withSQLConf( - CometConf.COMET_EXEC_RANGE_ENABLED.key -> "false", - CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false") { - spark.catalog.cacheTable("explain_cached_plan") - } - try { - val df = spark.sql("SELECT value, count(*) FROM explain_cached_plan GROUP BY value") - df.collect() - val plan = df.queryExecution.executedPlan - val scans = collect(plan) { case s: CometInMemoryTableScanExec => s } - assert(scans.size == 1, s"AQE $aqe: $plan") - val relation = scans.head.originalPlan.relation - - // The relation's line, then the cached plan's lines, follow the scan's line. - val lines = nodeLines(plan.treeString) - val scanAt = lines.indexWhere(_.startsWith("CometInMemoryTableScan ")) - assert(scanAt >= 0, s"AQE $aqe: ${plan.treeString}") - val below = nodeLines(relation.treeString) - assert(below.size > 1, relation.treeString) - assert( - lines.slice(scanAt + 1, scanAt + 1 + below.size) == below, - s"AQE $aqe: ${plan.treeString}") - - // EXPLAIN FORMATTED numbers the relation and draws it below the scan too. - val formatted = nodeLines(df.queryExecution.explainString(FormattedMode)) - val formattedScanAt = formatted.indexWhere(_.startsWith("CometInMemoryTableScan (")) - assert(formattedScanAt >= 0, s"AQE $aqe: ${formatted.mkString("\n")}") - assert( - formatted(formattedScanAt + 1).startsWith("InMemoryRelation ("), - s"AQE $aqe: ${formatted.mkString("\n")}") - - // Comet's own reporting leaves the cached plan out. - val info = new ExtendedExplainInfo() - val cachedReasons = info.getFallbackReasons(relation.cachedPlan) - assert(cachedReasons.nonEmpty, relation.cachedPlan.treeString) - val reasons = info.getFallbackReasons(plan) - assert(cachedReasons.forall(r => !reasons.contains(r)), s"AQE $aqe: $reasons") - val verbose = info.generateVerboseInfo(plan) - assert(!verbose.contains("InMemoryRelation"), s"AQE $aqe: $verbose") - } finally { - spark.catalog.clearCache() + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + withTempView("explain_cached_plan") { + spark + .range(1000) + .selectExpr("id AS key", "id % 8 AS value") + .createOrReplaceTempView("explain_cached_plan") + // The cached plan is planned here. Keeping its Range on Spark gives it a fallback + // reason that the reporting for a query reading the cache must leave out. + withSQLConf( + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "false", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false") { + spark.catalog.cacheTable("explain_cached_plan") + } + val df = spark.sql("SELECT value, count(*) FROM explain_cached_plan GROUP BY value") + df.collect() + val plan = df.queryExecution.executedPlan + val scans = collect(plan) { case s: CometInMemoryTableScanExec => s } + assert(scans.size == 1, plan) + val relation = scans.head.originalPlan.relation + + // The relation's line, then the cached plan's lines, follow the scan's line. + val lines = nodeLines(plan.treeString) + val below = nodeLines(relation.treeString) + assert(below.size > 1, relation) + val scanAt = lines.indexWhere(_.startsWith("CometInMemoryTableScan ")) + assert(lines.startsWith(below, scanAt + 1), plan) + + // EXPLAIN FORMATTED numbers the relation and draws it below the scan too. + val formatted = df.queryExecution.explainString(FormattedMode) + val formattedLines = nodeLines(formatted) + val formattedScanAt = + formattedLines.indexWhere(_.startsWith("CometInMemoryTableScan (")) + assert( + formattedLines(formattedScanAt + 1).startsWith("InMemoryRelation ("), + formatted) + + // Comet's own reporting leaves the cached plan out. + val info = new ExtendedExplainInfo() + val cachedReasons = info.getFallbackReasons(relation.cachedPlan) + assert(cachedReasons.nonEmpty, relation.cachedPlan) + val reasons = info.getFallbackReasons(plan) + assert(cachedReasons.intersect(reasons).isEmpty, reasons) + val verbose = info.generateVerboseInfo(plan) + assert(!verbose.contains("InMemoryRelation"), verbose) + } } } }