[SPARK-46741][SQL] Cache Table with CTE won't work

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/normalizer/WithCTENormalized.scala

@@ -0,0 +1,57 @@
+/*
+ * 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.normalizer
+
+import org.apache.spark.sql.catalyst.expressions.SubqueryExpression
+import org.apache.spark.sql.catalyst.plans.logical.{CacheTableAsSelect, CTERelationRef, LogicalPlan, WithCTE}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern.{CTE, PLAN_EXPRESSION}
+
+object WithCTENormalized extends Rule[LogicalPlan]{
+  override def apply(plan: LogicalPlan): LogicalPlan = {
+    val curId = new java.util.concurrent.atomic.AtomicLong()
+    plan transformDown {
+
+      case ctas @ CacheTableAsSelect(_, plan, _, _, _, _, _) =>
+        ctas.copy(plan = apply(plan))
+
+      case withCTE @ WithCTE(plan, cteDefs) =>
+        val defIdToNewId = withCTE.cteDefs.map(_.id).map((_, curId.getAndIncrement())).toMap
+        val normalizedPlan = canonicalizeCTE(plan, defIdToNewId)
+        val newCteDefs = cteDefs.map { cteDef =>
+          val normalizedCteDef = canonicalizeCTE(cteDef.child, defIdToNewId)
+          cteDef.copy(child = normalizedCteDef, id = defIdToNewId(cteDef.id))
+        }
+        withCTE.copy(plan = normalizedPlan, cteDefs = newCteDefs)
+    }
+  }
+
+  def canonicalizeCTE(plan: LogicalPlan, defIdToNewId: Map[Long, Long]): LogicalPlan = {
+    plan.transformDownWithPruning(
+      _.containsAnyPattern(CTE, PLAN_EXPRESSION)) {
+      // For nested WithCTE, if defIndex didn't contain the cteId,
+      // means it's not current WithCTE's ref.
+      case ref: CTERelationRef if defIdToNewId.contains(ref.cteId) =>
+        ref.copy(cteId = defIdToNewId(ref.cteId))
+      case other =>
+        other.transformExpressionsWithPruning(_.containsPattern(PLAN_EXPRESSION)) {
+          case e: SubqueryExpression => e.withNewPlan(canonicalizeCTE(e.plan, defIdToNewId))
+        }
+    }
+  }
+}

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

@@ -1391,7 +1391,8 @@ case class CacheTableAsSelect(
     isLazy: Boolean,
     options: Map[String, String],
     isAnalyzed: Boolean = false,
-    referredTempFunctions: Seq[String] = Seq.empty) extends AnalysisOnlyCommand {
+    referredTempFunctions: Seq[String] = Seq.empty)
+  extends AnalysisOnlyCommand with CTEInChildren {
   override protected def withNewChildrenInternal(
       newChildren: IndexedSeq[LogicalPlan]): CacheTableAsSelect = {
     assert(!isAnalyzed)
@@ -1406,6 +1407,10 @@ case class CacheTableAsSelect(
       // Collect the referred temporary functions from AnalysisContext
       referredTempFunctions = ac.referredTempFunctionNames.toSeq)
   }
+
+  override def withCTEDefs(cteDefs: Seq[CTERelationDef]): LogicalPlan = {
+    copy(plan = WithCTE(plan, cteDefs))
+  }
 }
 
 /**

sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala

@@ -22,6 +22,7 @@ import org.apache.spark.sql.artifact.ArtifactManager
 import org.apache.spark.sql.catalyst.analysis.{Analyzer, EvalSubqueriesForTimeTravel, FunctionRegistry, ReplaceCharWithVarchar, ResolveSessionCatalog, TableFunctionRegistry}
 import org.apache.spark.sql.catalyst.catalog.{FunctionExpressionBuilder, SessionCatalog}
 import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.normalizer.WithCTENormalized
 import org.apache.spark.sql.catalyst.optimizer.Optimizer
 import org.apache.spark.sql.catalyst.parser.ParserInterface
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
@@ -339,6 +340,7 @@ abstract class BaseSessionStateBuilder(
   }
 
   protected def planNormalizationRules: Seq[Rule[LogicalPlan]] = {
+    WithCTENormalized +:
     extensions.buildPlanNormalizationRules(session)
   }
 

sql/core/src/test/resources/sql-tests/analyzer-results/cache.sql.out

@@ -0,0 +1,184 @@
+-- Automatically generated by SQLQueryTestSuite
+-- !query
+CREATE TEMPORARY VIEW t1 AS SELECT * FROM VALUES (0, 0), (1, 1), (2, 2) AS t(c1, c2)
+-- !query analysis
+CreateViewCommand `t1`, SELECT * FROM VALUES (0, 0), (1, 1), (2, 2) AS t(c1, c2), false, false, LocalTempView, UNSUPPORTED, true
+   +- Project [c1#x, c2#x]
+      +- SubqueryAlias t
+         +- LocalRelation [c1#x, c2#x]
+
+
+-- !query
+CREATE TEMPORARY VIEW t2 AS
+WITH v as (
+  SELECT c1 + c1 c3 FROM t1
+)
+SELECT SUM(c3) s FROM v
+-- !query analysis
+CreateViewCommand `t2`, WITH v as (
+  SELECT c1 + c1 c3 FROM t1
+)
+SELECT SUM(c3) s FROM v, false, false, LocalTempView, UNSUPPORTED, true
+   +- WithCTE
+      :- CTERelationDef xxxx, false
+      :  +- SubqueryAlias v
+      :     +- Project [(c1#x + c1#x) AS c3#x]
+      :        +- SubqueryAlias t1
+      :           +- View (`t1`, [c1#x, c2#x])
+      :              +- Project [cast(c1#x as int) AS c1#x, cast(c2#x as int) AS c2#x]
+      :                 +- Project [c1#x, c2#x]
+      :                    +- SubqueryAlias t
+      :                       +- LocalRelation [c1#x, c2#x]
+      +- Aggregate [sum(c3#x) AS s#xL]
+         +- SubqueryAlias v
+            +- CTERelationRef xxxx, true, [c3#x], false
+
+
+-- !query
+CACHE TABLE cache_table
+WITH
+t2 AS (SELECT 1)
+SELECT * FROM t2
+-- !query analysis
+CacheTableAsSelect cache_table, WITH
+t2 AS (SELECT 1)
+SELECT * FROM t2, false, true
+   +- WithCTE
+      :- CTERelationDef xxxx, false
+      :  +- SubqueryAlias t2
+      :     +- Project [1 AS 1#x]
+      :        +- OneRowRelation
+      +- Project [1#x]
+         +- SubqueryAlias t2
+            +- CTERelationRef xxxx, true, [1#x], false
+
+
+-- !query
+SELECT * FROM cache_table
+-- !query analysis
+Project [1#x]
++- SubqueryAlias cache_table
+   +- View (`cache_table`, [1#x])
+      +- Project [cast(1#x as int) AS 1#x]
+         +- WithCTE
+            :- CTERelationDef xxxx, false
+            :  +- SubqueryAlias t2
+            :     +- Project [1 AS 1#x]
+            :        +- OneRowRelation
+            +- Project [1#x]
+               +- SubqueryAlias t2
+                  +- CTERelationRef xxxx, true, [1#x], false
+
+
+-- !query
+EXPLAIN EXTENDED SELECT * FROM cache_table
+-- !query analysis
+ExplainCommand 'Project [*], ExtendedMode
+
+
+-- !query
+CACHE TABLE cache_nested_cte_table
+WITH
+v AS (
+  SELECT c1 * c2 c3 from t1
+)
+SELECT SUM(c3) FROM v
+EXCEPT
+SELECT s FROM t2
+-- !query analysis
+CacheTableAsSelect cache_nested_cte_table, WITH
+v AS (
+  SELECT c1 * c2 c3 from t1
+)
+SELECT SUM(c3) FROM v
+EXCEPT
+SELECT s FROM t2, false, true
+   +- WithCTE
+      :- CTERelationDef xxxx, false
+      :  +- SubqueryAlias v
+      :     +- Project [(c1#x * c2#x) AS c3#x]
+      :        +- SubqueryAlias t1
+      :           +- View (`t1`, [c1#x, c2#x])
+      :              +- Project [cast(c1#x as int) AS c1#x, cast(c2#x as int) AS c2#x]
+      :                 +- Project [c1#x, c2#x]
+      :                    +- SubqueryAlias t
+      :                       +- LocalRelation [c1#x, c2#x]
+      +- Except false
+         :- Aggregate [sum(c3#x) AS sum(c3)#xL]
+         :  +- SubqueryAlias v
+         :     +- CTERelationRef xxxx, true, [c3#x], false
+         +- Project [s#xL]
+            +- SubqueryAlias t2
+               +- View (`t2`, [s#xL])
+                  +- Project [cast(s#xL as bigint) AS s#xL]
+                     +- WithCTE
+                        :- CTERelationDef xxxx, false
+                        :  +- SubqueryAlias v
+                        :     +- Project [(c1#x + c1#x) AS c3#x]
+                        :        +- SubqueryAlias t1
+                        :           +- View (`t1`, [c1#x, c2#x])
+                        :              +- Project [cast(c1#x as int) AS c1#x, cast(c2#x as int) AS c2#x]
+                        :                 +- Project [c1#x, c2#x]
+                        :                    +- SubqueryAlias t
+                        :                       +- LocalRelation [c1#x, c2#x]
+                        +- Aggregate [sum(c3#x) AS s#xL]
+                           +- SubqueryAlias v
+                              +- CTERelationRef xxxx, true, [c3#x], false
+
+
+-- !query
+SELECT * FROM cache_nested_cte_table
+-- !query analysis
+Project [sum(c3)#xL]
++- SubqueryAlias cache_nested_cte_table
+   +- View (`cache_nested_cte_table`, [sum(c3)#xL])
+      +- Project [cast(sum(c3)#xL as bigint) AS sum(c3)#xL]
+         +- WithCTE
+            :- CTERelationDef xxxx, false
+            :  +- SubqueryAlias v
+            :     +- Project [(c1#x * c2#x) AS c3#x]
+            :        +- SubqueryAlias t1
+            :           +- View (`t1`, [c1#x, c2#x])
+            :              +- Project [cast(c1#x as int) AS c1#x, cast(c2#x as int) AS c2#x]
+            :                 +- Project [c1#x, c2#x]
+            :                    +- SubqueryAlias t
+            :                       +- LocalRelation [c1#x, c2#x]
+            +- Except false
+               :- Aggregate [sum(c3#x) AS sum(c3)#xL]
+               :  +- SubqueryAlias v
+               :     +- CTERelationRef xxxx, true, [c3#x], false
+               +- Project [s#xL]
+                  +- SubqueryAlias t2
+                     +- View (`t2`, [s#xL])
+                        +- Project [cast(s#xL as bigint) AS s#xL]
+                           +- WithCTE
+                              :- CTERelationDef xxxx, false
+                              :  +- SubqueryAlias v
+                              :     +- Project [(c1#x + c1#x) AS c3#x]
+                              :        +- SubqueryAlias t1
+                              :           +- View (`t1`, [c1#x, c2#x])
+                              :              +- Project [cast(c1#x as int) AS c1#x, cast(c2#x as int) AS c2#x]
+                              :                 +- Project [c1#x, c2#x]
+                              :                    +- SubqueryAlias t
+                              :                       +- LocalRelation [c1#x, c2#x]
+                              +- Aggregate [sum(c3#x) AS s#xL]
+                                 +- SubqueryAlias v
+                                    +- CTERelationRef xxxx, true, [c3#x], false
+
+
+-- !query
+EXPLAIN EXTENDED SELECT * FROM cache_nested_cte_table
+-- !query analysis
+ExplainCommand 'Project [*], ExtendedMode
+
+
+-- !query
+DROP TABLE IF EXISTS t1
+-- !query analysis
+DropTempViewCommand t1
+
+
+-- !query
+DROP TABLE IF EXISTS t2
+-- !query analysis
+DropTempViewCommand t2

sql/core/src/test/resources/sql-tests/inputs/cache.sql

@@ -0,0 +1,33 @@
+CREATE TEMPORARY VIEW t1 AS SELECT * FROM VALUES (0, 0), (1, 1), (2, 2) AS t(c1, c2);
+CREATE TEMPORARY VIEW t2 AS
+WITH v as (
+  SELECT c1 + c1 c3 FROM t1
+)
+SELECT SUM(c3) s FROM v;
+
+CACHE TABLE cache_table
+WITH
+t2 AS (SELECT 1)
+SELECT * FROM t2;
+
+SELECT * FROM cache_table;
+
+EXPLAIN EXTENDED SELECT * FROM cache_table;
+
+-- Nested WithCTE
+CACHE TABLE cache_nested_cte_table
+WITH
+v AS (
+  SELECT c1 * c2 c3 from t1
+)
+SELECT SUM(c3) FROM v
+EXCEPT
+SELECT s FROM t2;
+
+SELECT * FROM cache_nested_cte_table;
+
+EXPLAIN EXTENDED SELECT * FROM cache_nested_cte_table;
+
+-- Clean up
+DROP TABLE IF EXISTS t1;
+DROP TABLE IF EXISTS t2;