spark-test commented on code in PR #60061:
URL: https://github.com/apache/spark/pull/60061#discussion_r4471355397
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveWithCTE.scala:
##########
@@ -303,22 +322,52 @@ object ResolveWithCTE extends Rule[LogicalPlan] {
// Assumes that `anchor` is already resolved. If `columnNames` is set (which
happens if an
// UnresolvedSubqueryColumnAliases is atop a Union) also places an
UnresolvedSubqueryColumnAliases
// node above UnionLoopRef. At some later stage
UnresolvedSubqueryColumnAliases gets resolved to
- // a Project node.
+ // a Project node. `refTypes` overrides the anchor's data types for the
UnionLoopRef output.
private def rewriteRecursiveCTERefs(
recursion: LogicalPlan,
anchor: LogicalPlan,
cteDefId: Long,
- columnNames: Option[Seq[String]]) = {
+ columnNames: Option[Seq[String]],
+ refTypes: Option[Seq[DataType]] = None) = {
+ val types = refTypes.getOrElse(anchor.output.map(_.dataType))
recursion.transformUpWithSubqueriesAndPruning(_.containsPattern(CTE)) {
case r: CTERelationRef if r.recursive && r.cteId == cteDefId =>
// We mark the output of UnionLoopRef as nullable as it may become
NULL in the following
// iterations.
- val ref =
- UnionLoopRef(r.cteId,
anchor.output.map(_.newInstance().withNullability(true)), false)
+ val output = anchor.output.zip(types).map { case (a, dt) =>
+ a.newInstance().withNullability(true).withDataType(dt)
+ }
+ val ref = UnionLoopRef(r.cteId, output, false)
columnNames.map(UnresolvedSubqueryColumnAliases(_, ref)).getOrElse(ref)
}
}
+ // The anchor's types, unless widenSelfReference has found wider ones.
+ private def selfReferenceTypes(cteDef: CTERelationDef, anchor: LogicalPlan):
Seq[DataType] =
+
cteDef.getTagValue(SELF_REFERENCE_TYPES).getOrElse(anchor.output.map(_.dataType))
+
+ // A recursion can write NULLs into nested types that the self-reference
declares non-nullable,
+ // e.g. ARRAY(CAST(NULL AS INT)) for an ARRAY(1) anchor. Once the loop is
resolved, widens the
+ // nullability of such self-reference types to the recursion's and rebuilds
the loop from the
+ // original definition. This repeats until no type widens.
+ private def widenSelfReference(cteDef: CTERelationDef): CTERelationDef = {
+ val loop = cteDef.child.collectFirst { case l: UnionLoop if l.id ==
cteDef.id => l }
+ loop.filter(_.resolved).flatMap(l => l.getTagValue(LOOP_ORIGIN).map(l ->
_)) match {
Review Comment:
**Non-blocking (P2):** `LOOP_ORIGIN` carries the recorded types and the
rebuild source between Resolution-batch iterations, and when it is missing the
`None` branch below returns the definition unchanged. Tags survive only through
the `copyTagsFrom` calls in the transform helpers, so a rule that rebuilds the
`UnionLoop` with a case-class copy drops them. `LimitPushDown` and
`InlineCTE.setUnlimitedRecursion` already rebuild this node that way, although
they run after analysis. No Resolution-batch rule in Spark does this today, but
nothing prevents one, for example a rule injected through
`injectResolutionRule`. If one does, NULL-writing recursions read the NULLs
back as 0 again with no error. In the coercion case the base branch failed with
`CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE`, so tag loss would turn that explicit
error into a silent wrong result.
Reading the current types from the loop's own `UnionLoopRef` output, and
raising an internal error when the loop needs wider types but the rebuild
source is missing, would make this fail closed.
**Verification:**
- **Regression:** Analyzing a recursive CTE whose resolved loop has lost its
rebuild source, while its recursion is wider than its self-reference, fails
with an internal error. It must not produce a plan whose self-reference has the
narrower types.
- **Behavior:** Existing widening cases keep the same self-reference and
output types: a pass-through keeps the anchor's types, and a column that copies
a widened column widens on the second round.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveWithCTE.scala:
##########
@@ -303,22 +322,52 @@ object ResolveWithCTE extends Rule[LogicalPlan] {
// Assumes that `anchor` is already resolved. If `columnNames` is set (which
happens if an
// UnresolvedSubqueryColumnAliases is atop a Union) also places an
UnresolvedSubqueryColumnAliases
// node above UnionLoopRef. At some later stage
UnresolvedSubqueryColumnAliases gets resolved to
- // a Project node.
+ // a Project node. `refTypes` overrides the anchor's data types for the
UnionLoopRef output.
private def rewriteRecursiveCTERefs(
recursion: LogicalPlan,
anchor: LogicalPlan,
cteDefId: Long,
- columnNames: Option[Seq[String]]) = {
+ columnNames: Option[Seq[String]],
+ refTypes: Option[Seq[DataType]] = None) = {
+ val types = refTypes.getOrElse(anchor.output.map(_.dataType))
recursion.transformUpWithSubqueriesAndPruning(_.containsPattern(CTE)) {
case r: CTERelationRef if r.recursive && r.cteId == cteDefId =>
// We mark the output of UnionLoopRef as nullable as it may become
NULL in the following
// iterations.
- val ref =
- UnionLoopRef(r.cteId,
anchor.output.map(_.newInstance().withNullability(true)), false)
+ val output = anchor.output.zip(types).map { case (a, dt) =>
+ a.newInstance().withNullability(true).withDataType(dt)
+ }
+ val ref = UnionLoopRef(r.cteId, output, false)
columnNames.map(UnresolvedSubqueryColumnAliases(_, ref)).getOrElse(ref)
}
}
+ // The anchor's types, unless widenSelfReference has found wider ones.
+ private def selfReferenceTypes(cteDef: CTERelationDef, anchor: LogicalPlan):
Seq[DataType] =
+
cteDef.getTagValue(SELF_REFERENCE_TYPES).getOrElse(anchor.output.map(_.dataType))
+
+ // A recursion can write NULLs into nested types that the self-reference
declares non-nullable,
+ // e.g. ARRAY(CAST(NULL AS INT)) for an ARRAY(1) anchor. Once the loop is
resolved, widens the
+ // nullability of such self-reference types to the recursion's and rebuilds
the loop from the
+ // original definition. This repeats until no type widens.
+ private def widenSelfReference(cteDef: CTERelationDef): CTERelationDef = {
+ val loop = cteDef.child.collectFirst { case l: UnionLoop if l.id ==
cteDef.id => l }
+ loop.filter(_.resolved).flatMap(l => l.getTagValue(LOOP_ORIGIN).map(l ->
_)) match {
+ case Some((l, (refTypes, child))) =>
+ val widened = refTypes.zip(l.recursion.output).map { case (dt, a) =>
Review Comment:
**Non-blocking (P2):** This comparison runs only in the Resolution batch,
but the later "ScalaUDF Null Handling" batch can still raise the nested
nullability of both loop children. A Scala UDF with primitive types is
non-nullable during resolution (its nullability comes from the output encoder),
and `HandleNullInputsForUDF` only later wraps it in `If(IsNull(input), null,
udf)`. `CreateArray` then re-derives `containsNull = true`, while the
`UnionLoopRef` keeps the types recorded here, and `UpdateAttributeNullability`
skips leaf nodes.
For example, with `spark.udf.register("f", (x: Long) => x + 1)`:
```sql
WITH RECURSIVE t(lvl, arr) AS (
SELECT 0, ARRAY(id) FROM range(200)
UNION ALL
SELECT lvl + 1, ARRAY(f(IF(lvl = 0, CAST(NULL AS BIGINT), arr[0]))) FROM t
WHERE lvl < 2
)
SELECT * FROM t
```
Level 1 writes `[null]`. Level 2 reads `arr[0]` through the narrower
self-reference, and `GetArrayItem` codegen skips `isNullAt` when `containsNull`
is false, so it gets `[1]` instead of `[null]`. A NULL written by a UDF in the
anchor is read back the same way. Small steps that are folded into a local
relation use interpreted evaluation, which still checks `isNullAt`.
Because the analyzed loop then no longer matches the recorded types, a later
analysis of the same plan restores the stored definition from `LOOP_ORIGIN`.
That happens on the deduplicated side of `df.join(df)`, so the two sides of
such a self-join can return different rows. This is traced from the source; it
was not run.
**Recommended change:** Make the self-reference widening see the nullability
that Scala UDF null handling adds later in analysis. Move
HandleNullInputsForUDF out of the Analyzer class into a standalone rule object
in the same package, and keep the 'ScalaUDF Null Handling' batch using it. In
widenSelfReference, when the loop's anchor or recursion contains a Scala UDF,
run HandleNullInputsForUDF and UpdateAttributeNullability on it until neither
changes it, as that batch does. Merge each self-reference type with the
resulting anchor and recursion output types. This is a pure transform of copies
that leaves the live plan unchanged. Add recursive CTE tests in which a
primitive Scala UDF returns NULL inside an array, once in the anchor and once
in the recursion, and a later iteration reads it back.
**Why this works:** The rebuild already runs in the Resolution batch, where
the recursion can be re-resolved. Computing the merge from null-handled copies
of the anchor and recursion makes the recorded self-reference types include the
element, value and field nullability that the later batch introduces. The final
self-reference then covers everything written into it, and generated code keeps
its null checks. Because the analyzed loop then matches its recorded types, a
later analysis of the same plan no longer restores the stored definition.
**Scope:** Expose the UDF null-handling rule for reuse, apply the UDF
null-handling batch to copies of the anchor and recursion when ResolveWithCTE
computes the widened self-reference types, and add end-to-end coverage with
primitive Scala UDFs.
**Compatibility:** The analyzer still applies UDF null handling exactly once
to the real plan, and a recursion that only passes a column through keeps the
anchor's nested types.
**Risks:** Columns that copy a UDF-fed nested column through the
self-reference will now widen their output types, which adds to the
persisted-view schema change raised separately in this review. Reverting the
UDF-aware merge rolls this back. Each widening check now runs the null-handling
rules on copies of the resolved anchor and recursion when they contain Scala
UDFs. Only recursive CTEs with such UDFs pay this extra analysis work, which
grows with their size. Extension post-hoc resolution rules that raise nested
nullability after Resolution stay uncovered, because the widening cannot see
changes made after that batch.
**Constraints:** The live plan must not be rewritten through the copies; the
Analyzer batch still applies UDF null handling exactly once. The merge must
keep only OR-ing nullability, so the widening still terminates. The copies must
go through both rules until neither changes them. A UDF fed by another UDF's
output becomes nullable only after UpdateAttributeNullability updates the
reference between them.
**Exclusions:**
- Outside recursive CTE self-references, nested nullability that UDF null
handling adds is still not propagated to attribute references above the wrapped
UDF, including the main query's reference to the CTE output. That is a general
analyzer behavior and stays outside this change.
**Success:** When the anchor or the recursion builds an array from a
primitive Scala UDF that returns NULL for a NULL input, the self-reference
declares nullable elements, and a later iteration that reads the element gets
NULL rather than 0. After analysis, merging the recorded self-reference types
with the loop's anchor and recursion output types changes nothing, so a later
analysis of the same plan restores no stored definition. Recursions without
Scala UDFs keep exactly the self-reference types this change already produces.
**Verification:**
- **Regression:** Recursive CTEs that wrap a primitive Scala UDF (Long =>
Long) in ARRAY(...) and pass it a NULL return NULL elements at every later
level that reads them, on the materialized-step path with code generation.
Cover the UDF in the anchor, in the recursion, and fed by another UDF's output
in a lower projection.
- **Regression:** A self-join of such a recursive CTE's DataFrame returns
the same values for the CTE on both sides.
- **Behavior:** For those queries, the analyzed self-reference declares a
nullable element type, and recursions without UDFs keep the same analyzed
self-reference and output types.
- **Compatibility:** Queries that use Scala UDFs outside recursive CTEs
analyze and run exactly as before, with UDF null handling applied once.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveWithCTE.scala:
##########
@@ -94,18 +105,20 @@ object ResolveWithCTE extends Rule[LogicalPlan] {
if (!anchor.resolved) {
cteDef
} else {
+ val refTypes = selfReferenceTypes(cteDef, anchor)
// We need to check whether any of the inner CTEs has a self
reference and replace
// it if needed
val newInnerCteDefs = innerCteDefs.map { innerCteDef =>
innerCteDef.copy(child = rewriteRecursiveCTERefs(
- innerCteDef.child, anchor, cteDef.id, None))
+ innerCteDef.child, anchor, cteDef.id, None,
Some(refTypes)))
}
val loop = UnionLoop(
cteDef.id,
anchor,
- rewriteRecursiveCTERefs(recursion, anchor, cteDef.id,
None),
+ rewriteRecursiveCTERefs(recursion, anchor, cteDef.id,
None, Some(refTypes)),
Review Comment:
**Nit (P3):** This arm, and the column-list variant further down, now pass
`refTypes` to the recursion's `UnionLoopRef`, but no test reaches either one
with nested types. The new suite cases build only `SubqueryAlias(Union)`, the
new `cte-recursion.sql` queries use only the column-list shape without an inner
`WITH`, and the existing inner-`WITH` recursive queries use INT columns only.
If either arm went back to the anchor's types, NULLs would read back as 0 again
and CI would stay green.
A `cte-recursion.sql` case like the following, with and without the column
list, would cover them (regenerate the golden files rather than editing them):
```sql
WITH RECURSIVE t(lvl, arr) AS (
WITH x AS (SELECT 1 AS one)
SELECT 0, ARRAY(one) FROM x
UNION ALL
SELECT lvl + 1, IF(lvl = 0, ARRAY(CAST(NULL AS INT)), ARRAY(arr[0] + 1))
FROM t WHERE lvl < 2
)
SELECT * FROM t
```
**Verification:**
- **Regression:** Take a recursive CTE whose definition begins with an inner
WITH, with and without a column list, run on the materialized-step path. A NULL
written into an array element at one level reads back as NULL at the next
level, and the case fails if that arm uses the anchor's types for its
self-reference.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveWithCTE.scala:
##########
@@ -303,22 +322,52 @@ object ResolveWithCTE extends Rule[LogicalPlan] {
// Assumes that `anchor` is already resolved. If `columnNames` is set (which
happens if an
// UnresolvedSubqueryColumnAliases is atop a Union) also places an
UnresolvedSubqueryColumnAliases
// node above UnionLoopRef. At some later stage
UnresolvedSubqueryColumnAliases gets resolved to
- // a Project node.
+ // a Project node. `refTypes` overrides the anchor's data types for the
UnionLoopRef output.
private def rewriteRecursiveCTERefs(
recursion: LogicalPlan,
anchor: LogicalPlan,
cteDefId: Long,
- columnNames: Option[Seq[String]]) = {
+ columnNames: Option[Seq[String]],
+ refTypes: Option[Seq[DataType]] = None) = {
+ val types = refTypes.getOrElse(anchor.output.map(_.dataType))
recursion.transformUpWithSubqueriesAndPruning(_.containsPattern(CTE)) {
case r: CTERelationRef if r.recursive && r.cteId == cteDefId =>
// We mark the output of UnionLoopRef as nullable as it may become
NULL in the following
// iterations.
- val ref =
- UnionLoopRef(r.cteId,
anchor.output.map(_.newInstance().withNullability(true)), false)
+ val output = anchor.output.zip(types).map { case (a, dt) =>
+ a.newInstance().withNullability(true).withDataType(dt)
+ }
+ val ref = UnionLoopRef(r.cteId, output, false)
columnNames.map(UnresolvedSubqueryColumnAliases(_, ref)).getOrElse(ref)
}
}
+ // The anchor's types, unless widenSelfReference has found wider ones.
+ private def selfReferenceTypes(cteDef: CTERelationDef, anchor: LogicalPlan):
Seq[DataType] =
+
cteDef.getTagValue(SELF_REFERENCE_TYPES).getOrElse(anchor.output.map(_.dataType))
+
+ // A recursion can write NULLs into nested types that the self-reference
declares non-nullable,
+ // e.g. ARRAY(CAST(NULL AS INT)) for an ARRAY(1) anchor. Once the loop is
resolved, widens the
+ // nullability of such self-reference types to the recursion's and rebuilds
the loop from the
+ // original definition. This repeats until no type widens.
+ private def widenSelfReference(cteDef: CTERelationDef): CTERelationDef = {
+ val loop = cteDef.child.collectFirst { case l: UnionLoop if l.id ==
cteDef.id => l }
Review Comment:
**Nit (P3):** This `collectFirst` walks the whole definition without pruning
(`UnionLoop` has no tree pattern). It runs for every definition whenever the
rule visits an unanalyzed `WithCTE`: on every Resolution iteration, and again
for nested `WITH`s, because `WithCTE.children` lists the definitions first. For
non-recursive CTEs, the most common case, it can never match, while the base
branch spent O(1) per stable definition through cached lazy vals. Each visit is
cheap, but the cost scales with definition size times iterations for every CTE
query. Skipping definitions whose cached `hasSelfReferenceAsUnionLoopRef` is
false would give the same result without the repeated walk.
**Verification:**
- **Behavior:** Every recursive CTE widening case (plain, column list, inner
WITH, two-round) keeps the same self-reference and output types.
- **Inspection:** The loop search runs only for definitions whose cached
self-reference check finds a UnionLoopRef. That check runs once per definition
instance, so an unchanged definition without one is not walked again on later
passes.
--
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]
