This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 172ebf16ce1c [SPARK-58006][SQL] Add forceSkipInline flag to
CTERelationDef to force CTE materialization in logical planning
172ebf16ce1c is described below
commit 172ebf16ce1c4debe8158aee0001188b5ff572bb
Author: Avery <[email protected]>
AuthorDate: Wed Jul 8 23:27:31 2026 +0800
[SPARK-58006][SQL] Add forceSkipInline flag to CTERelationDef to force CTE
materialization in logical planning
### What changes were proposed in this pull request?
This PR adds a `forceSkipInline` field to `CTERelationDef`.
When set to true, `InlineCTE` never inlines the CTE relation, even in
`alwaysInline` mode or when the CTE is deterministic and referenced only once.
This lets a producer guarantee that a CTE is materialized rather than
duplicated.
The flag is self-gating: it defaults to false, so there is no behavior
change unless a caller explicitly sets it.
Because a force-materialized CTE is evaluated on its own, it cannot carry
an outer reference across its boundary (there is no surrounding operator to
resolve it against after materialization). `InlineCTE` now validates this
invariant and throws an `INTERNAL_ERROR` with a descriptive message when it is
violated, for both direct `OuterReference`s and subqueries with outer-scope
references.
`PushdownPredicatesAndPruneColumnsForCTEDef` is updated to preserve the new
field via `copy` instead of reconstructing `CTERelationDef`, and
`CTERelationDef.stringArgs` keeps the default string representation stable when
the flag is not set.
### Why are the changes needed?
Previously the default operation for CTE in spark is to inline it.
However, there are many internal features requiring to generate a
materialized CTE which is guaranteed to be executed only once. For example,
DSv2 command rewrites (e.g. MERGE INTO) can wrap a source in a CTE to avoid
duplicating it. For correctness, if such CTE contains non-deterministic
operation, such a CTE must be materialized and never inlined. There is
currently no way to express that requirement explicitly on `CTERelationDef`.
Besides, we plan to reuse CTE results by introducing a shuffle and reuse
the shuffle files. For usecases like DSv2 commands, the best CTE partitioning
is already defined by the consumers. (In DSv2 case, the partitioning is the
join keys for the join between source and target.) Providing a field for the
suggested partitioning for the materialized CTE is necessary. The suggested
partitioning field will be added in later prs for the changes for shuffle
planning and query execution.
### Does this PR introduce any user-facing change?
No. The new field defaults to false and is internal only.
### How was this patch tested?
New unit tests in `InlineCTESuite` covering: the flag keeps a
single-reference deterministic CTE materialized, the same CTE is inlined when
the flag is unset, and validation fails with `INTERNAL_ERROR` when a
force-materialized CTE carries an outer reference across its boundary.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: claude-code-Opus 4.8 (1M context)
Closes #57091 from AveryQi115/SPARK-58006.
Authored-by: Avery <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit 176249b739567f320a4ae14725754f81599045b7)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../spark/sql/catalyst/optimizer/InlineCTE.scala | 62 ++++++++++++++++++----
...ushdownPredicatesAndPruneColumnsForCTEDef.scala | 10 ++--
.../sql/catalyst/plans/logical/cteOperators.scala | 17 +++++-
.../sql/catalyst/optimizer/InlineCTESuite.scala | 58 ++++++++++++++++++++
4 files changed, 132 insertions(+), 15 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala
index 44b31d3a05b4..c8bd01ed8849 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala
@@ -19,8 +19,9 @@ package org.apache.spark.sql.catalyst.optimizer
import scala.collection.mutable
+import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.analysis.DeduplicateRelations
-import org.apache.spark.sql.catalyst.expressions.{Alias, OuterReference,
SubqueryExpression}
+import org.apache.spark.sql.catalyst.expressions.{Alias, OuterReference,
OuterScopeReference, SubqueryExpression}
import org.apache.spark.sql.catalyst.plans.Inner
import org.apache.spark.sql.catalyst.plans.logical.{CTERelationDef,
CTERelationRef, Join, JoinHint, LogicalPlan, Project, Subquery, UnionLoop,
WithCTE}
import org.apache.spark.sql.catalyst.rules.Rule
@@ -49,20 +50,61 @@ case class InlineCTE(
val cteMap = mutable.SortedMap.empty[Long, CTEReferenceInfo]
buildCTEMap(plan, cteMap)
cleanCTEMap(cteMap)
- inlineCTE(plan, cteMap)
+ val inlined = inlineCTE(plan, cteMap)
+ // A CTE that opts out of inlining must be self-contained: it cannot
carry an outer
+ // reference across its boundary, because after the CTE is materialized
there is no
+ // surrounding operator to resolve that reference against.
+ inlined.foreachWithSubqueries {
+ case cteDef: CTERelationDef if cteDef.forceSkipInline =>
+ validateNoOuterReferencesAcrossCTEBoundary(cteDef)
+ case _ =>
+ }
+ inlined
} else {
plan
}
}
- private def shouldInline(cteDef: CTERelationDef, refCount: Int): Boolean =
alwaysInline || {
- // We do not need to check enclosed `CTERelationRef`s for `deterministic`
or `OuterReference`,
- // because:
- // 1) It is fine to inline a CTE if it references another CTE that is
non-deterministic;
- // 2) Any `CTERelationRef` that contains `OuterReference` would have been
inlined first.
- refCount == 1 ||
- cteDef.deterministic ||
- cteDef.child.exists(_.expressions.exists(_.isInstanceOf[OuterReference]))
+ private def shouldInline(cteDef: CTERelationDef, refCount: Int): Boolean = {
+ // A CTE definition that requests to skip inlining is never inlined, even
in `alwaysInline`
+ // mode, so that a producer can guarantee the CTE is materialized rather
than duplicated.
+ !cteDef.forceSkipInline && (alwaysInline || {
+ // We do not need to check enclosed `CTERelationRef`s for
`deterministic` or
+ // `OuterReference`, because:
+ // 1) It is fine to inline a CTE if it references another CTE that is
non-deterministic;
+ // 2) Any `CTERelationRef` that contains `OuterReference` would have
been inlined first.
+ refCount == 1 ||
+ cteDef.deterministic ||
+
cteDef.child.exists(_.expressions.exists(_.isInstanceOf[OuterReference]))
+ })
+ }
+
+ private def validateNoOuterReferencesAcrossCTEBoundary(cteDef:
CTERelationDef): Unit = {
+ val outerRefs = cteDef.child.flatMap(
+ _.expressions.flatMap(_.collect { case o: OuterReference => o }))
+ if (outerRefs.nonEmpty) {
+ throw SparkException.internalError(
+ "A force-materialized CTE cannot carry an outer reference across its
boundary, but " +
+ s"found outer reference '${outerRefs.head.name}' in the CTE
definition " +
+ s"(cteId=${cteDef.id}).")
+ }
+
+ val outerScopeSubqueries = cteDef.child.flatMap(
+ _.expressions.flatMap(_.collect {
+ case s: SubqueryExpression if s.outerScopeAttrs.nonEmpty => s
+ }))
+ if (outerScopeSubqueries.nonEmpty) {
+ // `outerScopeAttrs` are expressions that contain an
`OuterScopeReference` and may be
+ // compound (e.g. `Add(OuterScopeReference(a), Literal(1))`), so collect
the reference node
+ // rather than assuming the attribute itself is one.
+ val outerScopeRef = outerScopeSubqueries.head.outerScopeAttrs
+ .flatMap(_.collect { case r: OuterScopeReference => r }).head
+ throw SparkException.internalError(
+ "A force-materialized CTE cannot carry an outer reference across its
boundary, but " +
+ "found a subquery with outer-scope reference " +
+ s"'${outerScopeRef.name}' in the CTE definition " +
+ s"(cteId=${cteDef.id}).")
+ }
}
/**
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
index 421c47f12ddb..2339596920b0 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
@@ -122,7 +122,7 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends
Rule[LogicalPlan] {
private def pushdownPredicatesAndAttributes(
plan: LogicalPlan,
cteMap: CTEMap): LogicalPlan = plan.transformWithSubqueries {
- case cteDef @ CTERelationDef(child, id, originalPlanWithPredicates, _, _)
=>
+ case cteDef @ CTERelationDef(child, id, originalPlanWithPredicates, _, _,
_) =>
val (_, _, newPreds, newAttrSet) = cteMap(id)
val originalPlan = originalPlanWithPredicates.map(_._1).getOrElse(child)
val preds = originalPlanWithPredicates.map(_._2).getOrElse(Seq.empty)
@@ -134,9 +134,11 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends
Rule[LogicalPlan] {
} else {
originalPlan
}
- CTERelationDef(Filter(newCombinedPred, newChild), id,
Some((originalPlan, newPreds)))
+ cteDef.copy(child = Filter(newCombinedPred, newChild),
+ originalPlanWithPredicates = Some((originalPlan, newPreds)))
} else if (needsPruning(cteDef.child, newAttrSet)) {
- CTERelationDef(Project(newAttrSet.toSeq, cteDef.child), id,
Some((originalPlan, preds)))
+ cteDef.copy(child = Project(newAttrSet.toSeq, cteDef.child),
+ originalPlanWithPredicates = Some((originalPlan, preds)))
} else {
cteDef
}
@@ -170,7 +172,7 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends
Rule[LogicalPlan] {
object CleanUpTempCTEInfo extends Rule[LogicalPlan] {
override def apply(plan: LogicalPlan): LogicalPlan =
plan.transformWithPruning(_.containsPattern(CTE)) {
- case cteDef @ CTERelationDef(_, _, Some(_), _, _) =>
+ case cteDef @ CTERelationDef(_, _, Some(_), _, _, _) =>
cteDef.copy(originalPlanWithPredicates = None)
}
}
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala
index 4114bc290312..d7a2f0e76e37 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala
@@ -106,16 +106,31 @@ case class UnionLoopRef(
* @param underSubquery If true, it means we don't need to add a shuffle for
this CTE relation as
* subquery reuse will be applied to reuse CTE relation
output.
* @param maxDepth The maximal depth of a recursion in a recursive CTE.
+ * @param forceSkipInline If true, this CTE relation will never be inlined by
[[InlineCTE]],
+ * regardless of determinism or reference count. This
lets a producer
+ * force the CTE to be materialized instead of
duplicated, e.g. when the
+ * CTE wraps a non-deterministic source that must be
evaluated exactly once.
*/
case class CTERelationDef(
child: LogicalPlan,
id: Long = CTERelationDef.newId,
originalPlanWithPredicates: Option[(LogicalPlan, Seq[Expression])] = None,
underSubquery: Boolean = false,
- maxDepth: Option[Int] = None) extends UnaryNode {
+ maxDepth: Option[Int] = None,
+ forceSkipInline: Boolean = false) extends UnaryNode {
final override val nodePatterns: Seq[TreePattern] = Seq(CTE)
+ // Keep the default string representation stable when `forceSkipInline` is
not set, so that
+ // existing plan comparisons and golden files are unaffected by the new
field.
+ override def stringArgs: Iterator[Any] = {
+ if (forceSkipInline) {
+ super.stringArgs
+ } else {
+ super.stringArgs.toArray.dropRight(1).iterator
+ }
+ }
+
override def maxRows: Option[Long] = if
(conf.getConf(SQLConf.CTE_RELATION_DEF_MAX_ROWS)) {
child.maxRows
} else {
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala
index e23803c6f688..0d515c482401 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala
@@ -21,6 +21,7 @@ import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.analysis.TestRelation
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.{OuterReference,
OuterScopeReference, ScalarSubquery}
import org.apache.spark.sql.catalyst.plans.PlanTest
import org.apache.spark.sql.catalyst.plans.logical.{AppendData,
CTERelationDef, CTERelationRef, LogicalPlan, OneRowRelation, WithCTE}
import org.apache.spark.sql.catalyst.rules.RuleExecutor
@@ -50,4 +51,61 @@ class InlineCTESuite extends PlanTest {
}
assert(e.getMessage.contains("No CTERelationDef found"))
}
+
+ test("SPARK-58006: forceSkipInline keeps a single-reference deterministic
CTE") {
+ // A single-reference deterministic CTE is normally inlined, but
`forceSkipInline` should
+ // keep it materialized in the `WithCTE` node.
+ val cteDef = CTERelationDef(
+ TestRelation(Seq($"a".int)).select($"a"), forceSkipInline = true)
+ val cteRef = CTERelationRef(cteDef.id, cteDef.resolved, cteDef.output,
cteDef.isStreaming)
+ val plan = WithCTE(cteRef.select($"a"), Seq(cteDef)).analyze
+ val optimized = Optimize.execute(plan)
+ assert(optimized.collectFirst { case _: WithCTE => true }.isDefined,
+ "CTE with forceSkipInline should not be inlined")
+ }
+
+ test("SPARK-58006: forceSkipInline is inlined normally when not set") {
+ // The same single-reference deterministic CTE is inlined when
`forceSkipInline` is false.
+ val cteDef = CTERelationDef(TestRelation(Seq($"a".int)).select($"a"))
+ val cteRef = CTERelationRef(cteDef.id, cteDef.resolved, cteDef.output,
cteDef.isStreaming)
+ val plan = WithCTE(cteRef.select($"a"), Seq(cteDef)).analyze
+ val optimized = Optimize.execute(plan)
+ assert(optimized.collectFirst { case _: WithCTE => true }.isEmpty,
+ "CTE without forceSkipInline should be inlined")
+ }
+
+ test("SPARK-58006: forceSkipInline CTE with an outer reference across its
boundary fails") {
+ // A force-materialized CTE cannot carry an outer reference across its
boundary, because after
+ // materialization there is no surrounding operator to resolve it against.
+ val relation = TestRelation(Seq($"a".int))
+ val cteChild = relation.where(OuterReference($"a".int) === 1)
+ val cteDef = CTERelationDef(cteChild, forceSkipInline = true)
+ val cteRef = CTERelationRef(cteDef.id, cteDef.resolved, cteDef.output,
cteDef.isStreaming)
+ val plan = WithCTE(cteRef.select($"a"), Seq(cteDef))
+ val e = intercept[SparkException] {
+ Optimize.execute(plan)
+ }
+ assert(e.getCondition == "INTERNAL_ERROR")
+ assert(e.getMessage.contains(
+ "A force-materialized CTE cannot carry an outer reference across its
boundary"))
+ }
+
+ test("SPARK-58006: forceSkipInline CTE with an outer-scope subquery
reference fails") {
+ // Exercises the second validation branch: a subquery inside the CTE child
carries an
+ // outer-scope reference (nested correlation), which also cannot cross the
CTE boundary.
+ val relation = TestRelation(Seq($"a".int))
+ val subquery = ScalarSubquery(
+ TestRelation(Seq($"c".int)).select($"c"),
+ outerAttrs = Seq(OuterScopeReference($"a".int)))
+ val cteChild = relation.select(subquery.as("s"))
+ val cteDef = CTERelationDef(cteChild, forceSkipInline = true)
+ val cteRef = CTERelationRef(cteDef.id, cteDef.resolved, cteDef.output,
cteDef.isStreaming)
+ val plan = WithCTE(cteRef.select($"s"), Seq(cteDef))
+ val e = intercept[SparkException] {
+ Optimize.execute(plan)
+ }
+ assert(e.getCondition == "INTERNAL_ERROR")
+ assert(e.getMessage.contains(
+ "found a subquery with outer-scope reference"))
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
