diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala index fd3b9c484d..07a7f57ccd 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 d3f1ac03d3..ac1ca3854b 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,11 @@ 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.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 // 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 b2c3764fd6..f755374528 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. Seq( plan.treeString, - df.queryExecution.explainString(org.apache.spark.sql.execution.FormattedMode)) + if (isSpark40Plus) df.queryExecution.explainString(FormattedMode) + else scans.head.verboseStringWithOperatorId()) .foreach { text => assert(!text.contains("CachedRDDBuilder"), text) assert(!text.contains(classOf[ArrowCachedBatchSerializer].getName), text) @@ -365,6 +369,66 @@ 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 => + withClue(s"AQE $aqe: ") { + withSQLConf( + 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) + } + } + } + } + } + test("Comet in-memory cache disabled keeps SparkToColumnar fallback path") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",