peter-toth commented on code in PR #58661: URL: https://github.com/apache/spark/pull/58661#discussion_r3979312119
########## 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: **Finding 5.** Closing the boundary puts a referenced definition's body wherever the reference sits. When the reference sits inside a subquery of the materialized definition, that definition's `OuterReference` lands inside the subquery plan — exactly where the scan below does not look. So the comment on the next line does not hold: it is true for subqueries the definition wrote itself, not for a definition inlined into one. Three shapes, all accepted by analysis on `0364b66b093`, with `t` a two-column temp view: ```sql -- scalar subquery SELECT * FROM t o WHERE EXISTS ( WITH s AS (SELECT i.c1 FROM t i WHERE i.c1 = o.c1), v AS MATERIALIZED (SELECT a.c1, (SELECT count(*) FROM s) AS n FROM t a) SELECT * FROM v) -- LATERAL SELECT * FROM t o WHERE EXISTS ( WITH s AS (SELECT i.c1 FROM t i WHERE i.c1 = o.c1), v AS MATERIALIZED (SELECT a.c1, l.n FROM t a, LATERAL (SELECT count(*) n FROM s) l) SELECT * FROM v) -- IN subquery SELECT * FROM t o WHERE EXISTS ( WITH s AS (SELECT i.c1 FROM t i WHERE i.c1 = o.c1), v AS MATERIALIZED (SELECT a.c1 FROM t a WHERE a.c1 IN (SELECT c1 FROM s)) SELECT * FROM v) ``` The optimized plan puts the materializing `RepartitionByExpression` above a subtree that still holds `outer(c1#9)`. The last two die at execution with `[INTERNAL_ERROR] Cannot generate code for expression: outer(c1#9)`; the first fails inside `awaitResult`. Dropping `MATERIALIZED` returns rows. Both analyzers accept all three, so the dual-run hybrid analyzer does not flag it either — on `7e742bb73a4` the single-pass resolver still rejected them. The scan has to reach each referenced definition's own operator tree. `closed` already is that tree, since `WithCTE.children` is `cteDefs :+ plan`, so the smallest fix is to scan it and drop the inlining: ```scala val closed = if (referencedDefs.isEmpty) cteDef.child else WithCTE(cteDef.child, referencedDefs) // The scan stays out of subquery plans on purpose: a correlation a definition keeps inside // its own subquery targets that definition, not the enclosing query. closed.foreach(_.expressions.foreach(_.foreach { ... ``` (the `InlineCTE` call and its import then go away). I measured that arm: all three queries are rejected with `MATERIALIZED_CTE_WITH_OUTER_REFERENCE`, and `AnalysisErrorSuite`, `InlineCTESuite`, `CTEInlineSuiteAEOn` and `CTEInlineSuiteAEOff` are green, 188 tests. The `nested` filter in `collectReferencedDefs` still earns its keep — it keeps a definition nested in the body from being hoisted into the top-level `WithCTE`, which is what keeps the inner-correlated-CTE cases passing. Worth adding all three as negatives. Every current negative correlates from the definition's own operator level, so the suite cannot tell "the definition is correlated" from "something the definition references is correlated". ########## 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: **Finding 6.** `MATERIALIZED_CTE_ALWAYS_INLINED` is a hard error reachable from the syntax this page documents, and nothing on the page hints at it: ```sql WITH v AS MATERIALIZED (SELECT c1 FROM t) FROM v INSERT INTO a SELECT c1 INSERT INTO b SELECT c1 ``` The reply at [r3976569156](https://github.com/apache/spark/pull/58661#discussion_r3976569156) sets the bar for this page as observable behaviour, and an error is observable. It also does not need implementation terms — a sentence beside the outer-query restriction is enough: `MATERIALIZED` is not supported in a statement whose common table expressions are always inlined, such as a multiple-INSERT statement. `NOT MATERIALIZED` is 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: **Finding 3.** Gating on `cteDef.resolved` covers a typo inside the definition, but the call still sits above `checkAnalysis0`, so it preempts an unresolved column anywhere else in the query: ```sql SELECT * FROM t o WHERE EXISTS ( WITH v AS MATERIALIZED (SELECT 1 FROM t i WHERE i.c1 = o.c1) SELECT * FROM v ) AND no_such_col = 1 ``` Measured on `0364b66b093`: fixed-point reports `MATERIALIZED_CTE_WITH_OUTER_REFERENCE` naming `c1`, single-pass reports `UNRESOLVED_COLUMN.WITH_SUGGESTION` for `no_such_col`. Same query, two errors, and the useful one is the second. Dropping `MATERIALIZED` gives `UNRESOLVED_COLUMN` on both. The constraint is the un-inlined plan, not the position in the method — `plan` is still around after `checkAnalysis0`: ```scala } finally { preemptedError.clear() } // Check MATERIALIZED CTE relations on the original plan, as `inlinedPlan` has them inlined. MaterializedCTECheck(plan) plan.setAnalyzed() ``` I measured that: the query above reports `UNRESOLVED_COLUMN` on both analyzers, and `AnalysisErrorSuite`, `CTEInlineSuiteAEOn` and `CTEInlineSuiteAEOff` stay green, 179 tests. Keeping the call outside the `try` also keeps the error a plain `AnalysisException`, which is what the golden files record. ########## 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: **Finding 7.** This is the only place `RECURSIVE` meets the option, and it stops at the parse tree. I ran both end to end on `0364b66b093` and they behave: `WITH RECURSIVE r(n) AS MATERIALIZED (SELECT 1 UNION ALL SELECT n + 1 FROM r WHERE n < 5) SELECT * FROM r a JOIN r b ON a.n = b.n` returns 1..5 paired, with one `ReusedExchange` over the `UnionLoop`, and `NOT MATERIALIZED` inlines the loop. Since the recursive semantics were a deliberate call at [r3976569502](https://github.com/apache/spark/pull/58661#discussion_r3976569502), a `CTEInlineSuite` case pinning them is cheap and stops a later `InlineCTE` change from flipping it silently. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala: ########## @@ -288,29 +292,44 @@ trait CTEInChildren extends LogicalPlan { } } +/** + * A named common table expression (CTE) as defined in a WITH clause, before analysis. + * + * @param name The name of the CTE. + * @param plan The CTE definition query plan, aliased by `name`. + * @param maxDepth The optional `MAX RECURSION LEVEL` of a recursive CTE. + * @param materialized The materialization option: `Some(true)` for `MATERIALIZED`, + * `Some(false)` for `NOT MATERIALIZED`, `None` if unspecified. + */ +case class CTERelation( Review Comment: **Finding 8.** `CTERelation` sits between `CTERelationDef` and `CTERelationRef` in this file but is not a plan node, and `CTESubstitution` now matches `case CTERelation(name, relation, _, _)` a few lines from a `CTERelationDef(...)` construction. `UnresolvedCTERelation` would say which side of analysis it lives on, matching the `UnresolvedWith` that holds it. -- 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]
