AveryQi115 commented on code in PR #58924:
URL: https://github.com/apache/spark/pull/58924#discussion_r4462595391
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -6869,6 +6869,33 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val REPLACE_CTE_REF_WITH_CTE_REUSE =
+ buildConf("spark.sql.optimizer.replaceCTERefWithCTEReuse.enabled")
+ .internal()
+ .doc("When true, plan-reuse repartitions (those carrying a nonzero
repartition id, " +
+ "currently produced only for CTE definitions marked forceSkipInline)
are converted to " +
+ "CTEReuse nodes instead of plain Repartition nodes; ordinary retained
or user " +
+ "MATERIALIZED CTEs keep plain repartitions and are unaffected.
CTEReuse enables " +
+ "guaranteed exchange reuse in AQE by sharing a single inner
AdaptiveSparkPlanExec " +
+ "across all references.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(true)
+
+ val FAIL_ON_CTE_REUSE_WITHOUT_AQE =
+ buildConf("spark.sql.optimizer.failOnCTEReuseWithoutAQE.enabled")
+ .internal()
+ .doc("When true, throw an error if guaranteed CTE/subplan shuffle reuse
fails with AQE " +
+ "off, i.e. two or more shuffles tagged with the same cteId survive
after " +
+ "ReuseExchangeAndSubquery (they were not deduplicated into a
ReusedExchangeExec). " +
+ "When false, only log the plan, error, and cteIds. Enabled by default
in tests to " +
+ "catch reuse regressions.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
Review Comment:
You're right that the canonical-equality gate means a cteId-keyed owner
would only merge already-validated bodies — thanks, I'd mis-framed the risk. I
still think the mandatory non-AQE owner (and retiring the fail-open verifier)
is a larger change than fits here: the only current producer of
`forceSkipInline` CTEs is the DSv2 MERGE path, which runs with exchange reuse
on, so the duplicate-materialization case isn't reachable for it today. I'd
like to land this with the AQE-on guarantee + best-effort non-AQE reuse and do
the fail-closed cteId owner as a focused follow-up. If you'd rather it block
this PR, I'll pull it in here.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala:
##########
@@ -1222,6 +1222,34 @@ abstract class SparkStrategies extends
QueryPlanner[SparkPlan] {
f, key, lObj, rObj, lGroup, rGroup, lAttr, rAttr, lOrder, rOrder,
oAttr,
planLater(left), planLater(right)) :: Nil
+ case r: logical.CTEReuseRelation =>
+ // Plan the shared subplan eagerly rather than using planLater, because
+ // CTEReuseExchange is a LeafExecNode whose subplan is not visible via
+ // children. QueryPlanner.collectPlaceholders would never find a
PlanLater
+ // inside the subplan, leaving it unresolved.
+ //
+ // sharedSubplan is a RepartitionByExpression wrapping the CTE body.
Planning it produces a
+ // plain ShuffleExchangeExec; we re-stamp its origin as
LOCAL_SHUFFLE_FOR_CTE(cteId) so the
+ // shuffle is recognized as a reuse boundary
(`isCreatedForSubplanReuse`) by
+ // EnsureRequirements, the AQE-off unwrap path, and VerifyCTEReuse. We
use this shuffle
+ // directly instead of adding a separate exchange -- the Repartition
serves as the logical
+ // plan for the inner AQE, enabling correct re-optimization (broadcast
stages are findable
+ // as descendants, and re-planning preserves the root exchange).
+ val physicalSubplan = SparkStrategies.this.plan(r.sharedSubplan).next()
+ val shuffle = physicalSubplan match {
+ case s: ShuffleExchangeExec => s.copy(shuffleOrigin =
LOCAL_SHUFFLE_FOR_CTE(r.cteId))
Review Comment:
Good catch — fixed. The root shuffle lost its `LOCAL_SHUFFLE_FOR_CTE` origin
because the inner AQE re-plans the bare `sharedSubplan` (a
`RepartitionByExpression`), while the origin was only stamped post-planning in
the `CTEReuseRelation` strategy — so re-planning produced a plain
`REPARTITION_BY_COL` shuffle that `EnsureRequirements` then elided once the
body already satisfied the partitioning.
I moved the stamp into the `RepartitionByExpression` strategy itself: it now
emits `LOCAL_SHUFFLE_FOR_CTE(repartitionId)` for a plan-reuse repartition
(`isForPlanReuse`), so the origin is derived from the logical plan and is
reproduced on every replan. The `CTEReuseRelation` strategy uses that stamped
shuffle directly. Added an AQE regression test (a `REPARTITION(k)` body with
`GROUP BY k`, self-joined) that crashed before.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala:
##########
@@ -69,6 +69,18 @@ trait ShuffleExchangeLike extends Exchange {
*/
def shuffleOrigin: ShuffleOrigin
+ /**
+ * Whether this shuffle was inserted to guarantee CTE (subplan) reuse.
+ *
+ * When true, shuffle-planning rules must treat this exchange as immutable
so that all
+ * references to the same CTE stay canonically equal and dedup into a single
+ * `ReusedExchangeExec`. In particular `EnsureRequirements` must not rewrite
its
+ * partitioning in place; a consumer needing a different partitioning wraps
it in a new
+ * shuffle instead. A shuffle qualifies when it carries the
`LOCAL_SHUFFLE_FOR_CTE`
+ * origin, which is only produced for CTE reuse.
+ */
+ def isCreatedForSubplanReuse: Boolean =
shuffleOrigin.isInstanceOf[LOCAL_SHUFFLE_FOR_CTE]
Review Comment:
Thanks — this only triggers under the opt-in
`spark.sql.shuffle.localPipelined.enabled`, so I'd like to handle it as a
follow-up: teach `AQEEnablePipelinedShuffle` about the `LOCAL_SHUFFLE_FOR_CTE`
boundary (and let `hasRegularPrefix` see through the CTE stage), rather than
expand this PR.
##########
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:
Acknowledged — this is real. It's the same class of issue as subquery-stage
execution: in-flight inner work isn't torn down when the outer query fails. I'd
rather fix it together with the planned subquery-stage cancellation upgrade
than add a CTE-specific teardown here, so I'm tracking it as a follow-up rather
than blocking this PR.
##########
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:
Acknowledged — with `cteMaterialization.maxThreadThreshold=1` a nested
forced CTE can deterministically starve, the same
bounded-pool-with-blocking-waits shape subquery stages already have. I'd like
to address progress-guarantee scheduling together with the broader
subquery-stage work rather than bolt a CTE-specific fix on here; tracking it as
a follow-up.
--
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]
