[SPARK-28228][SQL] Better support for WITH clause

docs/sql-migration-guide-upgrade.md

@@ -147,6 +147,8 @@ license: |
 
   - Since Spark 3.0, if files or subdirectories disappear during recursive directory listing (i.e. they appear in an intermediate listing but then cannot be read or listed during later phases of the recursive directory listing, due to either concurrent file deletions or object store consistency issues) then the listing will fail with an exception unless `spark.sql.files.ignoreMissingFiles` is `true` (default `false`). In previous versions, these missing files or subdirectories would be ignored. Note that this change of behavior only applies during initial table file listing (or during `REFRESH TABLE`), not during query execution: the net change is that `spark.sql.files.ignoreMissingFiles` is now obeyed during table file listing / query planning, not only at query execution time.
 
+  - Since Spark 3.0, substitution order of nested WITH clauses is changed and an inner CTE definition takes precedence over an outer. The previous behaviour can be restored by setting `spark.sql.legacy.cte.substitution.enabled` to `true`.
+
 ## Upgrading from Spark SQL 2.4 to 2.4.1
 
   - The value of `spark.executor.heartbeatInterval`, when specified without units like "30" rather than "30s", was

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

@@ -209,38 +209,6 @@ class Analyzer(
       CleanupAliases)
   )
 
-  /**
-   * Analyze cte definitions and substitute child plan with analyzed cte definitions.
-   */
-  object CTESubstitution extends Rule[LogicalPlan] {
-    def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsUp {
-      case With(child, relations) =>
-        // substitute CTE expressions right-to-left to resolve references to previous CTEs:
-        // with a as (select * from t), b as (select * from a) select * from b
-        relations.foldRight(child) {
-          case ((cteName, ctePlan), currentPlan) =>
-            substituteCTE(currentPlan, cteName, ctePlan)
-        }
-      case other => other
-    }
-
-    private def substituteCTE(
-        plan: LogicalPlan,
-        cteName: String,
-        ctePlan: LogicalPlan): LogicalPlan = {
-      plan resolveOperatorsUp {
-        case UnresolvedRelation(Seq(table)) if resolver(cteName, table) =>
-          ctePlan
-        case other =>
-          // This cannot be done in ResolveSubquery because ResolveSubquery does not know the CTE.
-          other transformExpressions {
-            case e: SubqueryExpression =>
-              e.withNewPlan(substituteCTE(e.plan, cteName, ctePlan))
-          }
-      }
-    }
-  }
-
   /**
    * Substitute child plan with WindowSpecDefinitions.
    */

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

