stevomitric commented on code in PR #60061:
URL: https://github.com/apache/spark/pull/60061#discussion_r4483925940
##########
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:
done
##########
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:
done
##########
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:
done
--
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]
