sunchao commented on code in PR #5394:
URL: https://github.com/apache/datafusion-comet/pull/5394#discussion_r3867069780
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -573,7 +686,93 @@ case class CometExecRule(session: SparkSession)
newPlan
}
- private def _apply(plan: SparkPlan): SparkPlan = {
+ /**
+ * Build the Comet plan we would have executed and log it. Called from
`_apply` in plan-only
+ * mode; the built plan is discarded.
+ */
+ private def reportPlanOnlyCoverage(plan: SparkPlan): Unit = {
+ val preview = buildPreview(plan, topLevel = true)
+ logWarning(s"[Comet plan-only]\n${new
ExtendedExplainInfo().generateExtendedInfo(preview)}")
+ }
+
+ /**
+ * The plan Comet would have executed for `plan`. Passes `forPreview = true`
through the nested
+ * calls so both conversion rules run their normal transforms instead of
short-circuiting.
+ *
+ * Conversion is only the first half of Comet planning. Normally 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, rather than the pre-transition conversion
result.
+ *
+ * `RevertNativeForTransitionHeavyStages` is applied with `applyToAllStages`
because the preview
+ * holds the whole plan at once, whereas under AQE Spark hands that rule one
stage at a time.
+ *
+ * @param topLevel
+ * false when previewing the plan behind a subquery expression.
`ReuseExchangeAndSubquery` is
+ * the last step of Spark's preparation and `QueryExecution.preparations`
omits it for a
+ * subquery, so the preview follows suit.
+ */
+ private def buildPreview(plan: SparkPlan, topLevel: Boolean): SparkPlan = {
+ val converted =
+ _apply(CometScanRule(session)._apply(previewSubqueriesOf(plan)),
forPreview = true)
+ val withTransitions =
+ ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar =
false).apply(converted)
+ val reverted =
RevertNativeForTransitionHeavyStages(session).applyToAllStages(withTransitions)
+ val preview = EliminateRedundantTransitions(session).apply(reverted)
+ if (topLevel) ReuseExchangeAndSubquery.apply(preview) else preview
+ }
+
+ /**
+ * `plan` with the plan behind each of its subquery expressions replaced by
that plan's own
+ * preview.
+ *
+ * Extended explain walks a node's `innerChildren`, which for a `SparkPlan`
are the plans owned
+ * by its expressions, and counts their operators towards the report. Normal
planning has
+ * already converted those plans by the time the outer plan reaches this
rule - Spark prepares a
+ * scalar subquery through the full preparation sequence, columnar rules
included, before
+ * substituting it into the outer plan - so leaving them untouched here
would report every
+ * subquery operator as un-accelerated Spark and understate coverage
relative to what Comet
+ * really executes.
+ *
+ * Each subquery is also reported in its own right, because Spark prepares
it as a top-level
+ * plan of its own; those reports and the counts here therefore describe
overlapping sets of
+ * operators.
+ */
+ private def previewSubqueriesOf(plan: SparkPlan): SparkPlan = {
+ plan.transformAllExpressions { case subquery: ExecSubqueryExpression =>
+ subquery.withNewPlan(previewSubquery(subquery.plan))
+ }
+ }
+
+ private def previewSubquery(subquery: BaseSubqueryExec): BaseSubqueryExec =
subquery match {
+ // Reuse bookkeeping: the plan to preview is one level further down.
+ case reused: ReusedSubqueryExec => reused.copy(child =
previewSubquery(reused.child))
+ case other =>
+ val preview = buildPreview(stripPreparation(other.child), topLevel =
false)
Review Comment:
[P2] Apply stage reversion to the prepared DPP child
Could the DPP broadcast child be previewed through its own preparation
boundary before restoring the exchange wrapper? With AQE disabled, Spark's
`PlanDynamicPruningFilters` prepares the child before constructing
`BroadcastExchangeExec`. This path instead passes that exchange into
`buildPreview`, so `applyToAllStages` treats it as a boundary and leaves the
reconverted child unreverted. I reproduced this on the PR head with a
partitioned Parquet fact table joined to a filtered Parquet dimension,
`spark.comet.exec.transitionRevert.enabled=true`,
`spark.comet.exec.transitionRevert.maxTransitions=0`, and
`spark.comet.exec.project.enabled=false`: the outer report says 4/12
accelerated (33%), while `CometCoverageStats` for the normally executed plan
says 2/12 (16%). Both executions returned the same ten correct rows. The
separate DPP report correctly shows 0/3, but its child is counted as
accelerated again in the outer report. Please preserve the child's normal
preparation scope and add
a non-AQE DPP coverage comparison.
##########
spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala:
##########
@@ -797,4 +799,238 @@ class CometExecRuleSuite extends CometTestBase {
}
}
+ /**
+ * Run `sql` with plan-only mode enabled and assert nothing was offloaded to
native. `useV1`
+ * toggles between `USE_V1_SOURCE_LIST=parquet` (V1 `CometScanExec` path) and
+ * `USE_V1_SOURCE_LIST=""` (V2 `CometBatchScanExec` path).
+ */
+ private def runPlanOnlyAndAssertReverted(
+ sql: String,
+ useV1: Boolean = true,
+ aqe: Boolean = true): Unit = {
+ withSQLConf(
+ SQLConf.USE_V1_SOURCE_LIST.key -> (if (useV1) "parquet" else ""),
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.key -> "true") {
+ val executed = spark.sql(sql).queryExecution.executedPlan
+ val cometNodes = stripAQEPlan(executed).collect { case p: CometPlan => p
}
+ assert(
+ cometNodes.isEmpty,
+ s"plan-only mode must not offload; found Comet operators: $cometNodes")
+ }
+ }
+
+ for {
+ useV1 <- Seq(true, false)
+ aqe <- Seq(true, false)
+ } {
+ val label = s"${if (useV1) "V1" else "V2"} scan, AQE=$aqe"
+ test(s"plan-only mode: $label") {
+ withParquetTable((0 until 100).map(i => (i, i % 5)), "tbl") {
Review Comment:
[P3] Select V2 before creating the Parquet fixture
Could `USE_V1_SOURCE_LIST` be configured outside `withParquetTable`? The
fixture reads Parquet and registers `tbl` before `runPlanOnlyAndAssertReverted`
changes that setting. Changing it afterward does not replace the existing
logical relation, so both tests labeled "V2 scan" still exercise
`FileSourceScanExec`. A Spark 4.1.3 probe confirmed this with AQE on and off;
creating the fixture after the config change produces `BatchScanExec`. Actual
V2 plan-only execution passed separate checks, but these committed tests do not
cover that path.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -115,12 +116,124 @@ object CometExecRule {
*/
val SKIP_COMET_BROADCAST_TAG:
org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit] =
org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometBroadcast")
+
+ /**
+ * A bounded set of keys, used for plan-only reporting state. Evicts in LRU
order once `limit`
+ * keys are held, so a long-lived driver retains a fixed amount of reporting
state. Same
+ * synchronized-`LinkedHashMap` pattern used by
`IcebergPlanDataInjector.commonCache`.
+ */
+ private class BoundedKeySet(limit: Int) {
+ private val keys: java.util.Map[String, java.lang.Boolean] =
+ java.util.Collections.synchronizedMap(
+ new java.util.LinkedHashMap[String, java.lang.Boolean](16, 0.75f,
true) {
+ override def removeEldestEntry(
+ eldest: java.util.Map.Entry[String, java.lang.Boolean]): Boolean
= size() > limit
+ })
+
+ /** Adds `key`, returning true if it was not already present. */
+ def add(key: String): Boolean = keys.put(key, java.lang.Boolean.TRUE) ==
null
+ }
+
+ private val PLAN_ONLY_REPORTED_LIMIT = 1024
+
+ /** `executionId:planFingerprint` keys that plan-only mode has already
reported. */
+ private val planOnlyReportedPlans = new
BoundedKeySet(PLAN_ONLY_REPORTED_LIMIT)
+
+ /**
+ * Set on the root of every plan that plan-only mode has reported. Catalyst
copies a node's tags
+ * onto the node that replaces it (`TreeNode.copyTagsFrom`), so the mark
survives the rewrites
+ * Spark applies between one application of this rule and the next, which is
what lets a later
+ * application recognize a plan it has already described.
+ */
+ private val PLAN_ONLY_REPORTED: TreeNodeTag[Unit] =
TreeNodeTag[Unit]("comet.planOnlyReported")
+
+ /**
+ * Whether `plan` is an application of this rule to a plan already reported,
rather than to a
+ * plan the user is asking about. All the state lives on the plan itself, so
this holds whether
+ * or not the application carries a SQL execution ID - `df.rdd.count()` and
reading
+ * `executedPlan` without an action both plan, and in the first case execute
AQE stages, with no
+ * execution ID set.
+ *
+ * Under AQE one query reaches the conversion rules five times, and only the
first is the plan
+ * being evaluated:
+ *
+ * - the initial plan, through the query-stage-prep rule - the one to
report;
+ * - the same plan again as a columnar rule, now wrapped in
`AdaptiveSparkPlanExec`;
+ * - each query stage as it is created, again as a columnar rule, rooted
at the `Exchange` the
+ * stage was cut at;
+ * - the re-optimized plan after each stage materializes, through the prep
rule;
+ * - the final plan once every stage has materialized, as a columnar rule.
+ *
+ * `AdaptiveSparkPlanExec` is a leaf as far as `exists` is concerned, so the
wrapper is matched
+ * by name rather than through its contents. The two re-optimization passes
hold
+ * `QueryStageExec` nodes for the stages already materialized. The remaining
two - a stage, and
+ * a final plan for a query AQE never had to cut into stages - carry the
mark left by the prep
+ * rule.
+ *
+ * @param queryStagePrep
+ * whether the calling rule instance is registered as a query-stage-prep
rule.
+ * @param aqeEnabled
+ * whether AQE is enabled for the session. A plan rooted at an `Exchange`
is a stage only when
+ * AQE cuts stages; with AQE off it is an ordinary plan
(`df.repartition(n)`, say) and must
+ * still be reported.
+ */
+ private def isReapplication(
+ plan: SparkPlan,
+ queryStagePrep: Boolean,
+ aqeEnabled: Boolean): Boolean = {
+ plan.exists(p =>
+ p.isInstanceOf[QueryStageExec] || p.isInstanceOf[AdaptiveSparkPlanExec]
||
+ p.getTagValue(PLAN_ONLY_REPORTED).isDefined) ||
+ (aqeEnabled && !queryStagePrep && plan.isInstanceOf[Exchange])
Review Comment:
[P2] Keep report state when AQE replaces an empty plan
Could report ownership survive replacement of the entire physical tree? With
plan-only mode and AQE enabled, `spark.sql("SELECT id % 2 AS k, count(*) AS n
FROM range(20) WHERE id < 0 GROUP BY id % 2").collect()` emits an 83% report
for the initial aggregate and then a second 0% report for `EmptyRelation`.
After the empty shuffle materializes, AQE creates a fresh plan containing
neither `QueryStageExec` nor the reporting tag. Its structural hash also
differs, so both guards admit another report under the same execution ID. I
reproduced the two actual Comet warnings on Spark 4.1.3; AQE disabled emits
one. Please retain query/subquery reporting ownership across this rewrite and
cover an empty adaptive query.
--
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]