@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.analysis
+
+import org.apache.spark.sql.catalyst.expressions.SubqueryExpression
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, With}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Analyze WITH nodes and substitute child plan with CTE definitions.
+ */
+object CTESubstitution extends Rule[LogicalPlan] {
+  def apply(plan: LogicalPlan): LogicalPlan = if (SQLConf.get.legacyCTESubstitutionEnabled) {
+    legacyTraverseAndSubstituteCTE(plan)
+  } else {
+    traverseAndSubstituteCTE(plan, false)
+  }
+
+  private def legacyTraverseAndSubstituteCTE(plan: LogicalPlan): LogicalPlan = {
+    plan.resolveOperatorsUp {
+      case With(child, relations) =>
+        // substitute CTE expressions right-to-left to resolve references to previous CTEs:
+        // with a as (select * from t), b as (select * from a) select * from b
+        relations.foldRight(child) {
+          case ((cteName, ctePlan), currentPlan) => substituteCTE(currentPlan, cteName, ctePlan)
+        }
+    }
+  }
+
+  /**
+   * Traverse the plan and expression nodes as a tree and replace matching references to CTE
+   * definitions.
+   * - If the rule encounters a WITH node then it substitutes the child of the node with CTE
+   *   definitions of the node right-to-left order as a definition can reference to a previous
+   *   one.
+   *   For example the following query is valid:
+   *   WITH
+   *     t AS (SELECT 1),
+   *     t2 AS (SELECT * FROM t)
+   *   SELECT * FROM t2
+   * - If a CTE definition contains an inner WITH node then substitution of inner should take
+   *   precedence because it can shadow an outer CTE definition.
+   *   For example the following query should return 2:
+   *   WITH
+   *     t AS (SELECT 1),
+   *     t2 AS (
+   *       WITH t AS (SELECT 2)
+   *       SELECT * FROM t
+   *     )
+   *   SELECT * FROM t2
+   * - If a CTE definition contains a subquery that contains an inner WITH node then substitution
+   *   of inner should take precedence because it can shadow an outer CTE definition.
+   *   For example the following query should return 2:
+   *   WITH t AS (SELECT 1 AS c)
+   *   SELECT max(c) FROM (
+   *     WITH t AS (SELECT 2 AS c)
+   *     SELECT * FROM t
+   *   )
+   * - If a CTE definition contains a subquery expression that contains an inner WITH node then
+   *   substitution of inner should take precedence because it can shadow an outer CTE
+   *   definition.
+   *   For example the following query should return 2:
+   *   WITH t AS (SELECT 1)
+   *   SELECT (
+   *     WITH t AS (SELECT 2)
+   *     SELECT * FROM t
+   *   )
+   * @param plan  the plan to be traversed
+   * @param inTraverse whether the current traverse is called from another traverse, only in this
+   *                   case name collision can occur
+   * @return then plan where CTE substitution is applied
+   */
+  private def traverseAndSubstituteCTE(plan: LogicalPlan, inTraverse: Boolean): LogicalPlan = {
+    plan.resolveOperatorsUp {
+      case With(child: LogicalPlan, relations) =>
+        val traversedChild = child transformExpressions {
+          case e: SubqueryExpression => e.withNewPlan(traverseAndSubstituteCTE(e.plan, true))
+        }
+
+        relations.foldRight(traversedChild) {
+          case ((cteName, ctePlan), currentPlan) =>
+            lazy val substitutedCTEPlan = traverseAndSubstituteCTE(ctePlan, true)
+            substituteCTE(currentPlan, cteName, substitutedCTEPlan)
+        }
+
+      case other if inTraverse =>
+        other.transformExpressions {
+          case e: SubqueryExpression => e.withNewPlan(traverseAndSubstituteCTE(e.plan, true))
+        }
+    }
+  }
+
+  // A CTE definition might not be used at all so ctePlan is call by name.
+  private def substituteCTE(
+      plan: LogicalPlan,
+      cteName: String,
+      ctePlan: => LogicalPlan): LogicalPlan =
+    plan resolveOperatorsUp {
+      case UnresolvedRelation(Seq(table)) if plan.conf.resolver(cteName, table) => ctePlan
+
+      case other =>
+        // This cannot be done in ResolveSubquery because ResolveSubquery does not know the CTE.
+        other transformExpressions {
+          case e: SubqueryExpression => e.withNewPlan(substituteCTE(e.plan, cteName, ctePlan))
+        }
+    }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1838,6 +1838,12 @@ object SQLConf {
     .booleanConf
     .createWithDefault(false)
 
+  val LEGACY_CTE_SUBSTITUTION_ENABLED = buildConf("spark.sql.legacy.cte.substitution.enabled")
+    .internal()
+    .doc("When true, outer CTE definitions takes precedence over inner definitions.")
+    .booleanConf
+    .createWithDefault(false)
+
   val LEGACY_ARRAY_EXISTS_FOLLOWS_THREE_VALUED_LOGIC =
     buildConf("spark.sql.legacy.arrayExistsFollowsThreeValuedLogic")
       .doc("When true, the ArrayExists will follow the three-valued boolean logic.")
@@ -2322,6 +2328,8 @@ class SQLConf extends Serializable with Logging {
 
   def defaultV2Catalog: Option[String] = getConf(DEFAULT_V2_CATALOG)
 
+  def legacyCTESubstitutionEnabled: Boolean = getConf(LEGACY_CTE_SUBSTITUTION_ENABLED)
+
   /** ********************** SQLConf functionality methods ************ */
 
   /** Set Spark SQL configuration properties. */

sql/core/src/test/resources/sql-tests/inputs/cte-legacy.sql

@@ -0,0 +1,115 @@
+create temporary view t as select * from values 0, 1, 2 as t(id);
+create temporary view t2 as select * from values 0, 1 as t(id);
+
+-- CTE legacy substitution
+SET spark.sql.legacy.cte.substitution.enabled=true;
+
+-- CTE in CTE definition
+WITH t as (
+  WITH t2 AS (SELECT 1)
+  SELECT * FROM t2
+)
+SELECT * FROM t;
+
+-- CTE in subquery
+SELECT max(c) FROM (
+  WITH t(c) AS (SELECT 1)
+  SELECT * FROM t
+);
+
+-- CTE in subquery expression
+SELECT (
+  WITH t AS (SELECT 1)
+  SELECT * FROM t
+);
+
+-- CTE in CTE definition shadows outer
+WITH
+  t AS (SELECT 1),
+  t2 AS (
+    WITH t AS (SELECT 2)
+    SELECT * FROM t
+  )
+SELECT * FROM t2;
+
+-- CTE in CTE definition shadows outer 2
+WITH
+  t(c) AS (SELECT 1),
+  t2 AS (
+    SELECT (
+      SELECT max(c) FROM (
+        WITH t(c) AS (SELECT 2)
+        SELECT * FROM t
+      )
+    )
+  )
+SELECT * FROM t2;
+
+-- CTE in CTE definition shadows outer 3
+WITH
+  t AS (SELECT 1),
+  t2 AS (
+    WITH t AS (SELECT 2),
+    t2 AS (
+      WITH t AS (SELECT 3)
+      SELECT * FROM t
+    )
+    SELECT * FROM t2
+  )
+SELECT * FROM t2;
+
+-- CTE in subquery shadows outer
+WITH t(c) AS (SELECT 1)
+SELECT max(c) FROM (
+  WITH t(c) AS (SELECT 2)
+  SELECT * FROM t
+);
+
+-- CTE in subquery shadows outer 2
+WITH t(c) AS (SELECT 1)
+SELECT sum(c) FROM (
+  SELECT max(c) AS c FROM (
+    WITH t(c) AS (SELECT 2)
+    SELECT * FROM t
+  )
+);
+
+-- CTE in subquery shadows outer 3
+WITH t(c) AS (SELECT 1)
+SELECT sum(c) FROM (
+  WITH t(c) AS (SELECT 2)
+  SELECT max(c) AS c FROM (
+    WITH t(c) AS (SELECT 3)
+    SELECT * FROM t
+  )
+);
+
+-- CTE in subquery expression shadows outer
+WITH t AS (SELECT 1)
+SELECT (
+  WITH t AS (SELECT 2)
+  SELECT * FROM t
+);
+
+-- CTE in subquery expression shadows outer 2
+WITH t AS (SELECT 1)
+SELECT (
+  SELECT (
+    WITH t AS (SELECT 2)
+    SELECT * FROM t
+  )
+);
+
+-- CTE in subquery expression shadows outer 3
+WITH t AS (SELECT 1)
+SELECT (
+  WITH t AS (SELECT 2)
+  SELECT (
+    WITH t AS (SELECT 3)
+    SELECT * FROM t
+  )
+);
+
+-- Clean up
+DROP VIEW IF EXISTS t;
+DROP VIEW IF EXISTS t2;