cloud-fan commented on code in PR #58924:
URL: https://github.com/apache/spark/pull/58924#discussion_r4068531005
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -301,7 +302,21 @@ case class AdaptiveSparkPlanExec(
/**
* Run `fun` on finalized physical plan
*/
- def withFinalPlanUpdate[T](fun: SparkPlan => T): T = lock.synchronized {
+ /**
+ * Drive this AQE to completion of all non-result stages without creating a
+ * [[ResultQueryStageExec]] on top. Used by [[CTEReuseQueryStageExec]] for
shared CTE
+ * materialization: all references to a CTE wrap this same inner AQE, whose
materialization runs
+ * exactly once via the lazy `materializeFuture`.
+ */
+ def materialize(): Future[Any] = materializeFuture
+
+ @transient private lazy val materializeFuture: Future[Any] = Future {
Review Comment:
**Blocking (P1):** This Future runs the inner AQE on a reusable pool without
capturing the initiating SQLExecution state. The worker can therefore submit
its shuffle jobs with absent or stale execution ids, job groups, scheduler
properties, and artifact context, so they are misattributed and can escape
query cancellation.
See **Shared repair plan 1** in the review body.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala:
##########
@@ -243,6 +243,18 @@ object StaticSQLConf {
.checkValue(thres => thres > 0 && thres <= 1024, "The threshold must be
in (0,1024].")
.createWithDefault(1024)
+ val CTE_MATERIALIZATION_MAX_THREAD_THRESHOLD =
+ buildStaticConf("spark.sql.cteMaterialization.maxThreadThreshold")
+ .internal()
+ .doc("The size of the AdaptiveSparkPlanExec thread pool used for CTE
materialization " +
+ "and reuse. This pool is isolated from the main QueryStageCreator
thread pool. A " +
+ "relatively large default size of 1024 is chosen to minimize the risk
of deadlocks " +
+ "caused by deeply nested CTEs exhausting the available threads.")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
Review Comment:
**Nit (P3):** This new config entry is missing the standard
`.version("4.4.0")` metadata used by the adjacent entries. Without it,
generated configuration and compatibility metadata cannot report when the
setting became available.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala:
##########
@@ -573,10 +573,41 @@ object LogicalPlanIntegrity {
.orElse(LogicalPlanIntegrity.validateNoDanglingReferences(currentPlan))
.orElse(LogicalPlanIntegrity.validateAggregateExpressions(currentPlan))
.orElse(LogicalPlanIntegrity.validateNullability(currentPlan))
+ .orElse(LogicalPlanIntegrity.validateCTEReuseRelations(currentPlan))
.map(err => s"${err}\nPrevious schema:${previousPlan.output.mkString(",
")}" +
s"\nPrevious plan: ${previousPlan.treeString}")
validation
}
+
+ /**
+ * Validates that CTEReuseRelation nodes are correctly formed.
+ * join must also reside within the same shared subplan. A subplan whose
runtime
+ * filter references an external join would fail when materialized
independently.
+ *
+ * Returns an error message if the check fails, or None if it succeeds.
+ */
+ def validateCTEReuseRelations(plan: LogicalPlan): Option[String] = {
+ if (!plan.containsPattern(CTE_REUSE)) {
+ None
+ } else {
+ val cteReuses = plan.collectWithSubqueries {
Review Comment:
**Blocking (P1):** `collectWithSubqueries` cannot descend into a
`CTEReuseRelation.sharedSubplan` because the relation is a leaf and exposes
that plan only through metadata. Nested same-id reuse relations can therefore
bypass this canonical-equality check; `cteAQERegistry.getOrElseUpdate` then
binds all of them to whichever body was registered first, which can silently
execute the wrong nested CTE.
##########
sql/core/src/test/scala/org/apache/spark/sql/CTEReuseWithAQESuite.scala:
##########
@@ -0,0 +1,558 @@
+/*
+ * 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.spark.sql
+
+import org.apache.spark.sql.catalyst.plans.logical.{CTERelationDef,
CTEReuseRelation}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
AdaptiveSparkPlanHelper}
+import org.apache.spark.sql.execution.adaptive.AQEShuffleReadExec
+import org.apache.spark.sql.execution.adaptive.CTEReuseQueryStageExec
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+/**
+ * End-to-end tests for CTE reuse through AQE: verifies that
+ * CTEReuseQueryStageExec instances share the same inner AQE,
+ * metrics are recorded correctly, and AQE re-optimization
+ * interacts properly with CTE reuse stages.
+ */
+class CTEReuseWithAQESuite
+ extends QueryTest with SharedSparkSession
+ with AdaptiveSparkPlanHelper {
+
+ private val cteReuseConf =
+ "spark.sql.optimizer.replaceCTERefWithCTEReuse.enabled"
+
+ private def withCTEReuseEnabled(f: => Unit): Unit = {
+ withSQLConf(
+ cteReuseConf -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true"
+ )(f)
+ }
+
+ /**
+ * Verify that all [[CTEReuseQueryStageExec]] with the same cteId
+ * share the same inner AQE by identity. This is the core
+ * single-materialization invariant: N refs create N stage wrappers
+ * but all point to the same [[AdaptiveSparkPlanExec]], so the CTE
+ * shuffle is materialized exactly once.
+ */
+ /**
+ * Verify single-materialization: each distinct CTE has exactly one
+ * inner AQE shared by all its refs. `expectedDistinctCTEs` is the
+ * number of CTE definitions (default 1).
+ */
+ private def assertInnerAQEShared(
+ df: DataFrame,
+ expectedDistinctCTEs: Int = 1): Unit = {
+ val executedPlan = df.queryExecution.executedPlan
+ val stages = collectWithSubqueries(executedPlan) {
+ case s: CTEReuseQueryStageExec => s
+ }
+ if (stages.isEmpty) return
Review Comment:
**Non-blocking (P2):** Returning when `stages` is empty makes this assertion
vacuous even though `expectedDistinctCTEs` defaults to 1. Most callers would
therefore stay green if the physical CTE reuse path stopped producing any
stages; only one metrics case has an independent stage-count assertion.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala:
##########
@@ -133,7 +133,7 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends
Rule[LogicalPlan] with
private def pushdownPredicatesAndAttributes(
plan: LogicalPlan,
cteMap: CTEMap): LogicalPlan = plan.transformWithSubqueries {
- case cteDef @ CTERelationDef(child, id, originalPlanWithPredicates, _, _,
_, _) =>
+ case cteDef @ CTERelationDef(child, id, originalPlanWithPredicates, _, _,
_, _, _) =>
Review Comment:
**Blocking (P1):** The new `forcePartitioning` value is ignored when this
rule computes `newAttrSet`. If every reference projects only `a` but the CTE is
pinned by `HashPartitioning(b, N)`, this rule can prune `b`; the later
repartition then contains a dangling input reference and the valid forced plan
fails optimization or planning.
See **Shared repair plan 2** in the review body.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCTERefWithRepartition.scala:
##########
@@ -55,12 +56,42 @@ object ReplaceCTERefWithRepartition extends
Rule[LogicalPlan] {
cteDefs.foreach { cteDef =>
val inlined = replaceWithRepartition(cteDef.child, cteMap)
val withRepartition =
- if (canSkipExtraRepartition(inlined) || cteDef.underSubquery) {
- // If the CTE definition plan itself is a repartition operation or
if it hosts a merged
- // scalar subquery, we do not need to add an extra repartition
shuffle.
- inlined
+ if (cteDef.forceSkipInline) {
Review Comment:
**Non-blocking (P2):** `forcePartitioning` is an independent materialization
contract, but this control flow reads it only when `forceSkipInline` is also
true. A `MATERIALIZED` or multi-reference non-deterministic CTE can survive
inlining with `forceSkipInline = false`, in which case its requested
HashPartitioning is silently discarded.
See **Shared repair plan 2** in the review body.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala:
##########
@@ -856,7 +859,9 @@ object QueryExecution {
(if (subquery) {
Nil
} else {
- Seq(ReuseExchangeAndSubquery)
+ // VerifyCTEReuse runs only on the main query (not per-subquery) after
+ // ReuseExchangeAndSubquery, to verify guaranteed CTE shuffle reuse
held (AQE off).
+ Seq(ReuseExchangeAndSubquery, VerifyCTEReuse(failOnReuseFailure =
false))
Review Comment:
**Non-blocking (P2):** This hard-coded `false` makes
`spark.sql.optimizer.failOnCTEReuseWithoutAQE.enabled` ineffective: even when
callers set it to true, a failed guaranteed reuse only logs. The true fail-fast
path is also unreachable from the new tests, which set only the false mode.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -6869,6 +6869,30 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val REPLACE_CTE_REF_WITH_CTE_REUSE =
+ buildConf("spark.sql.optimizer.replaceCTERefWithCTEReuse.enabled")
+ .internal()
+ .doc("When true, replaces CTE references and repartitions with CTEReuse
nodes " +
Review Comment:
**Nit (P3):** This description is broader than the implementation: only the
`forceSkipInline` branch assigns a nonzero plan-reuse id, so ordinary retained
and user `MATERIALIZED` CTEs keep plain repartitions that the new conversion
ignores. Please describe the forceSkipInline/plan-reuse activation boundary
explicitly.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCTERefWithRepartition.scala:
##########
@@ -55,12 +56,42 @@ object ReplaceCTERefWithRepartition extends
Rule[LogicalPlan] {
cteDefs.foreach { cteDef =>
val inlined = replaceWithRepartition(cteDef.child, cteMap)
val withRepartition =
- if (canSkipExtraRepartition(inlined) || cteDef.underSubquery) {
- // If the CTE definition plan itself is a repartition operation or
if it hosts a merged
- // scalar subquery, we do not need to add an extra repartition
shuffle.
- inlined
+ if (cteDef.forceSkipInline) {
+ cteDef.forcePartitioning match {
+ case Some(h: HashPartitioning) =>
+ // A pinned partitioning: materialize it as a plan-reuse
RepartitionByExpression on
+ // the pinned expressions, so guaranteed CTE reuse can pick it
up. This takes
+ // precedence over `canSkipExtraRepartition` -- the pin states
the partitioning we
+ // must produce.
+ RepartitionByExpression(h.expressions, inlined,
optNumPartitions = None)
Review Comment:
**Non-blocking (P2):** Passing `None` here drops `h.numPartitions`;
`RepartitionByExpression` substitutes the session shuffle-partition count. A
producer requesting `HashPartitioning(keys, N)` therefore receives a different
layout whenever `N` differs from the session default.
See **Shared repair plan 2** in the review body.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala:
##########
@@ -1219,6 +1260,13 @@ object AdaptiveSparkPlanExec {
private[adaptive] val executionContext =
ExecutionContext.fromExecutorService(
ThreadUtils.newDaemonCachedThreadPool("QueryStageCreator", 16))
+ // A separate, larger pool for CTE-reuse inner AQEs. Each inner AQE's
`materialize()` blocks a
+ // pool thread waiting on its child stages; sharing the 16-thread
`QueryStageCreator` pool would
+ // let deeply nested CTEs exhaust it and deadlock, so CTE inner AQEs get
their own pool.
+ private[adaptive] val cteExecutionContext =
ExecutionContext.fromExecutorService(
+ ThreadUtils.newDaemonCachedThreadPool("CTEReuseInnerAQE",
Review Comment:
**Blocking (P1):** A separate pool does not break the recursive dependency
cycle when that pool is itself bounded. With the supported threshold of 1, an
outer inner-AQE occupies the sole worker in `withFinalPlanUpdate`, queues its
nested CTE Future back to this executor, and waits forever for that queued
work. Concurrent parents can create the same starvation at larger thresholds.
See **Shared repair plan 1** in the review body.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/QueryStageExec.scala:
##########
@@ -312,6 +312,78 @@ case class TableCacheQueryStageExec(
override def getRuntimeStatistics: Statistics =
inMemoryTableScan.runtimeStatistics
}
+/**
+ * A query stage that wraps the shared inner [[AdaptiveSparkPlanExec]]
materializing a reused CTE.
+ *
+ * All references to one CTE share a single inner AQE (keyed by cteId in
+ * [[AdaptiveExecutionContext.cteAQERegistry]]); each reference gets its own
stage whose `output`
+ * remaps the inner AQE's (primary reference's) attribute ids into this
reference's id space,
+ * mirroring [[org.apache.spark.sql.execution.exchange.ReusedExchangeExec]].
Execution and stats
+ * delegate to the inner AQE's post-iteration `executedPlan` so that going
through the inner AQE's
+ * own `execute()` (which would re-run iteration and wrap a result stage) is
avoided.
+ */
+case class CTEReuseQueryStageExec(
+ override val id: Int,
+ innerAQE: AdaptiveSparkPlanExec,
+ override val output: Seq[Attribute]) extends QueryStageExec {
Review Comment:
**Non-blocking (P2):** This composite stage has no cancellation path, while
`AdaptiveSparkPlanExec.cleanUpAndThrowException` cancels only
`ExchangeQueryStageExec` instances in the outer tree. If a sibling stage fails
during CTE materialization, the hidden inner AQE can keep submitting and
running work after the query has already thrown.
See **Shared repair plan 1** in the review body.
--
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]
