This is an automated email from the ASF dual-hosted git repository.
cloud-fan 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 522736869aae [SPARK-57212][SQL][FOLLOWUP] Record AQE rule timing into
the shared tracker via a lock instead of per-node trackers
522736869aae is described below
commit 522736869aae4da5a8a1c6605b15d7d6f51f0d6c
Author: Wenchen Fan <[email protected]>
AuthorDate: Tue Jun 9 13:03:32 2026 -0700
[SPARK-57212][SQL][FOLLOWUP] Record AQE rule timing into the shared tracker
via a lock instead of per-node trackers
### What changes were proposed in this pull request?
This is a follow-up of #56275, addressing
https://github.com/apache/spark/pull/56275#discussion_r3376088031.
The original PR routes physical preparation and AQE rules through
`RuleExecutor` so their per-rule timing lands in `QueryPlanningTracker`. To
avoid concurrent writes to a single tracker from the threads that plan scalar /
IN / DPP subqueries, it gave each `AdaptiveSparkPlanExec` its own
`QueryPlanningTracker`, registered them in
`AdaptiveExecutionContext.planningTrackers`, and folded them all into the
query's shared tracker (`context.qe.tracker`) in `finalPlanUpdate` via a new
`QueryP [...]
This PR replaces that machinery with a lock on the single shared tracker.
The only mutation point is `QueryPlanningTracker.recordRuleInvocation`; making
it (and the `rules` / `topRulesByTime` read accessors) `synchronized` lets
every AQE node -- main and subquery -- record straight into the single
`context.qe.tracker`, which is already reachable from every node (sub-AQEs
reuse the same `AdaptiveExecutionContext`).
That lets the following pieces, which existed solely to avoid concurrent
writes, be deleted:
- the per-`AdaptiveSparkPlanExec` `tracker` field,
- `AdaptiveExecutionContext.planningTrackers`,
- `QueryPlanningTracker.merge`,
- the merge loop in `finalPlanUpdate`, and
- the `if (!isSubquery)` special-case there.
`PhysicalRuleExecutor`, the `applyPhysicalRules` signature change, and the
`withTracker` wrapping from the original PR are unchanged.
### Why are the changes needed?
Beyond the reduced surface area, the lock makes correctness local:
- As written, the merge is safe only because sub-AQEs complete (via
`waitForSubqueries` -> `awaitResult`) before `finalPlanUpdate` reads their
trackers, and because `finalPlanUpdate` is a `lazy val` that runs once. That is
a non-obvious chain. A lock makes the guarantee local to the one mutating
method and independent of when subqueries complete or when `finalPlanUpdate`
runs, so a future change to subquery scheduling or plan finalization cannot
silently reintroduce a race.
- Any future code that records a rule from another thread is automatically
covered by the lock; with per-node trackers each new path would have to wire up
its own tracker plus merge.
- The cost is negligible: recording is one `HashMap` put per rule
invocation, so an uncontended monitor on the single-threaded analyzer /
optimizer path is tens of ns against ms-scale planning, and it is contended
only by the rare concurrent subquery-planning threads.
- A shared tracker reflects AQE rules as they run, rather than only after
`finalPlanUpdate` at execution time.
### Does this PR introduce _any_ user-facing change?
No. This is an internal refactor of how the same per-rule timing is
recorded; `tracker.rules`, `tracker.topRulesByTime(...)`, and
`RuleExecutor.dumpTimeSpent()` cover the same rules as after #56275.
### How was this patch tested?
Existing tests in `QueryPlanningTrackerEndToEndSuite` (added by #56275),
including `SPARK-57212: Track sub-query AQE rules`, which exercises the
cross-thread subquery-recording path, continue to pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 4.8
Closes #56410 from cloud-fan/SPARK-57212-tracker-lock-followup.
Authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../spark/sql/catalyst/QueryPlanningTracker.scala | 36 ++++++++++------------
.../execution/adaptive/AdaptiveSparkPlanExec.scala | 29 +++--------------
.../QueryPlanningTrackerEndToEndSuite.scala | 2 +-
3 files changed, 23 insertions(+), 44 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 9915c03eaa33..2595eebceb9e 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
@@ -72,8 +72,11 @@ object QueryPlanningTracker {
}
/**
- * A thread local variable to implicitly pass the tracker around. This
assumes the query planner
- * is single-threaded, and avoids passing the same tracker context in every
function call.
+ * A thread local variable to implicitly pass the tracker around, avoiding
passing the same
+ * tracker context in every function call. The same tracker may be shared
across threads -- e.g.
+ * AQE plans scalar / IN / DPP subqueries on separate threads, all recording
into the query's
+ * tracker -- so the recording path
([[QueryPlanningTracker.recordRuleInvocation]] and the
+ * reporting accessors) is synchronized.
*/
private val localTracker = new ThreadLocal[QueryPlanningTracker]() {
override def initialValue: QueryPlanningTracker = null
@@ -82,7 +85,7 @@ object QueryPlanningTracker {
/** Returns the current tracker in scope, based on the thread local
variable. */
def get: Option[QueryPlanningTracker] = Option(localTracker.get())
- /** Sets the current tracker for the execution of function f. We assume f is
single-threaded. */
+ /** Sets the current tracker for the execution of function f. */
def withTracker[T](tracker: QueryPlanningTracker)(f: => T): T = {
val originalTracker = localTracker.get()
localTracker.set(tracker)
@@ -194,11 +197,17 @@ class QueryPlanningTracker(
/**
* Record a specific invocation of a rule.
*
+ * `synchronized` because the same tracker can be written from multiple
threads: AQE plans
+ * scalar / IN / DPP subqueries on separate threads, all recording into the
query's tracker.
+ * Recording is one `HashMap` put against ms-scale planning, so the lock is
uncontended on the
+ * single-threaded analyzer / optimizer path and contended only by the rare
concurrent
+ * subquery-planning threads.
+ *
* @param rule name of the rule
* @param timeNs time taken to run this invocation
* @param effective whether the invocation has resulted in a plan change
*/
- def recordRuleInvocation(rule: String, timeNs: Long, effective: Boolean):
Unit = {
+ def recordRuleInvocation(rule: String, timeNs: Long, effective: Boolean):
Unit = synchronized {
var s = rulesMap.get(rule)
if (s eq null) {
s = new RuleSummary
@@ -210,22 +219,11 @@ 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
+ // `rules` and `topRulesByTime` are `synchronized` to read a consistent
snapshot of `rulesMap`
+ // while concurrent subquery-planning threads may still be recording (see
`recordRuleInvocation`).
+ def rules: Map[String, RuleSummary] = synchronized { rulesMap.asScala.toMap }
def phases: Map[String, PhaseSummary] = phasesMap.asScala.toMap
@@ -233,7 +231,7 @@ class QueryPlanningTracker(
* Returns the top k most expensive rules (as measured by time). If k is
larger than the rules
* seen so far, return all the rules. If there is no rule seen so far or k
<= 0, return empty seq.
*/
- def topRulesByTime(k: Int): Seq[(String, RuleSummary)] = {
+ def topRulesByTime(k: Int): Seq[(String, RuleSummary)] = synchronized {
if (k <= 0) {
Seq.empty
} else {
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 f98eaefc8cf7..9a483076ff56 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, ConcurrentLinkedQueue,
LinkedBlockingQueue}
+import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue}
import scala.collection.concurrent.TrieMap
import scala.collection.mutable
@@ -86,11 +86,6 @@ 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
@@ -207,7 +202,7 @@ case class AdaptiveSparkPlanExec(
}
@transient val initialPlan = context.session.withActive {
- QueryPlanningTracker.withTracker(tracker) {
+ QueryPlanningTracker.withTracker(context.qe.tracker) {
applyPhysicalRules(
applyQueryPostPlannerStrategyRules(inputPlan),
queryStagePreparationRules,
@@ -279,9 +274,9 @@ 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 {
- // 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) {
+ // Record this node's planning rules into the query's shared tracker.
Subquery plans are
+ // planned on separate threads, so `recordRuleInvocation` is
synchronized.
+ QueryPlanningTracker.withTracker(context.qe.tracker) {
val executionId = getExecutionId
// Use inputPlan logicalLink here in case some top level physical
nodes may be removed
// during `initialPlan`
@@ -414,12 +409,6 @@ 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)}")
}
@@ -981,14 +970,6 @@ 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/test/scala/org/apache/spark/sql/execution/QueryPlanningTrackerEndToEndSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryPlanningTrackerEndToEndSuite.scala
index 40accf0df4a0..c23d4e354c8a 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
@@ -88,7 +88,7 @@ class QueryPlanningTrackerEndToEndSuite extends StreamTest {
// 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.
+ // the query tracker if every `AdaptiveSparkPlanExec` records into the
shared query tracker.
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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
