cloud-fan commented on code in PR #58269:
URL: https://github.com/apache/spark/pull/58269#discussion_r3862860872
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala:
##########
@@ -89,30 +89,45 @@ case class InlineCTE(
}
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}).")
- }
+ // Only an outer reference that actually escapes the CTE definition is
invalid. An outer
+ // reference that points to an attribute produced by an operator inside
the definition
+ // body (e.g. a correlated subquery whose correlated column lives in the
body) is
+ // self-contained and safe to materialize. So walk the whole definition
(main tree plus
+ // nested subquery plans) once, collect every attribute bound anywhere in
the definition,
+ // and reject only references that resolve to none of them.
+ val allNodes = cteDef.child.collectWithSubqueries { case n: LogicalPlan =>
n }
+ val boundExprIds =
allNodes.flatMap(_.output.filter(_.resolved).map(_.exprId)).toSet
- val outerScopeSubqueries = cteDef.child.flatMap(
- _.expressions.flatMap(_.collect {
- case s: SubqueryExpression if s.outerScopeAttrs.nonEmpty => s
+ // Locate either a direct `OuterReference` or a correlated subquery whose
outer-scope
+ // attributes escape the def -- i.e. resolve to none of `boundExprIds`.
+ val escapingOuterRef = allNodes.iterator
+ .flatMap(_.expressions.iterator.flatMap(_.collect {
+ case o: OuterReference if !boundExprIds.contains(o.exprId) => o
}))
- 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}).")
+ .toSeq
Review Comment:
**Non-blocking:**
Stop at the first escaping reference instead of materializing every match.
Both this pipeline and the one ending at line 115 build a `Seq` only to take
`headOption`, so call `nextOption()` directly on each iterator. Build
`boundExprIds` from `allNodes.iterator` as well to avoid its intermediate `Seq`.
--
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]
