[SQL] Guaranteed CTE shuffle reuse

[AveryQi115]

TL;DR

Guarantees that a CTE marked non-inlinable (CTERelationDef.forceSkipInline = true) is materialized exactly once and physically reused across all of its references. Each such CTE reference is turned into a CTEReuseExchange node wrapped over a single shuffle of the CTE body. Whether or not Adaptive Query Execution (AQE) is enabled, execution respects the CTEReuseExchange: every reference reads the same shuffle files rather than recomputing the CTE. This closes the gap where stock ReuseExchangeAndSubquery cannot guarantee reuse for a non-deterministic CTE body — non-determinism breaks canonical equality, so the stock rule may fail to dedup the copies and silently recompute.

Why it is needed

A CTE referenced multiple times can be either inlined (recomputed per reference) or materialized once. For a non-deterministic body (e.g. rand(), monotonically_increasing_id(), non-deterministic UDFs), recomputation is not just wasteful but semantically wrong — each reference would observe different rows. Stock reuse keys on canonical equality of the physical subtrees, which non-determinism and per-reference exprId remapping defeat. This change introduces a cteId-keyed mechanism so all references provably share one materialized exchange.

---

Life cycle

The feature threads one identity — the cteId — from the logical plan down to a single physical shuffle:

CTERelationDef / CTERelationRef            (analyzer; forceSkipInline = true)
        │  ReplaceCTERefWithRepartition
        ▼
RepartitionByExpression (plan-reuse id)    (optimizer; body inlined + deduplicated)
        │  ReplaceRepartitionWithCTEReuse
        ▼
CTEReuseRelation(cteId, partitioning, …)   (optimizer; only when refs > 1)
        │  SparkStrategies
        ▼
CTEReuseExchange(cteId, shuffle)           (physical leaf; shuffle held as metadata)
        │
        ├── AQE on  ─▶  CTEReuseQueryStageExec ─▶ inner AdaptiveSparkPlanExec ─▶ Shuffle
        └── AQE off ─▶  LOCAL_SHUFFLE_FOR_CTE shuffle ─▶ ReusedExchangeExec (canonical reuse)

Logical → physical (common to both modes)

• ReplaceCTERefWithRepartition replaces each reference to a forceSkipInline CTE with a plan-reuse RepartitionByExpression (a non-zero repartitionId, assigned by RepartitionIdGenerator) wrapping the inlined CTE body. If the def carries forcePartitioning = Some(HashPartitioning), the reference is materialized with that pinned partitioning; otherwise it falls back to the normal path (RepartitionByExpression(Seq.empty, …), unless the body already satisfies the partitioning or the reference is under a subquery). The inlined body is run through DeduplicateRelations so each copy gets fresh exprIds — the repartitionId, not exprId equality, is what ties the copies together.
• ReplaceRepartitionWithCTEReuse (main-query pass only) counts plan-reuse repartitions by repartitionId across the whole plan including subqueries. Every id with more than one consumer is sealed into a CTEReuseRelation(cteId, partitioning, sharedSubplan); single-consumer ids stay as ordinary repartitions and simply inline. sharedSubplan is deliberately not a child of CTEReuseRelation, so outer-plan re-optimization treats the reference as a leaf and never descends into (or diverges) the shared body.
• SparkStrategies plans CTEReuseRelation eagerly (not via planLater, since the shuffle lives in metadata that the placeholder collector cannot see) into a ShuffleExchangeExec over the body, re-stamps its origin as LOCAL_SHUFFLE_FOR_CTE(cteId), and wraps it as a CTEReuseExchange(cteId, shuffle) leaf. The shuffle is stored in innerChildren, not children, so AQE's bottom-up stage creation does not redundantly plan its substages.

AQE on: CTEReuseExchange → CTEReuseQueryStageExec → inner AdaptiveSparkPlanExec → Shuffle

