unikdahal commented on code in PR #5514:
URL: https://github.com/apache/datafusion-comet/pull/5514#discussion_r4064889890
##########
spark/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -756,6 +756,19 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val COMET_EXPLAIN_PLAN_ONLY_ENABLED: ConfigEntry[Boolean] =
+ conf("spark.comet.explain.planOnly.enabled")
+ .category(CATEGORY_EXEC_EXPLAIN)
+ .doc("When enabled, Comet leaves the query for Spark to execute and
afterwards logs the " +
+ "Comet plan it would have executed, with a coverage summary. Use this
to evaluate how " +
+ "much of a workload Comet would accelerate without changing execution.
The estimate is " +
+ "Scala-side only; the plan is never handed to DataFusion, so native
planning failures " +
+ "are not surfaced and the acceleration percentage can be optimistic.
Reported once per " +
+ "action, so a plan built but never executed is not reported. Requires
" +
Review Comment:
Small wording nit: should this say "once per SQL execution" instead of "once
per action"?
The RDD case documented in `CometPlanOnly` shows that the two aren't always
the same.
##########
spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala:
##########
@@ -0,0 +1,413 @@
+/*
+ * 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 scala.util.control.NonFatal
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.expressions.{Attribute, ExprId}
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions,
BaseSubqueryExec, ColumnarToRowExec, CommandResultExec, ExecSubqueryExpression,
InputAdapter, QueryExecution, ReusedSubqueryExec, RowToColumnarExec, SparkPlan,
WholeStageCodegenExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
AQEShuffleReadExec, QueryStageExec}
+import org.apache.spark.sql.execution.command.ExecutedCommandExec
+import org.apache.spark.sql.execution.datasources.v2.V2CommandExec
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec,
ReusedExchangeExec}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.QueryExecutionListener
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+import org.apache.comet.CometSparkSessionExtensions.isCometLoaded
+
+/**
+ * Plan-only mode: report the Comet plan Comet would have executed for a
query, without offloading
+ * any of it to Comet. See `spark.comet.explain.planOnly.enabled`.
+ *
+ * The conversion rules leave the plan alone while the mode is on, so Spark
plans and executes the
+ * query exactly as it would with Comet switched off. The report is built
afterwards, from the
+ * plan Spark actually executed, and thrown away. Comet code therefore cannot
reach the query: not
+ * by planning it, and not by failing while describing it.
+ *
+ * Reporting once the query is over, rather than while it is being planned, is
what keeps this
+ * simple. Spark applies a planner rule many times for one query - once per
query stage and once
+ * per adaptive re-optimization under AQE, and separately for every subquery
it prepares - and
+ * telling those applications apart takes a good deal of bookkeeping. A query
execution listener
+ * fires once per SQL execution, holding the finished plan, so there is
nothing to tell apart.
+ *
+ * Per SQL execution, not per action: Spark drives `QueryExecutionListener`
from the Dataset
+ * action path, so work taken through `df.rdd` is outside it. On Spark 4.0+
obtaining `df.rdd`
+ * runs a query of its own and is reported at that point, while the RDD
actions that follow are
+ * not; on 3.4/3.5 obtaining it is not reported either. Covering those would
mean a second
+ * reporting path keyed on job starts, and reconciling it against this one so
an ordinary query is
+ * not reported twice.
+ */
+object CometPlanOnly extends Logging {
+
+ private val REPORT_PREFIX = "[Comet plan-only]"
+
+ /**
+ * Sessions that already have a listener registered. Weakly held so a
session that goes away is
+ * not kept alive by this, and so a long-lived driver retains no more state
than the sessions it
+ * is running.
+ */
+ private val registeredSessions: java.util.Set[SparkSession] =
+ java.util.Collections.synchronizedSet(
+ java.util.Collections.newSetFromMap(
+ new java.util.WeakHashMap[SparkSession, java.lang.Boolean]()))
+
+ /**
+ * Registers this session's plan-only listener, if it does not have one yet.
+ *
+ * Called from `CometExecRule` rather than at session creation so that a
session never carries a
+ * listener unless plan-only mode is actually used, and so the config can be
turned on part way
+ * through a session.
+ */
+ def register(session: SparkSession): Unit = {
+ if (registeredSessions.add(session)) {
+ session.listenerManager.register(new CometPlanOnlyListener)
+ logInfo(s"$REPORT_PREFIX registered a plan-only reporter for this
session")
+ }
+ }
+
+ /**
+ * The settings that decide whether a query gets a report and how its
preview is built, as they
+ * stood while the query was being planned.
+ *
+ * Reporting is asynchronous, so by the time the listener runs the caller
may have restored or
+ * changed any of these - a `withSQLConf` block that runs `collect()` and
exits before the
+ * callback is delivered is enough. Reading them back off the session then
decides one query's
+ * report using another query's settings, which usually means dropping it.
Snapshotting at plan
+ * time keeps the decision with the query it belongs to.
+ *
+ * `sqlConfs` is every SQL conf that was explicitly set, not just Comet's.
The preview reruns
+ * the conversion rules, and those read far more than three flags:
`isCometLoaded` and the two
+ * rules take the session conf, `CometExecRule` reads per-operator gates and
the strict-fallback
+ * and shuffle settings off `op.conf`, and the serde reads Spark settings
such as ANSI mode and
+ * the session time zone. `op.conf` is `session.sessionState.conf` of the
session each node
+ * captured when it was constructed (`SparkPlan.conf`), so no thread-local
override and no
+ * cloned session can redirect it. The snapshot therefore cannot be
*injected* into the preview;
+ * it is used to check that the configuration the preview will read is still
the one the query
+ * was planned under, and the report is skipped when it is not. See
[[changedSince]].
+ */
+ private case class PlanOnlySettings(
+ enabled: Boolean,
+ cometLoaded: Boolean,
+ execEnabled: Boolean,
+ sqlConfs: Map[String, String]) {
+
+ /**
+ * The settings that have changed since this snapshot was taken, other
than the reporting gate
+ * itself.
+ *
+ * `spark.comet.explain.planOnly.enabled` is deliberately excluded:
turning plan-only mode
+ * off, or leaving the `withSQLConf` block that turned it on, must not
drop the report for a
+ * query that was planned while it was on. That flag only gates reporting,
and the snapshot is
+ * the authority for it. Everything else shapes the preview, so a
difference there means the
+ * preview would describe a configuration the query never ran under.
+ */
+ def changedSince(conf: SQLConf): Seq[String] = {
+ val now = conf.getAllConfs
+ val gate = CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key
+ (sqlConfs.keySet ++ now.keySet).toSeq.sorted
+ .filterNot(_ == gate)
+ .filter(key => sqlConfs.get(key) != now.get(key))
+ }
+ }
+
+ /**
+ * Set on the plan `CometExecRule` saw while plan-only mode was on. Read
back in `report`, which
+ * runs on the listener bus with no access to the planning thread's
configuration.
+ */
+ private val PLAN_ONLY_SETTINGS = new
TreeNodeTag[PlanOnlySettings]("CometPlanOnlySettings")
+
+ private def snapshot(conf: SQLConf): PlanOnlySettings =
+ PlanOnlySettings(
+ enabled = CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf),
+ cometLoaded = isCometLoaded(conf),
+ execEnabled = CometConf.COMET_EXEC_ENABLED.get(conf),
+ sqlConfs = conf.getAllConfs)
+
+ /**
+ * Record the plan-time settings on `plan` and return it unchanged.
+ *
+ * Called by `CometExecRule` on the plan it is declining to convert.
+ */
+ def tagSettings(session: SparkSession, plan: SparkPlan): SparkPlan = {
+ plan.setTagValue(PLAN_ONLY_SETTINGS, snapshot(session.sessionState.conf))
+ plan
+ }
+
+ /**
+ * The settings recorded for this query, searching the adaptive wrappers'
own plans as well:
+ * under AQE the tagged plan is the one AQE was handed, which hangs off
`AdaptiveSparkPlanExec`
+ * rather than appearing among its children.
+ *
+ * Falls back to reading the session when no tag is found, which covers a
plan that reached the
+ * listener without passing through `CometExecRule`.
+ */
+ private def settingsFor(qe: QueryExecution): PlanOnlySettings = {
+ def search(plan: SparkPlan): Option[PlanOnlySettings] =
+ plan
+ .getTagValue(PLAN_ONLY_SETTINGS)
+ .orElse((plan match {
+ case adaptive: AdaptiveSparkPlanExec =>
+ Seq(adaptive.inputPlan, adaptive.initialPlan,
adaptive.executedPlan)
+ case stage: QueryStageExec => Seq(stage.plan)
+ case _ => Seq.empty
+ }).flatMap(search).headOption)
+ .orElse(plan.children.flatMap(search).headOption)
+
+
search(qe.executedPlan).getOrElse(snapshot(qe.sparkSession.sessionState.conf))
+ }
Review Comment:
I think this fallback can be risky.
If a normal query runs with plan-only off, it won't have the tag. If
plan-only gets enabled before the async listener callback runs, this can pick
up the new session config and report that older query as plan-only.
Would it be safer to just skip reporting when the tag is missing?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
