[SPARK-50838][SQL]Performs additional checks inside recursive CTEs to throw an error if forbidden case is encountered

common/utils/src/main/resources/error/error-conditions.json

@@ -3024,6 +3024,29 @@
     ],
     "sqlState" : "42602"
   },
+  "INVALID_RECURSIVE_REFERENCE" : {
+    "message" : [
+      "Invalid recursive reference found."
+    ],
+    "subClass" : {
+      "DATA_TYPE" : {
+        "message" : [
+          "The data type of recursive references cannot change during resolution. Originally it was <fromDataType> but after resolution is <toDataType>."
+        ]
+      },
+      "NUMBER" : {
+        "message" : [
+          "Recursive references cannot be used multiple times."
+        ]
+      },
+      "PLACE" : {
+        "message" : [
+          "Recursive references cannot be used on the right side of left outer/semi/anti joins, on the left side of right outer joins, in full outer joins and in aggregates."
+        ]
+      }
+    },
+    "sqlState" : "42836"
+  },
   "INVALID_QUERY_MIXED_QUERY_PARAMETERS" : {
     "message" : [
       "Parameterized query must either use positional, or named parameters, but not both."

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala

@@ -932,6 +932,7 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB
           case _ => // Analysis successful!
         }
     }
+    checkRecursion(plan)
     checkCollectedMetrics(plan)
     extendedCheckRules.foreach(_(plan))
     plan.foreachUp {
@@ -1042,6 +1043,75 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB
     if (Utils.isTesting) scrubOutIds(result) else result
   }
 
+  /**
+   * Recursion, according to SQL standard, comes with several limitations:
+   * 1. Recursive term can contain one recursive reference only.
+   * 2. Recursive reference can't be used in some kinds of joins and aggregations.
+   * This rule checks that these restrictions are not violated.
+   */
+  private def checkRecursion(
+      plan: LogicalPlan,
+      references: mutable.Map[Long, (Int, Seq[DataType])] = mutable.Map.empty): Unit = {
+    plan match {
+      // The map is filled with UnionLoop id as key and 0 (number of Ref occasions) and datatype
+      // as value
+      case UnionLoop(id, anchor, recursion, _) =>
+        checkRecursion(anchor, references)
+        checkRecursion(recursion, references += id -> (0, anchor.output.map(_.dataType)))
+        references -= id
+      case r @ UnionLoopRef(loopId, output, false) =>
+        // If we encounter a recursive reference, it has to be present in the map
+        if (!references.contains(loopId)) {
+          r.failAnalysis(
+            errorClass = "INVALID_RECURSIVE_REFERENCE.PLACE",
+            messageParameters = Map.empty
+          )
+        }
+        val (count, dataType) = references(loopId)
+        if (count > 0) {
+          r.failAnalysis(
+            errorClass = "INVALID_RECURSIVE_REFERENCE.NUMBER",
+            messageParameters = Map.empty
+          )
+        }
+        val originalDataType = r.output.map(_.dataType)
+        if (!originalDataType.zip(dataType).forall {
+          case (odt, dt) => DataType.equalsStructurally(odt, dt, true)
+        }) {
+          r.failAnalysis(
+            errorClass = "INVALID_RECURSIVE_REFERENCE.DATA_TYPE",
+            messageParameters = Map(
+              "fromDataType" -> originalDataType.map(toSQLType).mkString(", "),
+              "toDataType" -> dataType.map(toSQLType).mkString(", ")
+            )
+          )
+        }
+        references(loopId) = (count + 1, dataType)
+      case Join(left, right, Inner, _, _) =>
+        checkRecursion(left, references)
+        checkRecursion(right, references)
+      // Reference not allowed in the right part of LeftOuter join
+      case Join(left, right, LeftOuter, _, _) =>
+        checkRecursion(left, references)
+        checkRecursion(right, mutable.Map.empty)
+      case Join(left, right, RightOuter, _, _) =>
+        checkRecursion(left, mutable.Map.empty)
+        checkRecursion(right, references)
+      case Join(left, right, LeftSemi, _, _) =>
+        checkRecursion(left, references)
+        checkRecursion(right, mutable.Map.empty)
+      case Join(left, right, LeftAnti, _, _) =>
+        checkRecursion(left, references)
+        checkRecursion(right, mutable.Map.empty)
+      case Join(left, right, _, _, _) =>
+        checkRecursion(left, mutable.Map.empty)
+        checkRecursion(right, mutable.Map.empty)
+      // Reference not allowed below an aggregate node
+      case Aggregate(_, _, child, _) => checkRecursion(child, mutable.Map.empty)
+      case o => o.children.foreach(checkRecursion(_, references))
+    }
+  }
+
   /**
    * Validates subquery expressions in the plan. Upon failure, returns an user facing error.
    */

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala

@@ -61,7 +61,8 @@ case class InlineCTE(
     // 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 ||
+      // Don't inline recursive CTEs if not necessary as recursion is very costly.
+      (cteDef.deterministic && !cteDef.recursive) ||
       cteDef.child.exists(_.expressions.exists(_.isInstanceOf[OuterReference]))
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -845,6 +845,11 @@ object LimitPushDown extends Rule[LogicalPlan] {
     case LocalLimit(exp, u: Union) =>
       LocalLimit(exp, u.copy(children = u.children.map(maybePushLocalLimit(exp, _))))
 
+    case l @ LocalLimit(IntegerLiteral(limit), p @ Project(_, u: UnionLoop)) =>
+      l.copy(child = p.copy(child = u.copy(limit = Some(limit))))
+    case l @ LocalLimit(IntegerLiteral(limit), u: UnionLoop) =>
+      l.copy(child = u.copy(limit = Some(limit)))
+
     // Add extra limits below JOIN:
     // 1. For LEFT OUTER and RIGHT OUTER JOIN, we push limits to the left and right sides
     //    respectively if join condition is not empty.
@@ -1028,6 +1033,9 @@ object ColumnPruning extends Rule[LogicalPlan] {
       } else {
         p
       }
+    // TODO: Pruning `UnionLoop`s needs to take into account both the outer `Project` and the inner
+    //  `UnionLoopRef` nodes.
+    case p @ Project(_, _: UnionLoop) => p
 
     // Prune unnecessary window expressions
     case p @ Project(_, w: Window) if !w.windowOutputSet.subsetOf(p.references) =>