sunchao commented on code in PR #5514: URL: https://github.com/apache/datafusion-comet/pull/5514#discussion_r3878090680
########## 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: [P2] Retain the action's settings for delayed reporting A caller can enable plan-only, run `collect()`, and restore the setting after the action returns while this asynchronous callback is still queued. This reads the session's later flag and silently drops the report for a query that ran in plan-only mode. Other Comet settings changed in that gap also affect the preview. A gated Spark 4.0.4 listener control produces one report when the flag stays enabled and zero when it is restored before callback delivery. The new tests drain the bus inside `withSQLConf`, so they hide this ordering. Keep the action's eligibility and planning settings with its report instead of re-reading mutable session state. ########## 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: [P2] Preserve reused exchange output IDs when expanding it In a Parquet self-join with exchange reuse enabled, the wrapper can expose fresh IDs such as `k#16` while its shared child produces `k#3`. Returning that child leaves the parent sort referring to the discarded ID; Comet's attribute binder then declines the sort and its consuming join stays on Spark, so the preview understates coverage for otherwise supported work. The actual Spark plan binds before normalization; invoking this unchanged method loses the sort/join bindings with AQE both off and on, while disabling reuse preserves them. Retain or remap the wrapper's output IDs while exposing its subtree for conversion. ########## 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: [P2] Cover RDD actions outside the named SQL callback path Registering only this listener leaves `spark.sql(...).rdd.count()` without a report on Spark 3.4/3.5: the new RDD test captures zero reports and is the sole failure in both the [3.4 job](https://github.com/apache/datafusion-comet/actions/runs/33123655409/job/98701512148) and [3.5 job](https://github.com/apache/datafusion-comet/actions/runs/33123655409/job/98701512065). Those RDD actions do not emit the named SQL completion event this listener requires. The 4.x single-count pass is not a per-action control either: in the Spark 4.0.4 component check, obtaining `df.rdd` emits one report, while two subsequent `count()` actions add none. Please cover the RDD execution lifecycle before promising one report per action. -- 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]