• PlanCTEReuse (a preprocessing rule alongside PlanAdaptiveSubqueries) builds one shared inner AdaptiveSparkPlanExec per cteId, registered in AdaptiveExecutionContext.cteAQERegistry. The registry lives on the shared execution context, so the identical inner AQE instance is reused across every reference — main query and subqueries — and across re-plans, preventing duplicate replanning or plan divergence of the shared body.
• createNonResultQueryStages replaces each CTEReuseExchange with a CTEReuseQueryStageExec(id, innerAQE, output). All references' stages point at the same inner AQE; each stage remaps the inner AQE's (primary reference's) attribute ids into its own reference's id space, mirroring ReusedExchangeExec. Execution and statistics delegate to the inner AQE's post-iteration executedPlan, so the shared body is iterated once.

AQE off: CTEReuseExchange → LOCAL_SHUFFLE_FOR_CTE shuffle → ReusedExchangeExec

• UnwrapCTEReuseExchange (early in the preparation pipeline, in both the main and subquery batches) replaces each CTEReuseExchange with its underlying LOCAL_SHUFFLE_FOR_CTE(cteId) shuffle. From there all normal preparation rules run per copy.
• Because every copy of a cteId is planned from the same canonicalized sharedSubplan and the shuffle is protected from EnsureRequirements, the copies stay canonically equal, and the stock ReuseExchangeAndSubquery deduplicates them into ReusedExchangeExec.
• VerifyCTEReuse runs at the very end (final main-query pass only) and asserts at most one live shuffle per cteId. Under spark.sql.optimizer.failOnCteReuseWithoutAqe (test default) a violation throws; in production it logs. Reuse is a best-effort optimization here, not a correctness invariant.

---

Design considerations and details

When AQE is on

1. Deadlock avoidance (dedicated thread pool). The inner AQE is materialized via materialize() / lazy materializeFuture, which calls withFinalPlanUpdate(skipResultStage = true) on a dedicated cteExecutionContext thread pool (sized by spark.sql.cteMaterialization.maxThreadThreshold, default 1024). The outer AQE's stage-materialization threads block waiting on the CTE stage to finish; if the inner AQE materialized on the same shared QueryStageCreator pool, nested materialization could exhaust that pool while outer threads hold slots waiting on it — a classic thread-pool deadlock. A separate pool breaks the dependency cycle. skipResultStage = true materializes the shared body's stages without wrapping a ResultQueryStageExec, since the CTE stage is an input to other stages, not a query result.

2. Runtime filters. The CTE body is materialized once and shared across all consumers, so consumer-specific runtime filters (dynamic partition pruning, bloom filters) must not be pushed below the reuse boundary — doing so would specialize the single shared result to one consumer and corrupt the others. Runtime filters are applied per consumer above the CTEReuseExchange; the shared subplan itself is never specialized. This is a deliberate limitation of guaranteed reuse: correctness of the shared materialization takes precedence over consumer-local filter pushdown.

3. Shuffle planning (EnsureRequirements). The reuse shuffle's partitioning is fixed at the logical level (CTEReuseRelation.partitioning, derived from the common distribution of all consumers, or a safe fallback). EnsureRequirements must not rewrite it, because re-deriving partitioning per consumer would make the copies diverge and defeat reuse. The guard is case s: ShuffleExchangeExec if !s.isCreatedForSubplanReuse => s.copy(outputPartitioning = …), where isCreatedForSubplanReuse is true for the LOCAL_SHUFFLE_FOR_CTE origin. The protected shuffle is an immutable boundary: consumer-dependent operators are added above it, keeping the subtree below canonically identical across copies. The same origin also keeps the shuffle out of the AQE supportedShuffleOrigins allowlists (coalesce / local read / skew join), so its partition layout stays frozen.

When AQE is off

1. Interaction with planSubqueries. With AQE off, subqueries are prepared separately — PlanSubqueries calls QueryExecution.prepareExecutedPlan → preparations(subquery = true), applying the preparation rules recursively to each subquery plan on its own. So UnwrapCTEReuseExchange is registered in both the main and subquery preparation batches, ensuring CTE references inside subqueries are unwrapped too. Conversely, VerifyCTEReuse runs only on the final main-query pass: during the per-subquery PlanSubqueries invocations, reuse across the whole plan is not yet complete, so verifying there would produce false negatives. Reuse itself is achieved entirely through the stock ReuseExchangeAndSubquery (canonical matching) — this change adds no new dedup engine for the AQE-off path, only the protected LOCAL_SHUFFLE_FOR_CTE boundary that keeps the copies canonically equal for the stock rule to find.

Testing

Adds CTEReuseWithAQESuite, CTEReuseWithoutAQESuite, and ReplaceCTERefAndRepartitionWithCTEReuseSuite covering the optimizer rewrites, the AQE-on inner-AQE path, and the AQE-off canonical-reuse path (including reuse across subqueries). All new confs default to off, so existing behavior is unchanged.

---

Co-authored-by: Maryann Xue maryann.xue@gmail.com
Co-authored-by: Juliusz Sompolski julek@databricks.com

This pull request and its description were written by Isaac.