andygrove commented on code in PR #5514: URL: https://github.com/apache/datafusion-comet/pull/5514#discussion_r3970565258
########## spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala: ########## @@ -0,0 +1,243 @@ +/* + * 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.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.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 action, holding the finished plan, so there is nothing to tell apart. + */ +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) Review Comment: Fixed in 8ccff51ab, by correcting the claim rather than by covering RDD actions — and I reproduced your CI result locally first, on both profiles: ``` === 4.1 === - an RDD action is reported once ... succeeded 1 === 3.5 === - an RDD action is reported once *** FAILED *** ArrayBuffer() had size 0 instead of expected size 1 ``` So the split is exactly as you described. Reporting is driven by `QueryExecutionListener`, which Spark fires from the Dataset action path, so it is per **SQL execution**, not per action. On 4.0+ obtaining `df.rdd` runs a query of its own and is reported at that point; on 3.4/3.5 it is not reported at all; and on both, the RDD actions that follow add nothing because no new SQL execution starts. I did not add a second reporting path. Doing it properly means keying on job starts and then reconciling against this listener so an ordinary query is not reported twice, which is a good deal more machinery than this mode justifies — so the honest move was to stop promising something it does not do. The user guide now says "once per SQL execution" with a paragraph on what that means for RDD work, and the class doc says the same and records why the RDD path is not covered. The test is now *RDD actions are outside the reported path*. It takes three actions on the same RDD and asserts the version-appropriate outcome (one report on 4.0+, none on 3.x), which pins the split rather than leaving it to be rediscovered, and asserts the part that is version-independent: the RDD's own actions contribute nothing. `CometPlanOnlySuite` is 22/22 on both the default 4.1 profile and `-Pspark-3.5`. ########## spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala: ########## @@ -0,0 +1,243 @@ +/* + * 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.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.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 action, holding the finished plan, so there is nothing to tell apart. + */ +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") + } + } + + /** + * Logs the Comet plan Comet would have executed for `qe`. + * + * Nothing here may fail the query, which has finished by this point but whose action would + * still see an exception thrown from a listener. Plan-only mode exists to let a workload be + * assessed without taking on risk, so a plan shape the preview mishandles has to cost the + * report rather than the query. + */ + private def report(qe: QueryExecution): Unit = { + val session = qe.sparkSession + // The listener bus thread has no active session, and the conversion rules read their configs + // from the active one. Without this the preview would be built from default config values. + val previous = SparkSession.getActiveSession + SparkSession.setActiveSession(session) + try { + val conf = session.sessionState.conf + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && isCometLoaded(conf) && Review Comment: Fixed in 8ccff51ab. The settings now travel with the query instead of being read back off the session. `CometExecRule` tags the plan it declines to convert with a `PlanOnlySettings` snapshot — plan-only enabled, Comet loaded, exec enabled — taken synchronously while the query is being planned, so it is inside whatever `withSQLConf` the caller wrapped the action in. `report` reads the tag off the executed plan, searching the adaptive wrappers' own plans too, since under AQE the tagged plan hangs off `AdaptiveSparkPlanExec` rather than appearing among its children. It falls back to the session only when no tag is found, which covers a plan that reached the listener without passing through the rule. **The test was the hard part, and you had already identified why.** My first attempt did what the existing tests do — leave the `withSQLConf` block, then drain — and it passed against the *unfixed* code, because the bus normally drains during `collect()` and the callback sees the flag still on. A test that passes either way is worse than none, so I did not keep it. The version that works forces the ordering. It runs on `spark.newSession()`, which gives a deterministic registration order: `CometPlanOnly` registers its listener lazily on first use, so on a session that has never run a plan-only query a gating `QueryExecutionListener` registered first also runs first, and can park the bus on a latch. The setting is restored while the callback sits in that latch, then the latch opens. Against the previous code it fails as you predicted: ``` List() had size 0 instead of expected size 1 expected the query planned under plan-only mode to still be reported ``` ########## spark/src/main/scala/org/apache/comet/rules/CometPlanOnly.scala: ########## @@ -0,0 +1,243 @@ +/* + * 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.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.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 action, holding the finished plan, so there is nothing to tell apart. + */ +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") + } + } + + /** + * Logs the Comet plan Comet would have executed for `qe`. + * + * Nothing here may fail the query, which has finished by this point but whose action would + * still see an exception thrown from a listener. Plan-only mode exists to let a workload be + * assessed without taking on risk, so a plan shape the preview mishandles has to cost the + * report rather than the query. + */ + private def report(qe: QueryExecution): Unit = { + val session = qe.sparkSession + // The listener bus thread has no active session, and the conversion rules read their configs + // from the active one. Without this the preview would be built from default config values. + val previous = SparkSession.getActiveSession + SparkSession.setActiveSession(session) + try { + val conf = session.sessionState.conf + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && isCometLoaded(conf) && + CometConf.COMET_EXEC_ENABLED.get(conf) && !isMetadataOnly(qe.executedPlan)) { + val preview = previewOf(session, qe.executedPlan) + logWarning(s"$REPORT_PREFIX\n${new ExtendedExplainInfo().generateExtendedInfo(preview)}") + } + } catch { + case NonFatal(e) => + logWarning(s"$REPORT_PREFIX could not build a coverage report for this query", e) + } finally { + previous match { + case Some(session) => SparkSession.setActiveSession(session) + case None => SparkSession.clearActiveSession() + } + } + } + + /** + * Whether `plan` only touches metadata - `CREATE VIEW`, `SHOW TABLES`, `SET`. + * + * There is nothing to accelerate in one, and a session runs enough of them that reporting each + * as 0% would bury the reports worth reading. A command that carries a query below it - `INSERT + * ... SELECT`, `CREATE TABLE AS SELECT`, a V2 append - has that query as a child and is + * reported. + */ + private def isMetadataOnly(plan: SparkPlan): Boolean = plan match { + case _: ExecutedCommandExec | _: CommandResultExec => true + case command: V2CommandExec => command.children.isEmpty + case _ => false + } + + /** + * The plan Comet would have executed for `plan`, which Spark has finished preparing and + * running. + * + * Conversion is only the first half of Comet planning. Spark then inserts the columnar + * transitions and runs Comet's post-columnar rules (see + * `CometSparkSessionExtensions.CometExecColumnar.postColumnarTransitions`), which can revert + * whole stages back to Spark and drop redundant transitions. Those steps run here too, so the + * report describes the plan that would really have executed and counts the transitions that + * would really have been there. + * + * `RevertNativeForTransitionHeavyStages` is applied with `applyToAllStages` because this holds + * a whole plan, whereas under AQE Spark hands that rule one stage at a time. + */ + private def previewOf(session: SparkSession, plan: SparkPlan): SparkPlan = { + val prepared = previewSubqueriesOf(session, stripPreparation(plan)) + val converted = CometExecRule(session)._apply(CometScanRule(session)._apply(prepared)) + val withTransitions = + ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(converted) + val reverted = RevertNativeForTransitionHeavyStages(session).applyToAllStages(withTransitions) + EliminateRedundantTransitions(session).apply(reverted) + } + + /** + * `plan` as the conversion rules would have seen it, with everything Spark added after them + * removed: the adaptive wrappers, the whole-stage codegen wrappers, and the columnar + * transitions. + * + * Taking the plan Spark executed and undoing this much of its preparation is what buys the + * accuracy this mode needs. The alternative, describing the plan as it stood before + * preparation, describes a plan AQE may have replanned beyond recognition: stages coalesced, + * joins switched from sort merge to broadcast, an empty side pruned away. + */ + private def stripPreparation(plan: SparkPlan): SparkPlan = plan match { + // Under AQE the executed plan is a wrapper holding the plan AQE settled on. Its query stages + // hold their own plans off to one side, out of `children`, so an ordinary transform would not + // reach into them. + case adaptive: AdaptiveSparkPlanExec => stripPreparation(adaptive.executedPlan) + case stage: QueryStageExec => stripPreparation(stage.plan) + // A runtime partition-coalescing wrapper over a shuffle stage. It has no counterpart in a plan + // that has not been through AQE, and the conversion rules judge a shuffle by the exchange, so + // it goes with the stage it wraps. + case read: AQEShuffleReadExec => stripPreparation(read.child) + // `ReuseExchangeAndSubquery` is the last thing Spark's preparation does, after the columnar + // rules, so in a real Comet run the exchange behind a `ReusedExchangeExec` has already been + // converted. Here it has not, and the wrapper is a leaf as far as a transform is concerned, so + // conversion would never reach the subtree while the coverage count - which unwraps the + // wrapper - still counts every operator in it as Spark. Undo the reuse and let both copies + // convert, which is what the counts of a real Comet run reflect. + case reused: ReusedExchangeExec => stripPreparation(reused.child) Review Comment: Fixed in 8ccff51ab. `stripPreparation` still undoes the reuse — that part was deliberate and the comment explains why — but it now carries the wrapper's output IDs across: ```scala case reused: ReusedExchangeExec => restoreReusedOutput(reused, stripPreparation(reused.child)) ``` `restoreReusedOutput` is a no-op when the IDs already agree, which is the common case; when they do not, it re-aliases positionally — the same correspondence `ReusedExchangeExec` itself relies on — using a `ProjectExec` of `Alias`es carrying the wrapper's `exprId`s. That is one extra operator in the report per re-aliased reuse, which I judged the better trade: a projection of aliases is something Comet converts, so it costs a point of denominator, whereas losing the sort and the join above it was costing the preview real coverage. The regression test is *a reused exchange keeps its output IDs so consumers still convert*: a self-join over two identical grouped subqueries with `EXCHANGE_REUSE_ENABLED=true` and broadcast disabled, asserting the join converts in the preview and that no `Sort` is left on Spark. Against the previous code it fails, so it reproduces the shape you described rather than just documenting it. I used AQE off for the test since that is the simpler of the two configurations you mentioned; the fix is in `stripPreparation`, which both paths go through. -- 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]
