You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/04/15 01:39:09 UTC
(spark) branch master updated: [SPARK-52818][SQL] Fix MergeSubplans creating nested WithCTE with cross-scope CTE references
This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 8a2b33ae1255 [SPARK-52818][SQL] Fix MergeSubplans creating nested WithCTE with cross-scope CTE references
8a2b33ae1255 is described below
commit 8a2b33ae1255630383a0a2586f4fc8acb4f51382
Author: Wenchen Fan <we...@databricks.com>
AuthorDate: Wed Apr 15 09:38:51 2026 +0800
[SPARK-52818][SQL] Fix MergeSubplans creating nested WithCTE with cross-scope CTE references
### What changes were proposed in this pull request?
When a non-deterministic CTE (e.g. using `monotonically_increasing_id()`) is referenced in scalar subqueries and the result is displayed with `.show()` (which adds a `Limit`), `MergeSubplans` can create an outer `WithCTE` whose CTE defs reference CTE defs from an inner `WithCTE`. This causes `ReplaceCTERefWithRepartition` to crash with `NoSuchElementException` because it processes the outer CTE defs before the inner CTE defs have been added to the map.
The root cause: `MergeSubplans.apply` checks `case _: WithCTE => plan` to skip plans with CTEs, but this only matches when `WithCTE` is the **top-level** node. When `.show()` wraps the plan in `GlobalLimit(LocalLimit(WithCTE(...)))`, the top-level node is `GlobalLimit`, so `MergeSubplans` runs and creates a new outer `WithCTE` around the existing inner one — producing nested `WithCTE` nodes with cross-scope references.
**Repro:**
```sql
WITH cte_base AS (
SELECT recid, givenname, surname,
concat('t_', cast(monotonically_increasing_id() AS string)) AS __unique_id
FROM t
),
cte_tf_givenname AS (
SELECT givenname,
CAST(COUNT(*) AS DOUBLE) / (SELECT COUNT(givenname) FROM cte_base) AS tf
FROM cte_base WHERE givenname IS NOT NULL GROUP BY givenname
),
cte_tf_surname AS (
SELECT surname,
CAST(COUNT(*) AS DOUBLE) / (SELECT COUNT(surname) FROM cte_base) AS tf
FROM cte_base WHERE surname IS NOT NULL GROUP BY surname
)
SELECT cte_base.givenname, cte_base.surname,
cte_tf_givenname.tf, cte_tf_surname.tf
FROM cte_base
LEFT JOIN cte_tf_givenname ON cte_base.givenname = cte_tf_givenname.givenname
LEFT JOIN cte_tf_surname ON cte_base.surname = cte_tf_surname.surname
```
Calling `.show()` on this crashes. Calling `.collect()` works (no `Limit` wrapper).
The fix has two parts:
1. **`MergeSubplans`**: Change `case _: WithCTE => plan` to `case _ if plan.containsPattern(CTE) => plan` to skip plans that contain `WithCTE` anywhere in the tree, not just at the top level.
2. **`ReplaceCTERefWithRepartition`**: improve error message
### Why are the changes needed?
Bug fix. Queries with non-deterministic CTEs referenced in scalar subqueries crash with `java.util.NoSuchElementException: key not found` when using `.show()`.
### Does this PR introduce _any_ user-facing change?
Yes, queries that previously crashed now work correctly.
### How was this patch tested?
New tests added:
- `CTEInlineSuite`: end-to-end test with the exact repro query using `.show()`
- `InlineCTESuite`: unit test verifying `ReplaceCTERefWithRepartition` handles orphaned `CTERelationRef` gracefully
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #55338 from cloud-fan/fix-cte-merge-scalar-subqueries.
Authored-by: Wenchen Fan <we...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../sql/catalyst/optimizer/MergeSubplans.scala | 9 +++++---
.../optimizer/ReplaceCTERefWithRepartition.scala | 6 ++++-
.../sql/catalyst/optimizer/InlineCTESuite.scala | 11 +++++++++
.../org/apache/spark/sql/CTEInlineSuite.scala | 27 ++++++++++++++++++++++
4 files changed, 49 insertions(+), 4 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala
index 94b7da8da21f..151e9e49e781 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala
@@ -22,7 +22,7 @@ import scala.collection.mutable.ArrayBuffer
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, CTERelationDef, CTERelationRef, LeafNode, LogicalPlan, OneRowRelation, Project, Subquery, WithCTE}
import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY, SCALAR_SUBQUERY_REFERENCE, TreePattern}
+import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, CTE, NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY, SCALAR_SUBQUERY_REFERENCE, TreePattern}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.DataType
@@ -157,8 +157,11 @@ object MergeSubplans extends Rule[LogicalPlan] {
// This rule does a whole plan traversal, no need to run on subqueries.
case _: Subquery => plan
- // Plans with CTEs are not supported for now.
- case _: WithCTE => plan
+ // Plans with CTEs are not supported for now. We check containsPattern instead of
+ // just the top-level node because a Limit or other operator may wrap the WithCTE
+ // (e.g. from Dataset.show()), and creating an outer WithCTE with CTE defs that
+ // reference inner CTE defs causes ReplaceCTERefWithRepartition to crash.
+ case _ if plan.containsPattern(CTE) => plan
case _ => extractCommonScalarSubqueries(plan)
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCTERefWithRepartition.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCTERefWithRepartition.scala
index c01372c71ab0..7949d6f3f6b0 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCTERefWithRepartition.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReplaceCTERefWithRepartition.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.optimizer
import scala.collection.mutable
+import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.analysis.DeduplicateRelations
import org.apache.spark.sql.catalyst.expressions.{Alias, SubqueryExpression}
import org.apache.spark.sql.catalyst.plans.Inner
@@ -66,7 +67,10 @@ object ReplaceCTERefWithRepartition extends Rule[LogicalPlan] {
replaceWithRepartition(child, cteMap)
case ref: CTERelationRef =>
- val cteDefPlan = cteMap(ref.cteId)
+ val cteDefPlan = cteMap.getOrElse(ref.cteId,
+ throw SparkException.internalError(
+ s"No CTERelationDef found for CTERelationRef(cteId=${ref.cteId})."))
+
if (ref.outputSet == cteDefPlan.outputSet) {
cteDefPlan
} else {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala
index 9d775a5335c6..e23803c6f688 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala
@@ -17,6 +17,7 @@
package org.apache.spark.sql.catalyst.optimizer
+import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.analysis.TestRelation
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.dsl.plans._
@@ -39,4 +40,14 @@ class InlineCTESuite extends PlanTest {
).analyze
comparePlans(Optimize.execute(plan), plan)
}
+
+ test("SPARK-52818: ReplaceCTERefWithRepartition asserts on orphaned CTERelationRef") {
+ val cteDef = CTERelationDef(OneRowRelation().select(rand(0).as("a")))
+ val cteRef = CTERelationRef(cteDef.id, cteDef.resolved, cteDef.output, cteDef.isStreaming)
+ // A CTERelationRef without an enclosing WithCTE should produce a clear error.
+ val e = intercept[SparkException] {
+ ReplaceCTERefWithRepartition(cteRef.select($"a"))
+ }
+ assert(e.getMessage.contains("No CTERelationDef found"))
+ }
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala
index c1a001117be9..1262d0f491f3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CTEInlineSuite.scala
@@ -847,6 +847,33 @@ abstract class CTEInlineSuiteBase
case _ => false
})
}
+
+ test("SPARK-52818: MergeSubplans should not create nested WithCTE with cross-scope refs") {
+ // A non-deterministic CTE referenced in multiple scalar subqueries is not inlined, leaving
+ // a WithCTE. .show() adds a Limit, so the top node is not WithCTE and MergeSubplans runs,
+ // merging scalar subqueries into a new outer WithCTE whose CTE defs reference the inner
+ // WithCTE's defs. ReplaceCTERefWithRepartition then crashes processing the outer defs
+ // before the inner ones are in the map.
+ withTempView("t") {
+ Seq(("a", "b"), ("c", "d")).toDF("c1", "c2").createOrReplaceTempView("t")
+ sql(
+ """WITH cte AS (
+ | SELECT c1, c2, monotonically_increasing_id() AS id FROM t
+ |),
+ |agg1 AS (
+ | SELECT c1, count(*) / (SELECT count(c1) FROM cte) AS r
+ | FROM cte WHERE c1 IS NOT NULL GROUP BY c1
+ |),
+ |agg2 AS (
+ | SELECT c2, count(*) / (SELECT count(c2) FROM cte) AS r
+ | FROM cte WHERE c2 IS NOT NULL GROUP BY c2
+ |)
+ |SELECT b.c1, a1.r, a2.r FROM cte b
+ |LEFT JOIN agg1 a1 ON b.c1 = a1.c1
+ |LEFT JOIN agg2 a2 ON b.c2 = a2.c2
+ |""".stripMargin).show()
+ }
+ }
}
class CTEInlineSuiteAEOff extends CTEInlineSuiteBase with DisableAdaptiveExecutionSuite
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org