pan3793 commented on code in PR #58661: URL: https://github.com/apache/spark/pull/58661#discussion_r3980641775
########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/MaterializedCTECheck.scala: ########## @@ -0,0 +1,100 @@ +/* + * 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 scala.collection.mutable + +import org.apache.spark.sql.catalyst.expressions.{OuterReference, OuterScopeReference, SubqueryExpression} +import org.apache.spark.sql.catalyst.optimizer.InlineCTE +import org.apache.spark.sql.catalyst.plans.logical.{CTERelationDef, CTERelationRef, LogicalPlan, WithCTE} +import org.apache.spark.sql.catalyst.trees.TreePattern.CTE +import org.apache.spark.sql.errors.QueryCompilationErrors + +/** + * Checks that a MATERIALIZED CTE does not reference the query enclosing it, as it is evaluated + * once on its own. The CTEs it references, transitively, are inlined into it for the check, unless + * they are MATERIALIZED themselves and form their own boundary. Once the boundary is closed this + * way, a correlation kept inside a nested subquery stays inside it, so every outer reference left + * at the operator level crosses the boundary. Unresolved definitions are skipped: the analysis + * checks that follow report them. The check covers the given plan and all its subqueries. + */ +object MaterializedCTECheck extends (LogicalPlan => Unit) { + override def apply(plan: LogicalPlan): Unit = { + if (plan.containsPattern(CTE)) { + // All CTE definitions, including those of nested subqueries, so that references from a + // MATERIALIZED CTE can be followed across subquery boundaries. + val cteDefs = mutable.LinkedHashMap.empty[Long, CTERelationDef] + plan.foreachWithSubqueries { + case cteDef: CTERelationDef => cteDefs(cteDef.id) = cteDef + case _ => + } + cteDefs.values.filter(d => d.materialized.contains(true) && d.resolved).foreach { cteDef => + checkMaterializedCTE(cteDef, cteDefs) + } + } + } + + private def checkMaterializedCTE( + cteDef: CTERelationDef, + cteDefs: collection.Map[Long, CTERelationDef]): Unit = { + val referencedDefs = collectReferencedDefs(cteDef, cteDefs) + val closed = if (referencedDefs.isEmpty) cteDef.child else WithCTE(cteDef.child, referencedDefs) + val boundary = InlineCTE(alwaysInline = true, isAnalysis = true).apply(closed) Review Comment: Fixed in 1c6d2a9d8a6 as suggested: no inlining, the check scans the materialized definition and the outside definitions it references transitively at their own operator level. The `nested` filter stays so a definition nested in a body is not hoisted. The three shapes are negatives in `CTEInlineSuite` and `cte.sql`, plus a unit case with the reference inside a scalar subquery. ########## docs/sql-ref-syntax-qry-select-cte.md: ########## @@ -40,6 +40,16 @@ expression_name [ ( column_name [ , ... ] ) ] [ AS ] ( query ) Specifies a name for the common table expression. +* **MATERIALIZED**, **NOT MATERIALIZED** + + Optionally specifies how the common table expression is evaluated. `MATERIALIZED` forces it + to be evaluated once and shared by all references. `NOT MATERIALIZED` forces it to be inlined, + so that each reference is planned and evaluated independently, and non-deterministic + expressions such as `rand()` may yield different values per reference. A `MATERIALIZED` + common table expression cannot reference columns of an outer query. Omit both to let Spark Review Comment: Added in 1c6d2a9d8a6, in your wording, next to the outer-query restriction; `NOT MATERIALIZED` is stated as supported there. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala: ########## @@ -316,6 +316,9 @@ trait CheckAnalysis extends LookupCatalog with QueryErrorsBase with PlanToString } def checkAnalysis(plan: LogicalPlan): Unit = { + // Check MATERIALIZED CTE relations before inlining, as `InlineCTE` below inlines them to + // restore the original plan shape. + MaterializedCTECheck(plan) Review Comment: Moved in 1c6d2a9d8a6 to after `checkAnalysis0`, outside the `try`, on the original plan, as you laid out. The resolved gate is gone since everything is resolved at that point. `CTEInlineSuite` pins an unresolved column inside the definition and one elsewhere in the query, both reporting `UNRESOLVED_COLUMN`. The unit cases call `MaterializedCTECheck` directly. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala: ########## @@ -2165,6 +2165,45 @@ class PlanParserSuite extends AnalysisTest { cte(table("t").select(star()), false, "t" -> ((table("a").select($"c"), Seq("x"))))) } + test("CTE with materialization option") { + def cteWithOption(materialized: Option[Boolean]): UnresolvedWith = { + UnresolvedWith( + table("t").select(star()), + Seq(CTERelation( + "t", SubqueryAlias("t", table("a").select($"c")), materialized = materialized))) + } + assertEqual( + "WITH t AS MATERIALIZED (SELECT c FROM a) SELECT * FROM t", + cteWithOption(Some(true))) + assertEqual( + "WITH t AS NOT MATERIALIZED (SELECT c FROM a) SELECT * FROM t", + cteWithOption(Some(false))) + // AS is optional. + assertEqual( + "WITH t MATERIALIZED (SELECT c FROM a) SELECT * FROM t", + cteWithOption(Some(true))) + assertEqual( + "WITH t NOT MATERIALIZED (SELECT c FROM a) SELECT * FROM t", + cteWithOption(Some(false))) + // Combined with column aliases and recursion options. + assertEqual( + "WITH RECURSIVE r(x) MAX RECURSION LEVEL 5 AS MATERIALIZED (SELECT c FROM a) " + Review Comment: Added in 1c6d2a9d8a6: `CTEInlineSuite` runs your recursive query with both options, `MATERIALIZED` with one `ReusedExchangeExec` over the loop and `NOT MATERIALIZED` with no repartition. -- 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]
