[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

@@ -3117,6 +3117,34 @@
     ],
     "sqlState" : "42602"
   },
+  "INVALID_RECURSIVE_REFERENCE" : {
+    "message" : [
+      "Invalid recursive reference found inside WITH RECURSIVE clause."
+    ],
+    "subClass" : {
+      "DATA_TYPE" : {
+        "message" : [
+          "The data type of recursive references cannot change during resolution. Check that anchor and recursion term actually output the same data type. Originally it was <fromDataType> but after resolution is <toDataType>."
+        ]
+      },
+      "NUMBER" : {
+        "message" : [
+          "Multiple self-references to one recursive CTE are not allowed."
+        ]
+      },
+      "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."
+        ]
+      },
+      "SUBQUERY" : {
+        "message" : [
+          "Recursive reference must not appear within a subquery."
+        ]
+      }
+    },
+    "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/CTESubstitution.scala

@@ -87,6 +87,11 @@ object CTESubstitution extends Rule[LogicalPlan] {
         case LegacyBehaviorPolicy.CORRECTED =>
           traverseAndSubstituteCTE(plan, forceInline, Seq.empty, cteDefs)
     }
+    // Recursive CTEs allow only one self-reference per CTE. Here, we perform this check and
+    // throw an error if it is not fulfilled
+    cteDefs.foreach { cteDef =>
+      checkNumberOfSelfReferences(cteDef)
+    }
     if (cteDefs.isEmpty) {
       substituted
     } else if (substituted eq firstSubstituted.get) {
@@ -402,7 +407,7 @@ object CTESubstitution extends Rule[LogicalPlan] {
         other.transformExpressionsWithPruning(_.containsPattern(PLAN_EXPRESSION)) {
           case e: SubqueryExpression =>
             e.withNewPlan(
-              apply(substituteCTE(e.plan, alwaysInline, cteRelations, None)))
+              apply(substituteCTE(e.plan, alwaysInline, cteRelations, recursiveCTERelation)))
         }
     }
   }
@@ -423,4 +428,19 @@ object CTESubstitution extends Rule[LogicalPlan] {
       case _ => WithCTE(p, cteDefs)
     }
   }
+
+  /**
+   * Counts number of self-references in a recursive CTE definition and throws an error
+   * if that number is bigger than 1.
+   */
+  private def checkNumberOfSelfReferences(cteDef: CTERelationDef): Unit = {
+    val numOfSelfRef = cteDef.collectWithSubqueries {
+      case ref: CTERelationRef if ref.cteId == cteDef.id => ref
+    }.length
+    if (numOfSelfRef > 1) {
+      cteDef.failAnalysis(
+        errorClass = "INVALID_RECURSIVE_REFERENCE.NUMBER",
+        messageParameters = Map.empty)
+    }
+  }
 }

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

@@ -22,6 +22,7 @@ import org.apache.spark.{SparkException, SparkThrowable}
 import org.apache.spark.internal.{Logging, LogKeys, MDC}
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.ExtendedAnalysisException
+import org.apache.spark.sql.catalyst.analysis.ResolveWithCTE.checkIfSelfReferenceIsPlacedCorrectly
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.SubExprUtils._
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction, ListAgg, Median, PercentileCont, PercentileDisc}
@@ -274,6 +275,10 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB
         checkTrailingCommaInSelect(proj)
       case agg: Aggregate =>
         checkTrailingCommaInSelect(agg)
+      case unionLoop: UnionLoop =>
+        // Recursive CTEs have already substituted Union to UnionLoop at this stage.
+        // Here we perform additional checks for them.
+        checkIfSelfReferenceIsPlacedCorrectly(unionLoop)
 
       case _ =>
     }

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

