This is an automated email from the ASF dual-hosted git repository.

cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 176249b73956 [SPARK-58006][SQL] Add forceSkipInline flag to 
CTERelationDef to force CTE materialization in logical planning
176249b73956 is described below

commit 176249b739567f320a4ae14725754f81599045b7
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]>
---
 .../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]

Reply via email to