gengliangwang commented on code in PR #57091:
URL: https://github.com/apache/spark/pull/57091#discussion_r3540143319
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala:
##########
@@ -49,20 +50,56 @@ 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) {
+ throw SparkException.internalError(
+ "A force-materialized CTE cannot carry an outer reference across its
boundary, but " +
+ "found a subquery with outer-scope reference " +
+ s"'${outerScopeSubqueries.head.outerScopeAttrs.head.name}' in the
CTE definition " +
Review Comment:
This doesn't compile: `outerScopeAttrs` is typed `Seq[Expression]`, and
`Expression` has no `.name` (only `NamedExpression` does). CI's "Precompile
Spark" fails right here — `value name is not a member of ...Expression` at
`InlineCTE.scala:100:63` — which is what cascades the other module builds red.
The elements of `outerScopeAttrs` are expressions that *contain* an
`OuterScopeReference` (the partition predicate is
`_.exists(_.isInstanceOf[OuterScopeReference])`), and may be compound (e.g.
`Add(OuterScopeReference(a), Literal(1))`), so a plain cast to
`NamedExpression` isn't safe either — collect the reference node:
```suggestion
if (outerScopeSubqueries.nonEmpty) {
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 " +
```
This needs `OuterScopeReference` added to the imports on line 24 (`import
org.apache.spark.sql.catalyst.expressions.{Alias, OuterReference,
OuterScopeReference, SubqueryExpression}`). Check 1 above is fine as-is —
`outerRefs` is `Seq[OuterReference]` and `OuterReference` is a
`NamedExpression`.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala:
##########
@@ -50,4 +51,41 @@ 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(OneRowRelation().select(1.as("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(OneRowRelation().select(1.as("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") {
Review Comment:
The three tests look good, but the only cross-boundary test constructs a
direct `OuterReference` on the CTE child, which exercises check 1. Check 2 (the
`outerScopeAttrs`-based subquery branch — the one that currently doesn't
compile) has no coverage. Worth adding a test with a nested-correlation
subquery inside a `forceSkipInline` CTE so `outerScopeAttrs` is non-empty; that
exercises the failing branch and guards the fix.
--
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]