@@ -21,6 +21,7 @@ import scala.collection.mutable
 
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.expressions.SubqueryExpression
+import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftOuter, LeftSemi, RightOuter}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreePattern.{CTE, PLAN_EXPRESSION}
@@ -48,13 +49,16 @@ object ResolveWithCTE extends Rule[LogicalPlan] {
       cteDefMap: mutable.HashMap[Long, CTERelationDef]): LogicalPlan = {
     plan.resolveOperatorsDownWithPruning(_.containsAllPatterns(CTE)) {
       case withCTE @ WithCTE(_, cteDefs) =>
+        // Check for self-reference in subqueries and throw an error if that is the case.
+        // We don't want to resolve any such reference, therefore we need to perform this check
+        // here.
+        cteDefs.foreach(checkForSelfReferenceInSubquery)
         val newCTEDefs = cteDefs.map {
-          // `cteDef.recursive` means "presence of a recursive CTERelationRef under cteDef". The
-          // side effect of node substitution below is that after CTERelationRef substitution
-          // its cteDef is no more considered `recursive`. This code path is common for `cteDef`
-          // that were non-recursive from the get go, as well as those that are no more recursive
-          // due to node substitution.
-          case cteDef if !cteDef.recursive =>
+          // cteDef in the first case is either recursive and all the recursive CTERelationRefs
+          // are already substituted to UnionLoopRef in the previous pass, or it is not recursive
+          // at all. In both cases we need to put it in the map in case it is resolved.
+          // Second case is performing the substitution of recursive CTERelationRefs.
+          case cteDef if !cteDef.hasSelfReferenceAsCTERef =>
             if (cteDef.resolved) {
               cteDefMap.put(cteDef.id, cteDef)
             }
@@ -183,4 +187,52 @@ object ResolveWithCTE extends Rule[LogicalPlan] {
         columnNames.map(UnresolvedSubqueryColumnAliases(_, ref)).getOrElse(ref)
     }
   }
+
+  /**
+   * Checks if there is any self-reference within subqueries and throws an error
+   * if that is the case.
+   */
+  private def checkForSelfReferenceInSubquery(cteDef: CTERelationDef): Unit = {
+    cteDef.subqueriesAll.foreach { subquery =>
+      subquery.foreach {
+        case CTERelationRef(cteDef.id, _, _, _, _, true) =>
+          cteDef.failAnalysis(
+            errorClass = "INVALID_RECURSIVE_REFERENCE.SUBQUERY",
+            messageParameters = Map.empty)
+        case _ =>
+      }
+    }
+  }
+
+  /**
+   * Throws error if self-reference is placed in places which are not allowed:
+   * right side of left outer/semi/anti joins, left side of right outer joins,
+   * in full outer joins and in aggregates
+   */
+  def checkIfSelfReferenceIsPlacedCorrectly(unionLoop: UnionLoop): Unit = {
+    def unionLoopRefNotAllowedUnderCurrentNode(currentNode: LogicalPlan) : Unit =
+      currentNode.foreach {
+        case UnionLoopRef(unionLoop.id, _, _) =>
+          throw new AnalysisException(
+            errorClass = "INVALID_RECURSIVE_REFERENCE.PLACE",
+            messageParameters = Map.empty)
+        case other =>
+      }
+    unionLoop.foreach {
+      case Join(left, right, LeftOuter, _, _) =>
+        unionLoopRefNotAllowedUnderCurrentNode(right)
+      case Join(left, right, RightOuter, _, _) =>
+        unionLoopRefNotAllowedUnderCurrentNode(left)
+      case Join(left, right, LeftSemi, _, _) =>
+        unionLoopRefNotAllowedUnderCurrentNode(right)
+      case Join(left, right, LeftAnti, _, _) =>
+        unionLoopRefNotAllowedUnderCurrentNode(right)
+      case Join(left, right, _, _, _) =>
+        unionLoopRefNotAllowedUnderCurrentNode(left)
+        unionLoopRefNotAllowedUnderCurrentNode(right)
+      case Aggregate(_, _, child, _) =>
+        unionLoopRefNotAllowedUnderCurrentNode(child)
+      case other =>
+    }
+  }
 }

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

@@ -61,7 +61,10 @@ 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.
+      // The check if cteDef is recursive is performed by checking if it contains
+      // a UnionLoopRef with the same ID.
+      (cteDef.deterministic && !cteDef.hasSelfReferenceAsUnionLoopRef) ||
       cteDef.child.exists(_.expressions.exists(_.isInstanceOf[OuterReference]))
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala

@@ -100,12 +100,14 @@ case class CTERelationDef(
 
   override def output: Seq[Attribute] = if (resolved) child.output else Nil
 
-  lazy val recursive: Boolean = child.exists{
-    // If the reference is found inside the child, referencing to this CTE definition,
-    // and already marked as recursive, then this CTE definition is recursive.
+  lazy val hasSelfReferenceAsCTERef: Boolean = child.exists{
     case CTERelationRef(this.id, _, _, _, _, true) => true
     case _ => false
   }
+  lazy val hasSelfReferenceAsUnionLoopRef: Boolean = child.exists{
+    case UnionLoopRef(this.id, _, _) => true
+    case _ => false
+  }
 }
 
 object CTERelationDef {