This is an automated email from the ASF dual-hosted git repository.
peter-toth pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 098057d96349 [SPARK-57212][SQL] Track preparation and AQE rule timing
in `QueryPlanningTracker`
098057d96349 is described below
commit 098057d963495e91e5093371ee8a3543f709df55
Author: Peter Toth <[email protected]>
AuthorDate: Tue Jun 9 15:25:59 2026 +0200
[SPARK-57212][SQL] Track preparation and AQE rule timing in
`QueryPlanningTracker`
### What changes were proposed in this pull request?
Physical preparation rules and AQE physical rules run through plain
`foldLeft` loops that bypass `RuleExecutor.execute`, so their per-rule timing
is never recorded by `QueryPlanningTracker` and therefore not reported by
`topRulesByTime` or `RuleExecutor.dumpTimeSpent()`. This PR records that timing
by routing those rules through `RuleExecutor`, the same machinery the analyzer
and optimizer already use to track theirs.
- Add `QueryExecution.PhysicalRuleExecutor`, a small
`RuleExecutor[SparkPlan]` that runs a given rule list as a single
`FixedPoint(1)` batch, and route `QueryExecution.prepareForExecution` and
`AdaptiveSparkPlanExec.applyPhysicalRules` through it. `RuleExecutor.execute`
already times each rule, calls `recordRuleInvocation` on the
`QueryPlanningTracker` in scope, and runs `PlanChangeLogger`. `FixedPoint(1)`
runs each rule exactly once and skips the idempotence re-check, since preparati
[...]
- Supply the tracker through `QueryPlanningTracker.withTracker` (a
`ThreadLocal`): `QueryExecution` wraps the `PLANNING` phase, and each
`AdaptiveSparkPlanExec` wraps `initialPlan` and the body of
`withFinalPlanUpdate` with its own tracker. `reOptimize` calls
`optimizer.execute`, which reads the same `ThreadLocal`.
- Give each `AdaptiveSparkPlanExec` -- the main query plan and the plan
built for every scalar / IN / DPP subquery -- its own `QueryPlanningTracker`,
registered in `AdaptiveExecutionContext.planningTrackers`. Sub-AQEs plan on
separate threads (`SubqueryExec` / `BroadcastExchangeExec` pools), but since
each one writes only its own tracker, the per-rule recording path needs no
synchronization. The main (non-subquery) node merges all of these trackers into
the query's `QueryPlanningTrack [...]
With this, `QueryPlanningTracker.rules` and `topRulesByTime` cover physical
preparation rules (`EnsureRequirements`, `CollapseCodegenStages`,
`ReuseExchangeAndSubquery`, ...) and AQE rules of both the main query and its
subqueries (`AdjustShuffleExchangePosition`, `ValidateSparkPlan`,
`OptimizeSkewedJoin`, the AQE re-optimizer rules, ...), alongside the analyzer
and optimizer rules they already track. These rules also appear in
`RuleExecutor.dumpTimeSpent()` / `QueryExecutionMetering` [...]
`AdaptiveSparkPlanExec.optimizeQueryStage`'s per-stage
`queryStageOptimizerRules` (`CoalesceShufflePartitions`,
`OptimizeShuffleWithLocalRead`, ...) are out of scope here: that loop has its
own `AQEShuffleReadRule` validate/rollback handling and does not fit a plain
`RuleExecutor`.
### Why are the changes needed?
A long-running preparation or AQE rule -- for example `EnsureRequirements`
over a key-grouped join with many partitions, or AQE rules applied per stage or
per subquery -- shows up only inside the `planning` phase total, never per
rule, which makes it hard to diagnose where planning time goes.
`QueryExecution.normalize` already records per-rule timing through
`QueryPlanningTracker`; this brings the same per-rule visibility to physical
preparation and AQE. Routing through `RuleExecutor` [...]
### Does this PR introduce _any_ user-facing change?
No query-result or public-API change -- `tracker.rules`,
`tracker.topRulesByTime(...)`, and `RuleExecutor.dumpTimeSpent()` simply gain
entries for physical preparation and AQE rules.
One side effect is worth calling out, as it goes slightly beyond
observability: `RuleExecutor.execute` runs plan-change validation, which the
old `foldLeft` loops did not.
With `spark.sql.lightweightPlanChangeValidation` enabled by default, a
lightweight validation now runs after each effective preparation/AQE rule -- a
small, always-on cost on the planning path, and a new potential
`PLAN_VALIDATION_FAILED_RULE_IN_BATCH` site should a rule ever produce an
invalid plan. Preparation and AQE rules are expected to produce valid plans, so
this is arguably a latent correctness improvement -- but it is more than pure
observability.
### How was this patch tested?
New tests in `QueryPlanningTrackerEndToEndSuite`:
- `SPARK-57212: Track preparation rules` -- a non-AQE query asserts
`EnsureRequirements` and `CollapseCodegenStages` appear in `tracker.rules`.
- `SPARK-57212: Track AQE-internal preparation rules` -- a shuffle query
asserts the AQE-only rules `AdjustShuffleExchangePosition` and
`ValidateSparkPlan` appear.
- `SPARK-57212: Track sub-query AQE rules` -- a query whose only shuffle is
inside a scalar subquery asserts that `DynamicJoinSelection` (an AQE
re-optimizer rule the shuffle-free main plan never runs) appears, confirming
the subquery plan's tracker is merged into the
query tracker.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 4.8
Closes #56275 from peter-toth/SPARK-57212-track-preparation-rules.
Authored-by: Peter Toth <[email protected]>
Signed-off-by: Peter Toth <[email protected]>
---
.../spark/sql/catalyst/QueryPlanningTracker.scala | 13 +
.../spark/sql/execution/QueryExecution.scala | 29 ++-
.../execution/adaptive/AdaptiveSparkPlanExec.scala | 266 +++++++++++----------
.../adaptive/InsertAdaptiveSparkPlan.scala | 5 +-
.../QueryPlanningTrackerEndToEndSuite.scala | 38 +++
5 files changed, 215 insertions(+), 136 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/QueryPlanningTracker.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/QueryPlanningTracker.scala
index d1007404158f..9915c03eaa33 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/QueryPlanningTracker.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/QueryPlanningTracker.scala
@@ -210,6 +210,19 @@ class QueryPlanningTracker(
s.numEffectiveInvocations += (if (effective) 1 else 0)
}
+ /**
+ * Merge the per-rule statistics from `other` into this tracker.
[[RuleSummary]]'s fields are
+ * additive, so the totals are simply accumulated.
+ */
+ def merge(other: QueryPlanningTracker): Unit = {
+ other.rulesMap.forEach { (rule, otherSummary) =>
+ val s = rulesMap.computeIfAbsent(rule, _ => new RuleSummary)
+ s.totalTimeNs += otherSummary.totalTimeNs
+ s.numInvocations += otherSummary.numInvocations
+ s.numEffectiveInvocations += otherSummary.numEffectiveInvocations
+ }
+ }
+
// ------------ reporting functions below ------------
def rules: Map[String, RuleSummary] = rulesMap.asScala.toMap
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index d843af844d94..e37588faf8fb 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
@@ -35,7 +35,7 @@ import org.apache.spark.sql.catalyst.analysis.{Analyzer,
LazyExpression, NamePar
import org.apache.spark.sql.catalyst.expressions.codegen.ByteCodeStats
import org.apache.spark.sql.catalyst.plans.QueryPlan
import org.apache.spark.sql.catalyst.plans.logical.{AppendData, Command,
CommandResult, CompoundBody, CreateTableAsSelect, LogicalPlan,
OverwriteByExpression, OverwritePartitionsDynamic, ReplaceTableAsSelect,
ReturnAnswer, Union, UnresolvedWith, WithCTE}
-import org.apache.spark.sql.catalyst.rules.{PlanChangeLogger, Rule}
+import org.apache.spark.sql.catalyst.rules.{PlanChangeLogger, Rule,
RuleExecutor}
import org.apache.spark.sql.catalyst.transactions.TransactionUtils
import org.apache.spark.sql.catalyst.util.StringUtils.PlanStringConcat
import org.apache.spark.sql.catalyst.util.truncatedString
@@ -357,7 +357,9 @@ class QueryExecution(
val plan = executePhase(QueryPlanningTracker.PLANNING) {
// clone the plan to avoid sharing the plan instance between different
stages like analyzing,
// optimizing and planning.
- QueryExecution.prepareForExecution(preparations, sparkPlan.clone())
+ QueryPlanningTracker.withTracker(tracker) {
+ QueryExecution.prepareForExecution(preparations, sparkPlan.clone())
+ }
}
// Note: For eagerly executed command it might have already been called in
// `eagerlyExecutedCommand` and is a noop here.
@@ -788,14 +790,21 @@ object QueryExecution {
private[execution] def prepareForExecution(
preparations: Seq[Rule[SparkPlan]],
plan: SparkPlan): SparkPlan = {
- val planChangeLogger = new PlanChangeLogger[SparkPlan]()
- val preparedPlan = preparations.foldLeft(plan) { case (sp, rule) =>
- val result = rule.apply(sp)
- planChangeLogger.logRule(rule.ruleName, sp, result)
- result
- }
- planChangeLogger.logBatch("Preparations", plan, preparedPlan)
- preparedPlan
+ new PhysicalRuleExecutor("Preparations", preparations).execute(plan)
+ }
+
+ /**
+ * A [[RuleExecutor]] that applies a fixed list of physical-plan rules once
each (a single
+ * `FixedPoint(1)` batch). Routing physical preparation and AQE rules
through a `RuleExecutor`
+ * reuses its built-in per-rule timing, which records into the active
[[QueryPlanningTracker]]
+ * exactly like the analyzer and optimizer, instead of re-implementing the
timing loop. A
+ * `FixedPoint(1)` strategy (rather than `Once`) is used so each rule runs
exactly once without
+ * triggering the idempotence re-check, as preparation rules are not
necessarily idempotent.
+ */
+ private[execution] class PhysicalRuleExecutor(
+ batchName: String,
+ rules: Seq[Rule[SparkPlan]]) extends RuleExecutor[SparkPlan] {
+ override protected def batches: Seq[Batch] = Seq(Batch(batchName,
FixedPoint(1), rules: _*))
}
/**
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
index 112ee82314c4..f98eaefc8cf7 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
@@ -18,7 +18,7 @@
package org.apache.spark.sql.execution.adaptive
import java.util
-import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue}
+import java.util.concurrent.{ConcurrentHashMap, ConcurrentLinkedQueue,
LinkedBlockingQueue}
import scala.collection.concurrent.TrieMap
import scala.collection.mutable
@@ -32,6 +32,7 @@ import org.apache.spark.internal.LogKeys._
import org.apache.spark.internal.MessageWithContext
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.QueryPlanningTracker
import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, ReturnAnswer}
import org.apache.spark.sql.catalyst.plans.physical.{Distribution,
UnspecifiedDistribution}
@@ -85,6 +86,11 @@ case class AdaptiveSparkPlanExec(
@transient private val optimizer = new AQEOptimizer(conf,
context.session.sessionState.adaptiveRulesHolder.runtimeOptimizerRules)
+ // Each `AdaptiveSparkPlanExec` records the physical-planning rules it runs
into its own tracker.
+ // The main node folds them all into the shared `context.qe.tracker` once
the final plan is ready.
+ @transient private val tracker = new QueryPlanningTracker()
+ context.planningTrackers.add(tracker)
+
// `EnsureRequirements` may remove user-specified repartition and assume the
query plan won't
// change its output partitioning. This assumption is not true in AQE. Here
we check the
// `inputPlan` which has not been processed by `EnsureRequirements` yet, to
find out the
@@ -196,15 +202,17 @@ case class AdaptiveSparkPlanExec(
applyPhysicalRules(
plan,
context.session.sessionState.adaptiveRulesHolder.queryPostPlannerStrategyRules,
- Some((planChangeLogger, "AQE Query Post Planner Strategy Rules"))
+ "AQE Query Post Planner Strategy Rules"
)
}
@transient val initialPlan = context.session.withActive {
- applyPhysicalRules(
- applyQueryPostPlannerStrategyRules(inputPlan),
- queryStagePreparationRules,
- Some((planChangeLogger, "AQE Preparations")))
+ QueryPlanningTracker.withTracker(tracker) {
+ applyPhysicalRules(
+ applyQueryPostPlannerStrategyRules(inputPlan),
+ queryStagePreparationRules,
+ "AQE Preparations")
+ }
}
@volatile private var currentPhysicalPlan = initialPlan
@@ -271,119 +279,124 @@ case class AdaptiveSparkPlanExec(
// `plan.queryExecution.rdd`, we need to set active session here as new
plan nodes can be
// created in the middle of the execution.
context.session.withActive {
- val executionId = getExecutionId
- // Use inputPlan logicalLink here in case some top level physical nodes
may be removed
- // during `initialPlan`
- var currentLogicalPlan = inputPlan.logicalLink.get
- var result = createQueryStages(fun, currentPhysicalPlan, firstRun = true)
- val events = new LinkedBlockingQueue[StageMaterializationEvent]()
- val errors = new mutable.ArrayBuffer[Throwable]()
- var stagesToReplace = Seq.empty[QueryStageExec]
- while (!result.allChildStagesMaterialized) {
- currentPhysicalPlan = result.newPlan
- if (result.newStages.nonEmpty) {
- stagesToReplace = result.newStages ++ stagesToReplace
- executionId.foreach(onUpdatePlan(_, result.newStages.map(_.plan)))
-
- // SPARK-33933: we should submit tasks of broadcast stages first, to
avoid waiting
- // for tasks to be scheduled and leading to broadcast timeout.
- // This partial fix only guarantees the start of materialization for
BroadcastQueryStage
- // is prior to others, but because the submission of collect job for
broadcasting is
- // running in another thread, the issue is not completely resolved.
- val reorderedNewStages = result.newStages
- .sortWith {
- case (_: BroadcastQueryStageExec, _: BroadcastQueryStageExec) =>
false
- case (_: BroadcastQueryStageExec, _) => true
- case _ => false
- }
+ // Record this node's planning rules into its own tracker; merged into
the query's tracker by
+ // the main node in `finalPlanUpdate`.
+ QueryPlanningTracker.withTracker(tracker) {
+ val executionId = getExecutionId
+ // Use inputPlan logicalLink here in case some top level physical
nodes may be removed
+ // during `initialPlan`
+ var currentLogicalPlan = inputPlan.logicalLink.get
+ var result = createQueryStages(fun, currentPhysicalPlan, firstRun =
true)
+ val events = new LinkedBlockingQueue[StageMaterializationEvent]()
+ val errors = new mutable.ArrayBuffer[Throwable]()
+ var stagesToReplace = Seq.empty[QueryStageExec]
+ while (!result.allChildStagesMaterialized) {
+ currentPhysicalPlan = result.newPlan
+ if (result.newStages.nonEmpty) {
+ stagesToReplace = result.newStages ++ stagesToReplace
+ executionId.foreach(onUpdatePlan(_, result.newStages.map(_.plan)))
+
+ // SPARK-33933: we should submit tasks of broadcast stages first,
to avoid waiting
+ // for tasks to be scheduled and leading to broadcast timeout.
+ // This partial fix only guarantees the start of materialization
for BroadcastQueryStage
+ // is prior to others, but because the submission of collect job
for broadcasting is
+ // running in another thread, the issue is not completely resolved.
+ val reorderedNewStages = result.newStages
+ .sortWith {
+ case (_: BroadcastQueryStageExec, _: BroadcastQueryStageExec)
=> false
+ case (_: BroadcastQueryStageExec, _) => true
+ case _ => false
+ }
- // Start materialization of all new stages and fail fast if any
stages failed eagerly
- reorderedNewStages.foreach { stage =>
- try {
- stage.materialize().onComplete { res =>
- if (res.isSuccess) {
- // record shuffle IDs for successful stages for cleanup
- stage.plan.collect {
- case s: ShuffleExchangeLike =>
- context.shuffleIds.put(s.shuffleId, true)
+ // Start materialization of all new stages and fail fast if any
stages failed eagerly
+ reorderedNewStages.foreach { stage =>
+ try {
+ stage.materialize().onComplete { res =>
+ if (res.isSuccess) {
+ // record shuffle IDs for successful stages for cleanup
+ stage.plan.collect {
+ case s: ShuffleExchangeLike =>
+ context.shuffleIds.put(s.shuffleId, true)
+ }
+ events.offer(StageSuccess(stage, res.get))
+ } else {
+ events.offer(StageFailure(stage, res.failed.get))
}
- events.offer(StageSuccess(stage, res.get))
- } else {
- events.offer(StageFailure(stage, res.failed.get))
- }
- // explicitly clean up the resources in this stage
- stage.cleanupResources()
- }(AdaptiveSparkPlanExec.executionContext)
- } catch {
- case e: Throwable =>
- stage.error.set(Some(e))
- cleanUpAndThrowException(Seq(e), Some(stage.id))
+ // explicitly clean up the resources in this stage
+ stage.cleanupResources()
+ }(AdaptiveSparkPlanExec.executionContext)
+ } catch {
+ case e: Throwable =>
+ stage.error.set(Some(e))
+ cleanUpAndThrowException(Seq(e), Some(stage.id))
+ }
}
}
- }
- // Wait on the next completed stage, which indicates new stats are
available and probably
- // new stages can be created. There might be other stages that finish
at around the same
- // time, so we process those stages too in order to reduce re-planning.
- val nextMsg = events.take()
- val rem = new util.ArrayList[StageMaterializationEvent]()
- events.drainTo(rem)
- (Seq(nextMsg) ++ rem.asScala).foreach {
- case StageSuccess(stage, res) =>
- stage.resultOption.set(Some(res))
- case StageFailure(stage, ex) =>
- stage.error.set(Some(ex))
- errors.append(ex)
- }
+ // Wait on the next completed stage, which indicates new stats are
available and probably
+ // new stages can be created. There might be other stages that
finish at around the same
+ // time, so we process those stages too in order to reduce
re-planning.
+ val nextMsg = events.take()
+ val rem = new util.ArrayList[StageMaterializationEvent]()
+ events.drainTo(rem)
+ (Seq(nextMsg) ++ rem.asScala).foreach {
+ case StageSuccess(stage, res) =>
+ stage.resultOption.set(Some(res))
+ case StageFailure(stage, ex) =>
+ stage.error.set(Some(ex))
+ errors.append(ex)
+ }
- // In case of errors, we cancel all running stages and throw exception.
- if (errors.nonEmpty) {
- cleanUpAndThrowException(errors.toSeq, None)
- }
- val testTriggerForceCancellation =
AQETestHelper.shouldForceCancellation(this)
-
- if (!currentPhysicalPlan.isInstanceOf[ResultQueryStageExec]) {
- // Try re-optimizing and re-planning. Adopt the new plan if its cost
is equal to or less
- // than that of the current plan; otherwise keep the current
physical plan together with
- // the current logical plan since the physical plan's logical links
point to the logical
- // plan it has originated from.
- // Meanwhile, we keep a list of the query stages that have been
created since last plan
- // update, which stands for the "semantic gap" between the current
logical and physical
- // plans. And each time before re-planning, we replace the
corresponding nodes in the
- // current logical plan with logical query stages to make it
semantically in sync with
- // the current physical plan. Once a new plan is adopted and both
logical and physical
- // plans are updated, we can clear the query stage list because at
this point the two
- // plans are semantically and physically in sync again.
- var logicalPlan =
replaceWithQueryStagesInLogicalPlan(currentLogicalPlan, stagesToReplace)
- if (testTriggerForceCancellation) {
- // Force unwrap all LogicalQueryStage so they get replanned.
- logicalPlan = logicalPlan.transformDown {
- case LogicalQueryStage(logical, _) => logical
- }
+ // In case of errors, we cancel all running stages and throw
exception.
+ if (errors.nonEmpty) {
+ cleanUpAndThrowException(errors.toSeq, None)
}
- val afterReOptimize = reOptimize(logicalPlan)
- if (afterReOptimize.isDefined) {
- val (newPhysicalPlan, newLogicalPlan) = afterReOptimize.get
- val origCost = costEvaluator.evaluateCost(currentPhysicalPlan)
- val newCost = costEvaluator.evaluateCost(newPhysicalPlan)
- if (newCost < origCost ||
- (newCost == origCost && currentPhysicalPlan != newPhysicalPlan)
||
- testTriggerForceCancellation) {
- lazy val plans = sideBySide(
- currentPhysicalPlan.treeString,
newPhysicalPlan.treeString).mkString("\n")
- logOnLevel(log"Plan changed:\n${MDC(QUERY_PLAN, plans)}")
- cleanUpTempTags(newPhysicalPlan)
- currentPhysicalPlan = newPhysicalPlan
- currentLogicalPlan = newLogicalPlan
- stagesToReplace = Seq.empty[QueryStageExec]
+ val testTriggerForceCancellation =
AQETestHelper.shouldForceCancellation(this)
+
+ if (!currentPhysicalPlan.isInstanceOf[ResultQueryStageExec]) {
+ // Try re-optimizing and re-planning. Adopt the new plan if its
cost is equal to or less
+ // than that of the current plan; otherwise keep the current
physical plan together with
+ // the current logical plan since the physical plan's logical
links point to the logical
+ // plan it has originated from.
+ // Meanwhile, we keep a list of the query stages that have been
created since last plan
+ // update, which stands for the "semantic gap" between the current
logical and physical
+ // plans. And each time before re-planning, we replace the
corresponding nodes in the
+ // current logical plan with logical query stages to make it
semantically in sync with
+ // the current physical plan. Once a new plan is adopted and both
logical and physical
+ // plans are updated, we can clear the query stage list because at
this point the two
+ // plans are semantically and physically in sync again.
+ var logicalPlan =
+ replaceWithQueryStagesInLogicalPlan(currentLogicalPlan,
stagesToReplace)
+ if (testTriggerForceCancellation) {
+ // Force unwrap all LogicalQueryStage so they get replanned.
+ logicalPlan = logicalPlan.transformDown {
+ case LogicalQueryStage(logical, _) => logical
+ }
+ }
+ val afterReOptimize = reOptimize(logicalPlan)
+ if (afterReOptimize.isDefined) {
+ val (newPhysicalPlan, newLogicalPlan) = afterReOptimize.get
+ val origCost = costEvaluator.evaluateCost(currentPhysicalPlan)
+ val newCost = costEvaluator.evaluateCost(newPhysicalPlan)
+ if (newCost < origCost ||
+ (newCost == origCost && currentPhysicalPlan !=
newPhysicalPlan) ||
+ testTriggerForceCancellation) {
+ lazy val plans = sideBySide(
+ currentPhysicalPlan.treeString,
newPhysicalPlan.treeString).mkString("\n")
+ logOnLevel(log"Plan changed:\n${MDC(QUERY_PLAN, plans)}")
+ cleanUpTempTags(newPhysicalPlan)
+ currentPhysicalPlan = newPhysicalPlan
+ currentLogicalPlan = newLogicalPlan
+ stagesToReplace = Seq.empty[QueryStageExec]
+ }
+ }
+ if (testTriggerForceCancellation) {
+ AQETestHelper.markForcedCancellationTriggeredForPlan(this)
}
}
- if (testTriggerForceCancellation) {
- AQETestHelper.markForcedCancellationTriggeredForPlan(this)
- }
+ // Now that some stages have finished, we can try creating new
stages.
+ result = createQueryStages(fun, currentPhysicalPlan, firstRun =
false)
}
- // Now that some stages have finished, we can try creating new stages.
- result = createQueryStages(fun, currentPhysicalPlan, firstRun = false)
}
}
_isFinalPlan = true
@@ -401,6 +414,12 @@ case class AdaptiveSparkPlanExec(
if (shouldUpdatePlan) {
getExecutionId.foreach(onUpdatePlan(_, Seq.empty))
}
+ // The main query node folds every `AdaptiveSparkPlanExec`'s tracker (its
own and all sub-query
+ // plans, which were planned on separate threads) into the shared query
tracker. This runs on
+ // the main query thread once the final plan is ready and all sub-queries
have completed.
+ if (!isSubquery) {
+ context.planningTrackers.asScala.foreach(context.qe.tracker.merge)
+ }
logOnLevel(log"Final plan:\n${MDC(QUERY_PLAN, currentPhysicalPlan)}")
}
@@ -667,7 +686,7 @@ case class AdaptiveSparkPlanExec(
val optimizedRootPlan = applyPhysicalRules(
optimizeQueryStage(plan, isFinalStage = true),
postStageCreationRules(supportsColumnar),
- Some((planChangeLogger, "AQE Post Stage Creation")))
+ "AQE Post Stage Creation")
val resultStage = ResultQueryStageExec(currentStageId, optimizedRootPlan,
resultHandler)
currentStageId += 1
setLogicalLinkForNewQueryStage(resultStage, plan)
@@ -681,7 +700,7 @@ case class AdaptiveSparkPlanExec(
val newPlan = applyPhysicalRules(
optimized,
postStageCreationRules(outputsColumnar = plan.supportsColumnar),
- Some((planChangeLogger, "AQE Post Stage Creation")))
+ "AQE Post Stage Creation")
if (e.isInstanceOf[ShuffleExchangeLike]) {
if (!newPlan.isInstanceOf[ShuffleExchangeLike]) {
throw SparkException.internalError(
@@ -810,7 +829,7 @@ case class AdaptiveSparkPlanExec(
val newPlan = applyPhysicalRules(
applyQueryPostPlannerStrategyRules(sparkPlan),
preprocessingRules ++ queryStagePreparationRules,
- Some((planChangeLogger, "AQE Replanning")))
+ "AQE Replanning")
// When both enabling AQE and DPP, `PlanAdaptiveDynamicPruningFilters`
rule will
// add the `BroadcastExchangeExec` node manually in the DPP subquery,
@@ -939,19 +958,8 @@ object AdaptiveSparkPlanExec {
def applyPhysicalRules(
plan: SparkPlan,
rules: Seq[Rule[SparkPlan]],
- loggerAndBatchName: Option[(PlanChangeLogger[SparkPlan], String)] =
None): SparkPlan = {
- if (loggerAndBatchName.isEmpty) {
- rules.foldLeft(plan) { case (sp, rule) => rule.apply(sp) }
- } else {
- val (logger, batchName) = loggerAndBatchName.get
- val newPlan = rules.foldLeft(plan) { case (sp, rule) =>
- val result = rule.apply(sp)
- logger.logRule(rule.ruleName, sp, result)
- result
- }
- logger.logBatch(batchName, plan, newPlan)
- newPlan
- }
+ batchName: String): SparkPlan = {
+ new QueryExecution.PhysicalRuleExecutor(batchName, rules).execute(plan)
}
}
@@ -973,6 +981,14 @@ case class AdaptiveExecutionContext(session: SparkSession,
qe: QueryExecution) {
new TrieMap[SparkPlan, ExchangeQueryStageExec]()
val shuffleIds: ConcurrentHashMap[Int, Boolean] = new ConcurrentHashMap[Int,
Boolean]()
+
+ /**
+ * The per-`AdaptiveSparkPlanExec` planning trackers for this query: the
main query plan and every
+ * sub-query plan, which are planned on separate threads. The main node
merges these into
+ * `qe.tracker` once its final plan is ready (see
`AdaptiveSparkPlanExec.finalPlanUpdate`).
+ */
+ val planningTrackers: ConcurrentLinkedQueue[QueryPlanningTracker] =
+ new ConcurrentLinkedQueue[QueryPlanningTracker]()
}
/**
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/InsertAdaptiveSparkPlan.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/InsertAdaptiveSparkPlan.scala
index d8f850ab2189..835505a5eaec 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/InsertAdaptiveSparkPlan.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/InsertAdaptiveSparkPlan.scala
@@ -78,7 +78,10 @@ case class InsertAdaptiveSparkPlan(
val preprocessingRules = Seq(
planSubqueriesRule)
// Run pre-processing rules.
- val newPlan = AdaptiveSparkPlanExec.applyPhysicalRules(plan,
preprocessingRules)
+ val newPlan = AdaptiveSparkPlanExec.applyPhysicalRules(
+ plan,
+ preprocessingRules,
+ "AQE Preprocessing")
logDebug(s"Adaptive execution enabled for plan: $plan")
AdaptiveSparkPlanExec(newPlan, adaptiveExecutionContext,
preprocessingRules, isSubquery)
} catch {
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryPlanningTrackerEndToEndSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryPlanningTrackerEndToEndSuite.scala
index aac68e638980..40accf0df4a0 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryPlanningTrackerEndToEndSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryPlanningTrackerEndToEndSuite.scala
@@ -18,6 +18,7 @@
package org.apache.spark.sql.execution
import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream,
StreamExecution}
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.StreamTest
class QueryPlanningTrackerEndToEndSuite extends StreamTest {
@@ -58,6 +59,43 @@ class QueryPlanningTrackerEndToEndSuite extends StreamTest {
StopStream)
}
+ test("SPARK-57212: Track preparation rules") {
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = spark.range(1000).selectExpr("count(*)")
+ df.collect()
+ val ruleNames = df.queryExecution.tracker.rules.keySet
+ assert(ruleNames.contains(classOf[
+ org.apache.spark.sql.execution.exchange.EnsureRequirements].getName))
+ assert(ruleNames.contains(classOf[
+ org.apache.spark.sql.execution.CollapseCodegenStages].getName))
+ }
+ }
+
+ test("SPARK-57212: Track AQE-internal preparation rules") {
+ // Run a query with a shuffle so AQE creates a query stage and runs
preparation rules.
+ val df = spark.range(1000).repartition(4).selectExpr("count(*)")
+ df.collect()
+ val ruleNames = df.queryExecution.tracker.rules.keySet
+ // AQE-only rules in queryStagePreparationRules; only reachable via
+ // AdaptiveSparkPlanExec.applyPhysicalRules.
+ assert(ruleNames.contains(
+ "org.apache.spark.sql.execution.adaptive.AdjustShuffleExchangePosition"))
+ assert(ruleNames.contains(
+ "org.apache.spark.sql.execution.adaptive.ValidateSparkPlan"))
+ }
+
+ test("SPARK-57212: Track sub-query AQE rules") {
+ // The main query has no shuffle, so its AQE never re-optimizes and thus
never runs the AQE
+ // logical optimizer. The scalar sub-query does have a shuffle, so its
sub-AQE (planned on a
+ // separate thread) re-optimizes and runs `DynamicJoinSelection`. That
rule is only visible in
+ // the query tracker if every `AdaptiveSparkPlanExec`'s tracker is merged
back into it.
+ val df = spark.sql("SELECT id FROM range(10) WHERE id > (SELECT count(*)
FROM range(1000))")
+ df.collect()
+ val ruleNames = df.queryExecution.tracker.rules.keySet
+ assert(ruleNames.contains(
+ "org.apache.spark.sql.execution.adaptive.DynamicJoinSelection"))
+ }
+
test("The start times should be in order: parsing <= analysis <=
optimization <= planning") {
val df = spark.sql("select count(*) from range(1)")
df.queryExecution.executedPlan
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
