cloud-fan commented on code in PR #58924:
URL: https://github.com/apache/spark/pull/58924#discussion_r4414361112
##########
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:
Thanks for explaining. ReplaceRepartitionWithCTEReuse already seals a cteId
only after validateCTEReuseRelations confirms that every same-id sharedSubplan
is canonically equal, so a cteId-keyed owner on the non-AQE path would merge
only bodies that have passed that check. Today, on any non-AQE route (AQE
disabled, AQE enabled but not applied, or the CTEAdaptiveNotSupportedException
fallback), a nondeterministic forced CTE can still be materialized once per
reference, for example with spark.sql.exchange.reuse=false, and VerifyCTEReuse
only logs by default. That is the failure this PR sets out to fix, so I'd still
like the non-AQE path to either share one owner per cteId or fail before
execution, rather than run duplicate materializations by default.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4237377074","thread_id":"inline:4237377074","verdict_sha256":"7971a94f9214bdf1857a39473550a560c323a2594d2e01d88f4a127ddd3021af"}
-->
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCTERefWithRepartition.scala:
##########
@@ -53,14 +54,48 @@ object ReplaceCTERefWithRepartition extends
Rule[LogicalPlan] {
case WithCTE(child, cteDefs) =>
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
- } else {
- RepartitionByExpression(Seq.empty, inlined, None)
- }
+ // `forcePartitioning` is only honored together with
`forceSkipInline`. A pin requests a
+ // sealed, guaranteed-reuse materialization, which only applies to
definitions `InlineCTE`
+ // keeps (i.e. `forceSkipInline`). A pin on a non-`forceSkipInline`
definition is ignored
+ // here, because such a definition may be inlined by `InlineCTE`
before this rule runs, so
+ // honoring the pin would be inconsistent with the inlining decision.
+ val withRepartition = (cteDef.forceSkipInline,
cteDef.forcePartitioning) match {
+ case (true, Some(h: HashPartitioning)) =>
+ // A pinned partitioning materializes the CTE as a plan-reuse
RepartitionByExpression
+ // on the pinned expressions, so guaranteed CTE reuse can pick it
up. It takes
+ // precedence over `canSkipExtraRepartition` -- the pin states the
exact partitioning
+ // (keys and number of partitions) we must produce.
+ RepartitionByExpression(
+ h.expressions, inlined, optNumPartitions = Some(h.numPartitions))
+ .addRepartitionId(reassign = true)
+ case (true, Some(other)) =>
+ // forcePartitioning is an internal dev API; only HashPartitioning
is supported
+ // today (see CTERelationDef.forcePartitioning).
+ throw new UnsupportedOperationException(
+ s"CTERelationDef.forcePartitioning supports only
HashPartitioning, but got " +
+ s"${other.getClass.getSimpleName}")
+ case (true, None) =>
+ // A forced CTE must always carry the plan-reuse identity so every
reference seals into
+ // a CTEReuseRelation and materializes once, whatever the body's
root shape. Stamp an
+ // existing reusable repartition in place (no extra shuffle); wrap
any other root -- a
+ // plain repartition, a rebalance, or a merged-subquery body -- in
a plan-reuse
+ // RepartitionByExpression.
+ inlined match {
+ case r: PlanReusableRepartition => r.addRepartitionId()
+ case _ =>
+ RepartitionByExpression(Seq.empty, inlined, None)
Review Comment:
**Non-blocking (P2):** This wrapper is added to every forced CTE whose root
is not a `RepartitionByExpression`, whatever
`spark.sql.optimizer.replaceCTERefWithCTEReuse.enabled` says, because only
`ReplaceRepartitionWithCTEReuse` reads the flag. With the flag off nothing is
sealed, but a forced CTE rooted at `REPARTITION(3)` or `REBALANCE` still gets
an extra round-robin shuffle over the user's shuffle. Measured locally against
the merge base, that is 2 live shuffles instead of 1 for a single reference,
and 4 instead of 3 for a self-join. So turning the flag off does not roll the
change back. The existing flag-off test uses a non-forced CTE, so it cannot see
this.
See **Shared repair plan 1** in the review body.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/CTEReuseExchange.scala:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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.execution.exchange
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder}
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+/**
+ * A physical leaf node representing a reference to a shared CTE subplan. The
[[subplan]]
+ * (a [[ShuffleExchangeExec]] over the CTE body, produced from the Repartition
in
+ * [[CTEReuseRelation.sharedSubplan]]) is stored as metadata, not as a tree
child, so
+ * AQE's bottom-up stage creation does not traverse into it and redundantly
plan its
+ * substages before CTEReuseExchange is resolved.
+ *
+ * Consumed before execution:
+ * - AQE on: `createNonResultQueryStages` replaces this with
[[CTEReuseQueryStageExec]],
+ * building a shared inner [[AdaptiveSparkPlanExec]] keyed by `cteId`.
+ * - AQE off: `UnwrapCTEReuseExchange` replaces each ref with its
`LOCAL_SHUFFLE_FOR_CTE`
+ * shuffle (tagged with `cteId`); the final `ReuseExchangeAndSubquery` then
reuses the
+ * canonically-equal copies via [[ReusedExchangeExec]], and
`VerifyCTEReuse` confirms it held.
+ *
+ * @param cteId Unique identifier for the CTE definition to be reused.
+ * @param subplan The [[ShuffleExchangeLike]] over the shared subplan. Its
logical link
+ * points to the Repartition in
[[CTEReuseRelation.sharedSubplan]] so the
+ * inner AQE can re-optimize correctly.
+ * @param innerAQE The shared inner [[AdaptiveSparkPlanExec]] for this
`cteId`, attached by
+ * `PlanCTEReuse` (AQE on only) from
`AdaptiveExecutionContext.cteAQERegistry`.
+ * It is auxiliary metadata carried on the node; the registry,
shared across the
+ * main query and its subqueries, remains the source of truth.
`None` before
+ * `PlanCTEReuse` runs and on the AQE-off path.
+ */
+case class CTEReuseExchange(
+ cteId: Long,
+ subplan: ShuffleExchangeLike,
+ innerAQE: Option[AdaptiveSparkPlanExec] = None
+) extends LeafExecNode {
Review Comment:
**Non-blocking (P2):** `CTEReuseExchange` is a `LeafExecNode` that keeps the
CTE shuffle in a field, so `InsertAdaptiveSparkPlan.shouldApplyAQE` finds no
`Exchange` in a plan whose only shuffles are inside forced CTEs. I checked this
locally: `WITH c AS (...) SELECT * FROM c UNION ALL SELECT * FROM c` runs under
`AdaptiveSparkPlanExec` on the merge base but not with this PR. Such queries
lose AQE for the whole plan, including inside the CTE body, and fall back to
the best-effort non-AQE reuse path, so the PR's AQE-on guarantee does not hold
for them: with `spark.sql.exchange.reuse=false`, the two references still saw
different values from a nondeterministic body. Treating `CTEReuseExchange` (or
the `CTE_REUSE_EXCHANGE` pattern) as an exchange in `shouldApplyAQE` would keep
these queries adaptive.
**Verification:**
- **Regression:** With AQE enabled, a UNION ALL of two references to a
forced CTE produces an adaptive plan whose CTE reuse query stages share one
inner adaptive plan.
- **Compatibility:** Queries without forced CTEs get the same AQE
applicability decision as before.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/CTEReuseExchange.scala:
##########
@@ -0,0 +1,86 @@
+/*
+ * 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.execution.exchange
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder}
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.catalyst.trees.TreePattern._
+import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+/**
+ * A physical leaf node representing a reference to a shared CTE subplan. The
[[subplan]]
+ * (a [[ShuffleExchangeExec]] over the CTE body, produced from the Repartition
in
+ * [[CTEReuseRelation.sharedSubplan]]) is stored as metadata, not as a tree
child, so
+ * AQE's bottom-up stage creation does not traverse into it and redundantly
plan its
+ * substages before CTEReuseExchange is resolved.
+ *
+ * Consumed before execution:
+ * - AQE on: `createNonResultQueryStages` replaces this with
[[CTEReuseQueryStageExec]],
+ * building a shared inner [[AdaptiveSparkPlanExec]] keyed by `cteId`.
+ * - AQE off: `UnwrapCTEReuseExchange` replaces each ref with its
`LOCAL_SHUFFLE_FOR_CTE`
+ * shuffle (tagged with `cteId`); the final `ReuseExchangeAndSubquery` then
reuses the
+ * canonically-equal copies via [[ReusedExchangeExec]], and
`VerifyCTEReuse` confirms it held.
+ *
+ * @param cteId Unique identifier for the CTE definition to be reused.
+ * @param subplan The [[ShuffleExchangeLike]] over the shared subplan. Its
logical link
+ * points to the Repartition in
[[CTEReuseRelation.sharedSubplan]] so the
+ * inner AQE can re-optimize correctly.
+ * @param innerAQE The shared inner [[AdaptiveSparkPlanExec]] for this
`cteId`, attached by
+ * `PlanCTEReuse` (AQE on only) from
`AdaptiveExecutionContext.cteAQERegistry`.
+ * It is auxiliary metadata carried on the node; the registry,
shared across the
+ * main query and its subqueries, remains the source of truth.
`None` before
+ * `PlanCTEReuse` runs and on the AQE-off path.
+ */
+case class CTEReuseExchange(
+ cteId: Long,
+ subplan: ShuffleExchangeLike,
+ innerAQE: Option[AdaptiveSparkPlanExec] = None
Review Comment:
**Nit (P3):** Nothing in Apache Spark reads `innerAQE`. Stage creation looks
the inner AQE up in `cteAQERegistry`, canonicalization clears the field, and
the only reader the comments name is 'the Photon compiler'. For no effect here,
it costs a copy of every `CTEReuseExchange` on each preprocessing pass and a
canonicalization special case. Please drop it, or give it an in-tree reader.
**Verification:**
- **Compile:** No code refers to the removed field.
- **Behavior:** AQE CTE reuse tests still show that all references to a CTE
share one inner adaptive plan.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/QueryStageExec.scala:
##########
@@ -312,6 +312,84 @@ 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 {
+
+ override val plan: SparkPlan = innerAQE
Review Comment:
**Non-blocking (P2):** `doExecute` and `doExecuteColumnar` read
`innerAQE.executedPlan`, but `executeCollect`, `executeTake`, `executeTail` and
`executeToIterator` are still inherited from `QueryStageExec` and call the same
methods on `plan = innerAQE`. That re-drives the shared inner AQE through
`withFinalPlanUpdate` and puts a `ResultQueryStageExec` at its root, so every
other reference to the CTE breaks. Reproduced locally with AQE on: `SELECT id
FROM c UNION ALL SELECT c.id FROM c JOIN (SELECT id FROM side WHERE id > 100) e
ON c.id = e.id` fails with `[INTERNAL_ERROR] Expected ShuffleQueryStageExec at
the root of CTE inner AQE, but got ResultQueryStageExec`. Here AQE empties the
join branch, which makes the CTE stage itself the result stage. With the flag
off, the same query returns its rows. Overriding these four methods to use
`innerAQE.executedPlan` would close the gap.
**Verification:**
- **Regression:** With AQE on, a UNION ALL of a forced CTE reference and a
join of another reference with a side that is empty at runtime returns the
first branch's rows on the first and on a repeated collect, and the shared
inner adaptive plan still ends in a shuffle query stage.
- **Behavior:** Calling executeCollect, executeTake, executeTail and
executeToIterator directly on a CTE reuse stage returns the expected rows and
leaves the statistics of every other reference to the same CTE readable.
##########
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:
**Non-blocking (P2):** `isCreatedForSubplanReuse` documents this shuffle as
immutable for shuffle-planning rules, but `AQEEnablePipelinedShuffle` never
checks it. Its `hasRegularPrefix` also stops at `CTEReuseQueryStageExec.plan`,
a leaf `AdaptiveSparkPlanExec`, so the rule cannot see the regular CTE shuffle
below a CTE stage. Reproduced locally in a single-executor channel-shuffle
session with `spark.sql.shuffle.localPipelined.enabled=true`: a forced CTE
self-join pipelines both join exchanges above the CTE stages. The first collect
works, but collecting the same DataFrame again fails with 'A job that mixes a
pipelined shuffle dependency with a regular shuffle dependency is only
supported when every regular shuffle is a fully-materialized prefix below the
pipelined shuffles'. With the reuse flag off, nothing is pipelined.
**Verification:**
- **Regression:** In a single-executor local session with the pipelined
shuffle enabled and AQE on, a forced CTE self-join returns the same rows on its
first and on repeated collects of the same DataFrame, its inner root is a
regular shuffle query stage, and no outer exchange above a CTE stage is
pipelined.
- **Compatibility:** Plans without CTE reuse are still rewritten to
pipelined exactly as before.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PlanCTEReuse.scala:
##########
@@ -0,0 +1,138 @@
+/*
+ * 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.execution.exchange
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{QUERY_ID, QUERY_PLAN}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.SparkPlan
+
+/**
+ * Unwraps [[CTEReuseExchange]] nodes early in the preparation pipeline (AQE
off).
+ *
+ * Each [[CTEReuseExchange]] is replaced by its `subplan` (a
[[ShuffleExchangeExec]] with a
+ * `LOCAL_SHUFFLE_FOR_CTE(cteId)` origin over the CTE body). The `cteId` is
carried by the origin
+ * itself (set in `SparkStrategies` when the shuffle is built), so
[[VerifyCTEReuse]] can later
+ * confirm the copies were reused without relying on a tree-node tag -- the
origin survives node
+ * reconstructions (e.g. photonization rebuilding the shuffle) that would drop
a tag.
+ *
+ * After unwrapping, all preparation rules (EnsureRequirements, columnar
transitions, codegen,
+ * etc.) process each CTE subplan copy normally, in its own consumer context.
+ * The `LOCAL_SHUFFLE_FOR_CTE` shuffle serves as a protected boundary:
`EnsureRequirements` /
+ * `EnsureRequirementsDP` treat it as immutable (see
`ShuffleExchangeLike.isCreatedForSubplanReuse`)
+ * and add consumer-dependent operators above it, so the subtree below stays
canonically equal
+ * across copies. The final `ReuseExchangeAndSubquery` then reuses the copies
by canonical equality.
+ *
+ * Runs in both main and subquery preparation batches so CTE references inside
subqueries are also
+ * unwrapped.
+ */
+case object UnwrapCTEReuseExchange extends Rule[SparkPlan] with Logging {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Plain transformUp (not transformUpWithSubqueries) is sufficient:
subquery plans are prepared
+ // separately by PlanSubqueries, which calls
QueryExecution.prepareExecutedPlan ->
+ // preparations(subquery = true), so this rule is applied recursively to
each subquery's plan
+ // on its own. We only need to handle CTEReuseExchange nodes in the
current plan tree here.
+ plan.transformUp {
+ case cte: CTEReuseExchange =>
+ // Recursively unwrap nested CTEReuseExchange nodes inside the subplan
+ // (e.g., cte_inner inside cte_outer's body).
+ val unwrapped = apply(cte.subplan)
+ unwrapped match {
+ case s: ShuffleExchangeLike if
s.shuffleOrigin.isInstanceOf[LOCAL_SHUFFLE_FOR_CTE] =>
+ // The cteId is already carried by the LOCAL_SHUFFLE_FOR_CTE
origin (SparkStrategies set
+ // it from the cteId in CTEReuseRelation); nothing extra to attach
here.
+ logDebug(log"CTE reuse (AQE off): unwrapped CTEReuseExchange for "
+
+ log"cteId=${MDC(QUERY_ID, cte.cteId)}")
Review Comment:
**Nit (P3):** `MDC(QUERY_ID, cte.cteId)` records a `RepartitionIdGenerator`
counter as a query id. `UnwrapCTEReuseExchange`'s error branch logs a node
name, and `VerifyCTEReuse` a cteId summary, under `QUERY_PLAN`. Structured-log
consumers that search those keys get the wrong values. A dedicated key such as
`CTE_ID`, with `QUERY_PLAN` kept for plan strings, would fix it.
**Verification:**
- **Inspection:** Every MDC in the CTE reuse rules labels its value with a
key that matches what the value is.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PlanCTEReuse.scala:
##########
@@ -0,0 +1,138 @@
+/*
+ * 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.execution.exchange
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{QUERY_ID, QUERY_PLAN}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.SparkPlan
+
+/**
+ * Unwraps [[CTEReuseExchange]] nodes early in the preparation pipeline (AQE
off).
+ *
+ * Each [[CTEReuseExchange]] is replaced by its `subplan` (a
[[ShuffleExchangeExec]] with a
+ * `LOCAL_SHUFFLE_FOR_CTE(cteId)` origin over the CTE body). The `cteId` is
carried by the origin
+ * itself (set in `SparkStrategies` when the shuffle is built), so
[[VerifyCTEReuse]] can later
+ * confirm the copies were reused without relying on a tree-node tag -- the
origin survives node
+ * reconstructions (e.g. photonization rebuilding the shuffle) that would drop
a tag.
+ *
+ * After unwrapping, all preparation rules (EnsureRequirements, columnar
transitions, codegen,
+ * etc.) process each CTE subplan copy normally, in its own consumer context.
+ * The `LOCAL_SHUFFLE_FOR_CTE` shuffle serves as a protected boundary:
`EnsureRequirements` /
+ * `EnsureRequirementsDP` treat it as immutable (see
`ShuffleExchangeLike.isCreatedForSubplanReuse`)
+ * and add consumer-dependent operators above it, so the subtree below stays
canonically equal
+ * across copies. The final `ReuseExchangeAndSubquery` then reuses the copies
by canonical equality.
+ *
+ * Runs in both main and subquery preparation batches so CTE references inside
subqueries are also
+ * unwrapped.
+ */
+case object UnwrapCTEReuseExchange extends Rule[SparkPlan] with Logging {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Plain transformUp (not transformUpWithSubqueries) is sufficient:
subquery plans are prepared
+ // separately by PlanSubqueries, which calls
QueryExecution.prepareExecutedPlan ->
+ // preparations(subquery = true), so this rule is applied recursively to
each subquery's plan
+ // on its own. We only need to handle CTEReuseExchange nodes in the
current plan tree here.
+ plan.transformUp {
Review Comment:
**Blocking (P1):** This `transformUp`, and `VerifyCTEReuse`'s
`foreachWithSubqueries` below, run without tree-pattern pruning in every
non-AQE preparation. So any query that AQE does not wrap pays two extra
full-plan traversals even when it has no forced CTE. That covers AQE off, plans
`InsertAdaptiveSparkPlan` skips, and every streaming micro-batch, since
`IncrementalExecution` inherits this rule list. Neither traversal existed
before, and ordinary queries gain nothing from them.
`transformUpWithPruning(_.containsPattern(CTE_REUSE_EXCHANGE))`, plus running
the verifier only when something was unwrapped, would restore the previous cost.
**Verification:**
- **Regression:** For a non-AQE plan without forced CTEs, the CTE unwrap and
verification rules visit no nodes beyond the root pattern check.
- **Compatibility:** Forced CTE references in the main query, in subqueries
and in nested bodies are still unwrapped and reused.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PlanCTEReuse.scala:
##########
@@ -0,0 +1,138 @@
+/*
+ * 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.execution.exchange
+
+import org.apache.spark.SparkException
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{QUERY_ID, QUERY_PLAN}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.SparkPlan
+
+/**
+ * Unwraps [[CTEReuseExchange]] nodes early in the preparation pipeline (AQE
off).
+ *
+ * Each [[CTEReuseExchange]] is replaced by its `subplan` (a
[[ShuffleExchangeExec]] with a
+ * `LOCAL_SHUFFLE_FOR_CTE(cteId)` origin over the CTE body). The `cteId` is
carried by the origin
+ * itself (set in `SparkStrategies` when the shuffle is built), so
[[VerifyCTEReuse]] can later
+ * confirm the copies were reused without relying on a tree-node tag -- the
origin survives node
+ * reconstructions (e.g. photonization rebuilding the shuffle) that would drop
a tag.
+ *
+ * After unwrapping, all preparation rules (EnsureRequirements, columnar
transitions, codegen,
+ * etc.) process each CTE subplan copy normally, in its own consumer context.
+ * The `LOCAL_SHUFFLE_FOR_CTE` shuffle serves as a protected boundary:
`EnsureRequirements` /
+ * `EnsureRequirementsDP` treat it as immutable (see
`ShuffleExchangeLike.isCreatedForSubplanReuse`)
+ * and add consumer-dependent operators above it, so the subtree below stays
canonically equal
+ * across copies. The final `ReuseExchangeAndSubquery` then reuses the copies
by canonical equality.
+ *
+ * Runs in both main and subquery preparation batches so CTE references inside
subqueries are also
+ * unwrapped.
+ */
+case object UnwrapCTEReuseExchange extends Rule[SparkPlan] with Logging {
Review Comment:
**Nit (P3):** This file defines only `UnwrapCTEReuseExchange` and
`VerifyCTEReuse`. The rule actually named `PlanCTEReuse` lives in the sibling
`adaptive/PlanCTEReuse.scala`, so two files with the same name hold unrelated
halves of the feature. Renaming this one, for example to
`UnwrapCTEReuseExchange.scala` or `CTEReuseRules.scala`, would make both easy
to find.
**Verification:**
- **Inspection:** No two files with the same name and unrelated contents
implement the two halves of the CTE reuse feature.
- **Compile:** Every reference to the AQE-off CTE reuse rules still resolves
after the rename.
##########
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:
**Non-blocking (P2):** The `LOCAL_SHUFFLE_FOR_CTE` origin is stamped only on
this initially planned shuffle. When a stage inside the CTE body materializes,
the inner AQE re-plans the root from the plain `RepartitionByExpression`, which
gets an ordinary origin. `EnsureRequirements` then removes a
`REPARTITION_BY_COL` or `REPARTITION_BY_NUM` hash shuffle that the body already
satisfies, and `CTEReuseQueryStageExec` finds no shuffle stage at the root.
Reproduced locally with AQE on: `WITH c AS (SELECT /*+ REPARTITION(k) */ k,
count(*) AS cnt FROM t GROUP BY k) SELECT ... FROM c c1 JOIN c c2 ON c1.k =
c2.k` fails with `[INTERNAL_ERROR] Expected ShuffleQueryStageExec at the root
of CTE inner AQE, but got HashAggregateExec`. The same body pinned with
`forcePartitioning` on `k` fails the same way. With the flag off, and on the
merge base, both return the right rows. Bodies whose root survives still lose
the CTE origin after the first replan; it became `REBALANCE_PARTITIONS_BY_NONE`.
**Recommended change:** Make the shared inner adaptive plan keep its root
reuse shuffle on every replan. When the inner plan's input is a
LOCAL_SHUFFLE_FOR_CTE shuffle, re-stamp the replanned root shuffle with the
same LOCAL_SHUFFLE_FOR_CTE(cteId) origin before the preparation rules run. If a
replan no longer has a shuffle at the root, reject it and keep the current
plan. Add AQE tests for pinned and hash-hinted forced CTEs whose body contains
a stage.
**Why this works:** Because the origin is restored before EnsureRequirements
runs, its optimize-out rule, which matches only REPARTITION_BY_COL/NUM origins,
cannot remove the root, and the AQE read rules keep excluding it.
CTEReuseQueryStageExec therefore always sees a shuffle-stage root. This follows
the existing restoration of a root BroadcastExchangeLike in reOptimize.
**Scope:** Keep the CTE reuse root intact across inner-AQE replans.
**Compatibility:** Forced CTE bodies without stages, and adaptive plans
without CTEs, replan exactly as before.
**Risks:** A replan that would have been adopted for a cheaper body plan is
rejected when it removes the root shuffle, so the CTE keeps one shuffle that
the merge target might have optimized out.
**Constraints:** The restored root must keep the sealed partitioning and the
cteId.
**Exclusions:**
- A root that collapses to an empty relation, which is raised separately as
a question.
**Success:** Under AQE, a forced CTE pinned on, or hinted to repartition by,
its body's grouping key materializes through a shuffle stage and returns
correct results. Inner-AQE replans never change the origin of the root shuffle.
**Verification:**
- **Regression:** With AQE on, a self-joined forced CTE with
forcePartitioning on its GROUP BY key and the default partition count completes
with the expected rows, and its inner adaptive plan ends in a shuffle query
stage.
- **Regression:** The same holds for an unpinned forced CTE whose body
carries a REPARTITION hint on its grouping key.
- **Behavior:** After every inner replan, the root shuffle's origin is still
the CTE reuse origin for that cteId.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/InsertAdaptiveSparkPlan.scala:
##########
@@ -48,7 +48,8 @@ case class InsertAdaptiveSparkPlan(
override def apply(plan: SparkPlan): SparkPlan = applyInternal(plan, false)
- private def applyInternal(plan: SparkPlan, isSubquery: Boolean): SparkPlan =
plan match {
+ // Not private: `PlanCTEReuse` calls this to wrap a `CTEReuseExchange`'s
subplan in an inner AQE.
+ def applyInternal(plan: SparkPlan, isSubquery: Boolean): SparkPlan = plan
match {
Review Comment:
**Nit (P3):** `applyInternal` goes from private to public for
`PlanCTEReuse`, which is in the same package, so `private[adaptive]` would do.
`AdaptiveSparkPlanExec` also gains a public `materialize()`, and a public
`skipResultStage` flag on `withFinalPlanUpdate[T]`. With that flag set, the
method never calls `fun` and returns `().asInstanceOf[T]`, so a caller outside
this package can get a mistyped value. All callers are in
`org.apache.spark.sql.execution.adaptive`.
**Verification:**
- **Compile:** Code outside the adaptive package cannot call applyInternal,
materialize, or a skip-result variant of withFinalPlanUpdate.
- **Behavior:** AQE CTE reuse tests still show shared CTE materialization
working.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -7195,6 +7195,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")
Review Comment:
**Nit (P3):** The three new keys sit in unrelated namespaces
(`spark.sql.optimizer.replaceCTERefWithCTEReuse.enabled`,
`spark.sql.optimizer.failOnCTEReuseWithoutAQE.enabled` and
`spark.sql.cteMaterialization.maxThreadThreshold`), and this boolean puts a
verb before `.enabled`. That goes against the naming guideline in
`ConfigEntry.scala`. The keys are internal and unreleased, so they can still
move under one feature namespace cheaply.
**Verification:**
- **Inspection:** Every new CTE reuse config key sits under one namespace,
and every boolean key is either the umbrella enabled switch or reads as a verb.
- **Compile:** No code, test or documentation still refers to an old key
name.
##########
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
+ * reusing one CTE share the same inner AdaptiveSparkPlanExec, and that 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 the single-materialization invariant: all
[[CTEReuseQueryStageExec]] stages reusing one
+ * CTE point to the same inner AdaptiveSparkPlanExec by identity, so a CTE
referenced N times
+ * creates N stage wrappers but is materialized exactly once.
`expectedDistinctCTEs` is the number
+ * of distinct CTE definitions (default 1).
+ */
+ private def assertInnerAQEShared(
+ df: DataFrame,
+ expectedDistinctCTEs: Int = 1,
+ allowEmpty: Boolean = false): Unit = {
+ val executedPlan = df.queryExecution.executedPlan
+ val stages = collectWithSubqueries(executedPlan) {
+ case s: CTEReuseQueryStageExec => s
+ }
+ if (stages.isEmpty) {
+ // Only callers that legitimately prune all CTE stages (e.g.
empty-relation propagation)
+ // may pass here; otherwise a missing-stages regression must fail, not
pass vacuously.
+ assert(allowEmpty,
+ s"Expected CTEReuseQueryStageExec stages but found
none:\n${executedPlan.treeString}")
+ return
+ }
+ val distinctAQEs = stages.map(s =>
+ System.identityHashCode(s.innerAQE)).toSet
+ assert(distinctAQEs.size == expectedDistinctCTEs,
+ s"Expected $expectedDistinctCTEs distinct inner AQE(s) " +
+ s"for ${stages.size} CTEReuseQueryStageExec, " +
+ s"but found ${distinctAQEs.size}.\n" +
+ s"Plan:\n${executedPlan.treeString}")
+ }
+
+ private def collectCTEReuseRelations(
+ df: DataFrame): Seq[CTEReuseRelation] = {
+ df.queryExecution.optimizedPlan.collectWithSubqueries {
+ case r: CTEReuseRelation => r
+ }
+ }
+
+ /**
+ * Runs `text` with every CTE forced to materialize (`forceSkipInline =
true`) so it opts into
+ * guaranteed CTE shuffle reuse. `forceSkipInline` has no SQL syntax, so tag
the analyzed plan's
+ * CTE definitions and re-run the tagged plan.
+ */
+ private def sqlWithForcedCTEReuse(text: String): DataFrame = {
+ val analyzed = spark.sql(text).queryExecution.analyzed
+ val tagged = analyzed.transformWithSubqueries {
+ case d: CTERelationDef => d.copy(forceSkipInline = true)
+ }
+ classic.Dataset.ofRows(spark, tagged)
+ }
+
+ // -----------------------------------------------------------------
+ // Inner-AQE sharing tests
+ // -----------------------------------------------------------------
+
+ test("CTE reuse: 2 refs share one inner AQE") {
+ withCTEReuseEnabled {
+ withTable("metric_src") {
+ sql("CREATE TABLE metric_src (id INT) USING parquet")
+ sql(
+ "INSERT INTO metric_src VALUES (1), (2), (3), (4), (5)")
+ val df = sqlWithForcedCTEReuse(
+ """WITH cte AS (
+ | SELECT id, rand() as r FROM metric_src
+ |)
+ |SELECT * FROM cte c1 JOIN cte c2 ON c1.id = c2.id
+ |""".stripMargin)
+ df.collect()
+ // Two refs create two CTEReuseQueryStageExec wrappers, both sharing
one inner AQE.
+ assert(collectWithSubqueries(df.queryExecution.executedPlan) {
+ case s: CTEReuseQueryStageExec => s
+ }.size == 2)
+ assertInnerAQEShared(df)
+ }
+ }
+ }
+
+ test("CTE reuse when subplan produces empty result") {
+ withCTEReuseEnabled {
+ withTable("metric_empty_src") {
+ sql("CREATE TABLE metric_empty_src (id INT) USING parquet")
+ sql("INSERT INTO metric_empty_src VALUES (1), (2), (3)")
+ val df = sqlWithForcedCTEReuse(
+ """WITH cte AS (
+ | SELECT id, rand() as r
+ | FROM metric_empty_src WHERE id > 100
+ |)
+ |SELECT * FROM cte c1 JOIN cte c2 ON c1.id = c2.id
+ |""".stripMargin)
+ checkAnswer(df, Seq.empty)
+ // Any surviving CTE stages (empty-relation propagation may prune
some) share one inner AQE.
+ assertInnerAQEShared(df, allowEmpty = true)
+ }
+ }
+ }
+
+ // -----------------------------------------------------------------
+ // AQE empty relation propagation
+ // -----------------------------------------------------------------
+
+ test("AQE empty relation propagation eliminates CTE consumers") {
+ withCTEReuseEnabled {
+ withSQLConf(
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ ("org.apache.spark.sql.catalyst.optimizer" +
+ ".PropagateEmptyRelation")
+ ) {
+ withTable("cte_nonempty", "join_side") {
+ sql(
+ "CREATE TABLE cte_nonempty (id INT) USING parquet")
+ sql(
+ "INSERT INTO cte_nonempty VALUES (1), (2), (3), (4), (5)")
+ sql("CREATE TABLE join_side (id INT) USING parquet")
+ sql("INSERT INTO join_side VALUES (1), (2), (3)")
+ val df = sqlWithForcedCTEReuse(
+ """WITH cte AS (
+ | SELECT id, rand() as r FROM cte_nonempty
+ |)
+ |SELECT c1.id FROM cte c1
+ |INNER JOIN (
+ | SELECT id FROM join_side WHERE id > 100
+ |) e1 ON c1.id = e1.id
+ |UNION ALL
+ |SELECT c2.id FROM cte c2
+ |INNER JOIN (
+ | SELECT id FROM join_side WHERE id > 100
+ |) e2 ON c2.id = e2.id
+ |""".stripMargin)
+ checkAnswer(df, Seq.empty)
+ val aqe = df.queryExecution.executedPlan
+ .asInstanceOf[AdaptiveSparkPlanExec]
+ val finalPlan = aqe.executedPlan
+
+ // 1. All CTEReuseQueryStageExec removed from final plan.
+ val remaining = finalPlan.collect {
+ case s: CTEReuseQueryStageExec => s
+ }
+ assert(remaining.isEmpty,
+ "Expected all CTEReuseQueryStageExec removed " +
+ "after empty relation propagation, " +
+ s"but found ${remaining.size}." +
+ s"\nFinal plan:\n${finalPlan.treeString}")
+
+ // 2. Inner AQE was created (registry entry exists).
+ val registry = aqe.context.cteAQERegistry
+ assert(registry.nonEmpty,
+ "Expected cteAQERegistry to contain inner AQE")
+ }
+ }
+ }
+ }
+
+ // -----------------------------------------------------------------
+ // q24-style: CTE in main query + subquery
+ // -----------------------------------------------------------------
+
+ test("q24-style: main query + HAVING subquery") {
+ withCTEReuseEnabled {
+ withTable("sales", "items") {
+ sql(
+ "CREATE TABLE sales (item_id INT, amount INT) USING parquet")
+ sql(
+ """INSERT INTO sales VALUES
+ |(1, 10), (1, 20), (2, 30), (2, 40),
+ |(3, 50), (3, 60)""".stripMargin)
+ sql(
+ "CREATE TABLE items (id INT, color STRING) USING parquet")
+ sql(
+ """INSERT INTO items VALUES
+ |(1, 'red'), (2, 'blue'), (3, 'red')""".stripMargin)
+
+ val df = sqlWithForcedCTEReuse(
+ """WITH ssales AS (
+ | SELECT s.item_id, i.color,
+ | sum(s.amount) as total, rand() as r
+ | FROM sales s JOIN items i ON s.item_id = i.id
+ | GROUP BY s.item_id, i.color
+ |)
+ |SELECT item_id, sum(total) as paid
+ |FROM ssales
+ |WHERE color = 'red'
+ |GROUP BY item_id
+ |HAVING sum(total) > (
+ | SELECT 0.05 * avg(total) FROM ssales
+ |)
+ |""".stripMargin)
+
+ val reuses = collectCTEReuseRelations(df)
+ assert(reuses.size >= 2,
+ s"Expected >= 2 CTEReuseRelation (main + subquery), " +
+ s"got ${reuses.size}.\nPlan:\n" +
+ df.queryExecution.optimizedPlan.treeString)
+ df.collect()
+ assertInnerAQEShared(df)
+ }
+ }
+ }
+
+ test("q24-style: main query + scalar subquery") {
+ withCTEReuseEnabled {
+ withTable("orders", "products") {
+ sql(
+ "CREATE TABLE orders (pid INT, qty INT) USING parquet")
+ sql(
+ """INSERT INTO orders VALUES
+ |(1, 5), (2, 10), (3, 15),
+ |(1, 20), (2, 25)""".stripMargin)
+ sql(
+ "CREATE TABLE products (id INT, cat STRING) USING parquet")
+ sql(
+ """INSERT INTO products VALUES
+ |(1, 'A'), (2, 'B'), (3, 'A')""".stripMargin)
+
+ val df = sqlWithForcedCTEReuse(
+ """WITH agg AS (
+ | SELECT o.pid, p.cat,
+ | sum(o.qty) as total, rand() as r
+ | FROM orders o JOIN products p ON o.pid = p.id
+ | GROUP BY o.pid, p.cat
+ |)
+ |SELECT pid, total FROM agg a1
+ |WHERE cat = 'A'
+ |AND total > (SELECT avg(total) FROM agg)
+ |""".stripMargin)
+
+ val reuses = collectCTEReuseRelations(df)
+ assert(reuses.size >= 2,
+ s"Expected >= 2 CTEReuseRelation (main + subquery), " +
+ s"got ${reuses.size}.\nPlan:\n" +
+ df.queryExecution.optimizedPlan.treeString)
+ df.collect()
+ assertInnerAQEShared(df)
+ }
+ }
+ }
+
+ // -----------------------------------------------------------------
+ // AQE multi-iteration replan
+ // -----------------------------------------------------------------
+
+ test("AQE multi-iteration replan preserves CTE reuse") {
+ // Multiple CTE refs across joins force multiple AQE iterations.
+ // Disable broadcast to force shuffle joins. Verifies that
+ // re-optimization across iterations does not break sharing.
+ //
+ // Iter 0 logical plan -- all leaves are CTEReuseRelation:
+ // Project [a, b, a2, b2]
+ // +- Join Inner, (a = b2)
+ // :- Join Inner, (a = a2)
+ // : :- Join Inner, (a = b)
+ // : : :- CTEReuseRelation cteId=1 (cte_a ref 1)
+ // : : +- CTEReuseRelation cteId=2 (cte_b ref 1)
+ // : +- CTEReuseRelation cteId=1 (cte_a ref 2)
+ // +- CTEReuseRelation cteId=2 (cte_b ref 2)
+ //
+ // Iter 0 physical plan -- all CTEReuseRelation converted to
+ // CTEReuseQueryStageExec in createQueryStages:
+ // SortMergeJoin [a], [b2], Inner
+ // :- SortMergeJoin [a], [a2], Inner
+ // : :- SortMergeJoin [a], [b], Inner
+ // : : :- Exchange ENSURE_REQUIREMENTS
+ // : : : +- CTEReuseQueryStage 0 (innerAQE-1)
+ // : : +- Exchange ENSURE_REQUIREMENTS
+ // : : +- CTEReuseQueryStage 1 (innerAQE-2)
+ // : +- Exchange ENSURE_REQUIREMENTS
+ // : +- CTEReuseQueryStage 2 (innerAQE-1)
+ // +- Exchange ENSURE_REQUIREMENTS
+ // +- CTEReuseQueryStage 3 (innerAQE-2)
+ //
+ // Iter 0 re-optimize input -- all CTEReuseRelation wrapped in
+ // LogicalQueryStage with runtime stats visible:
+ // Project [a, b, a2, b2]
+ // +- Join Inner, (a = b2)
+ // :- Join Inner, (a = a2)
+ // : :- Join Inner, (a = b)
+ // : : :- LogicalQueryStage CTEReuseRelation cteId=1,
+ // : : : CTEReuseQueryStage 0
+ // : : +- LogicalQueryStage CTEReuseRelation cteId=2,
+ // : : CTEReuseQueryStage 1
+ // : +- LogicalQueryStage CTEReuseRelation cteId=1,
+ // : CTEReuseQueryStage 2
+ // +- LogicalQueryStage CTEReuseRelation cteId=2,
+ // CTEReuseQueryStage 3
+ withCTEReuseEnabled {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_LOG_LEVEL.key -> "INFO",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "5"
+ ) {
+ withTable("replan_a", "replan_b") {
+ sql(
+ "CREATE TABLE replan_a (a INT, ka INT) USING parquet")
+ sql(
+ "CREATE TABLE replan_b (b INT, kb INT) USING parquet")
+ (1 to 50).foreach { i =>
+ sql(s"INSERT INTO replan_a VALUES ($i, $i)")
+ }
+ (1 to 30).foreach { i =>
+ sql(s"INSERT INTO replan_b VALUES ($i, $i)")
+ }
+
+ // CTE A referenced twice (c1 + c3), CTE B twice (c2 + c4).
+ // Join c1-c2 on a=b, then join with c3 on a=a, then join
+ // with c4 on a=b. Forces 2+ AQE iterations.
+ val df = sqlWithForcedCTEReuse(
+ """WITH
+ | cte_a AS (
+ | SELECT a, ka, rand() as r FROM replan_a
+ | ),
+ | cte_b AS (
+ | SELECT b, kb, rand() as r FROM replan_b
+ | )
+ |SELECT c1.a, c2.b, c3.a as a2, c4.b as b2
+ |FROM cte_a c1
+ |JOIN cte_b c2 ON c1.a = c2.b
+ |JOIN cte_a c3 ON c1.a = c3.a
+ |JOIN cte_b c4 ON c1.a = c4.b
+ |""".stripMargin)
+
+ val reuses = collectCTEReuseRelations(df)
+ assert(reuses.size >= 4,
+ s"Expected >= 4 CTEReuseRelation, " +
+ s"got ${reuses.size}.\nPlan:\n" +
+ df.queryExecution.optimizedPlan.treeString)
+
+ df.collect()
+ assertInnerAQEShared(df, expectedDistinctCTEs = 2)
+ }
+ }
+ }
+ }
+
+ // -----------------------------------------------------------------
+ // Nested CTEs
+ // -----------------------------------------------------------------
+
+ test("nested CTEs: outer CTE body references inner CTE") {
+ // Inner CTE (cte_inner) is defined inside outer CTE (cte_outer)'s
+ // body. Both have rand() to prevent inlining. The outer CTE is
+ // referenced twice in the main query; the inner CTE is referenced
+ // twice inside the outer CTE's definition. This produces 2
+ // distinct CTE materializations in cteAQERegistry.
+ withCTEReuseEnabled {
+ withTable("nested_src") {
+ sql(
+ "CREATE TABLE nested_src (id INT, v INT) USING parquet")
+ sql(
+ """INSERT INTO nested_src VALUES
+ |(1, 10), (2, 20), (3, 30),
+ |(4, 40), (5, 50)""".stripMargin)
+
+ val df = sqlWithForcedCTEReuse(
+ """WITH
+ | cte_inner AS (
+ | SELECT id, v, rand() as ri FROM nested_src
+ | ),
+ | cte_outer AS (
+ | SELECT a.id, a.v + b.v as total, rand() as ro
+ | FROM cte_inner a JOIN cte_inner b ON a.id = b.id
+ | )
+ |SELECT c1.id, c2.id as id2
+ |FROM cte_outer c1
+ |JOIN cte_outer c2 ON c1.id = c2.id
+ |""".stripMargin)
+
+ val reuses = collectCTEReuseRelations(df)
+ // Only the outer CTE's 2 refs are visible at the top level.
+ // The inner CTE's refs live inside the outer CTE's
+ // sharedSubplan and materialize inside the outer CTE's
+ // inner AQE.
+ assert(reuses.size >= 2,
+ s"Expected >= 2 CTEReuseRelation (outer CTE refs)" +
+ s", got ${reuses.size}.\nPlan:\n" +
+ df.queryExecution.optimizedPlan.treeString)
+
+ df.collect()
+ // collectWithSubqueries descends into QueryStageExec.plan
+ // and AdaptiveSparkPlanExec.executedPlan, so it finds both
+ // outer CTE stages (2) and inner CTE stages inside the
+ // outer's inner AQE. 2 distinct inner AQEs total.
+ assertInnerAQEShared(df, expectedDistinctCTEs = 2)
+ // Both outer and inner CTE are in the shared registry.
+ val aqe = df.queryExecution.executedPlan
+ .asInstanceOf[AdaptiveSparkPlanExec]
+ val registry = aqe.context.cteAQERegistry
+ assert(registry.size == 2,
+ s"Expected 2 entries in cteAQERegistry " +
+ s"(outer + inner CTE), got ${registry.size}")
+ }
+ }
+ }
+
+ // -----------------------------------------------------------------
+ // Deeply nested CTEs (6 levels)
+ // -----------------------------------------------------------------
+
+ // -----------------------------------------------------------------
+ // LocalPartition CTE feeding a global aggregate
+ // -----------------------------------------------------------------
+
+ test("global aggregate directly over a LocalPartition CTE") {
Review Comment:
**Nit (P3):** This test is built on a `LocalPartition` partitioning and a
`RepartitionForCTEStage` hook, and neither exists in this codebase. Here an
unpinned forced CTE seals with
`RoundRobinPartitioning(spark.sql.shuffle.partitions)`. With the default 200
partitions, that never satisfies `AllTuples`, so `EnsureRequirements` always
adds an exchange above the CTE stage, and `localReads.isEmpty` cannot fail for
the regression the comment describes. `checkAnswer` is the only live check.
**Verification:**
- **Inspection:** The test mentions no LocalPartition,
RepartitionForCTEStage or internal tracker, and its plan assertion would fail
if the guarded plan property changed.
##########
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)
Review Comment:
Thanks for updating the description. It now matches the default-on setting,
which resolves the inconsistency raised here. Some LocalPartition and Photon
references remain in code comments and tests; this review raises them
separately.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4237377076","thread_id":"inline:4237377076","verdict_sha256":"7971a94f9214bdf1857a39473550a560c323a2594d2e01d88f4a127ddd3021af"}
-->
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala:
##########
@@ -1196,6 +1196,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:
Thanks. The merge target also planned each inlined copy of the CTE body
separately, so per-reference planning is not a regression from this PR, and
handling the planning memo in a follow-up makes sense.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4237377080","thread_id":"inline:4237377080","verdict_sha256":"7971a94f9214bdf1857a39473550a560c323a2594d2e01d88f4a127ddd3021af"}
-->
--
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]
