From 185d1d22c3ae65b414a5079bf8ba84fc845faf38 Mon Sep 17 00:00:00 2001 From: Kush Vyas Date: Sun, 19 Jul 2026 14:20:44 +0530 Subject: [PATCH] [SPARK-57109][SQL] Fix exchange reuse failure for VIEW in UNION ALL branches via CTE --- .../spark/sql/catalyst/plans/QueryPlan.scala | 20 +++++++++++++ .../sql/catalyst/plans/QueryPlanSuite.scala | 19 +++++++++++++ .../ReuseExchangeAndSubquerySuite.scala | 28 ++++++++++++++++++- 3 files changed, 66 insertions(+), 1 deletion(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala index 8ca3f76c3b885..b746b8dc3f1ff 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala @@ -709,6 +709,26 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]] var id = -1 val allAttributesSeq = this.allAttributes mapExpressions { + // A no-op rename alias - i.e. `Alias(attr, attr.name)` with metadata that matches the + // child attribute's metadata - is exactly the shape that `RemoveRedundantAliases` strips + // whenever the underlying attribute is not on its exclude list. Such an alias can still + // survive to this point even though it is semantically redundant, e.g. because it sits on + // the first child of a `Union` (SPARK-39887 protects the first child's attributes for + // correctness) while the exact same alias was stripped from a sibling branch that + // references the same underlying attribute (e.g. via a de-duplicated CTE/self-join). + // That leaves two semantically-identical branches with different shapes: one still + // wrapped in the alias, the other already a bare `AttributeReference`. Left alone, this + // asymmetry makes their canonical forms differ, which defeats exchange/subquery reuse. + // Canonicalize the no-op alias exactly like the bare attribute it is equivalent to, so + // both shapes converge on the same canonical form. + case a @ Alias(attr: Attribute, name) if name == attr.name && a.metadata == attr.metadata => + if (allAttributesSeq.indexOf(attr.exprId) == -1) { + id += 1 + attr.withExprId(ExprId(id)).canonicalized + } else { + QueryPlan.normalizeExpressions(attr, allAttributesSeq) + } + case a: Alias => id += 1 // As the root of the expression, Alias will always take an arbitrary exprId, we need to diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala index 68c0e354950f3..39e5f90dd7c78 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala @@ -200,4 +200,23 @@ class QueryPlanSuite extends SparkFunSuite { assert(normalizedExpressions.map(_.dataType) == Seq(IntegerType, StringType), "normalizeExpressions preserves original expression order") } + + test("SPARK-57109: no-op rename Alias canonicalizes the same as a bare attribute") { + val relation = LocalRelation(AttributeReference("a", IntegerType)(ExprId(0))) + val a = relation.output.head + + // Branch 1: the attribute is still wrapped in a no-op rename alias, e.g. because it is + // protected from RemoveRedundantAliases (the Union first-child case in SPARK-39887, or a + // self-join dedup alias). + val aliasedBranch = Project(Seq(Alias(a, a.name)(ExprId(1))), relation) + + // Branch 2: the exact same underlying attribute, but the no-op alias has already been + // stripped down to a bare AttributeReference (e.g. because RemoveRedundantAliases was free + // to remove it here). + val bareBranch = Project(Seq(a), relation) + + assert(aliasedBranch.sameResult(bareBranch), + "a no-op rename alias and the bare attribute it wraps must canonicalize identically, " + + "or ReuseExchange/subquery reuse cannot match semantically equivalent Union branches") + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReuseExchangeAndSubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReuseExchangeAndSubquerySuite.scala index 16726547e3247..f446b517bb5bf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReuseExchangeAndSubquerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReuseExchangeAndSubquerySuite.scala @@ -17,12 +17,13 @@ package org.apache.spark.sql.execution +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} import org.apache.spark.sql.functions.col import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -class ReuseExchangeAndSubquerySuite extends SharedSparkSession { +class ReuseExchangeAndSubquerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper { val tableFormat: String = "parquet" @@ -112,4 +113,29 @@ class ReuseExchangeAndSubquerySuite extends SharedSparkSession { } } } + + test("SPARK-57109: Exchange reused when a VIEW is referenced via a CTE in both UNION branches") { + withTable("t1", "t2") { + withView("v") { + sql(s"CREATE TABLE t1 (a BIGINT, b INT, c INT) USING $tableFormat") + sql(s"CREATE TABLE t2 (a BIGINT, d BIGINT) USING $tableFormat") + sql("INSERT INTO t1 VALUES (1, 2, 3), (4, 5, 6)") + sql("INSERT INTO t2 VALUES (1, 7), (4, 8)") + sql("CREATE OR REPLACE VIEW v AS SELECT a, b FROM t1 WHERE c <> 6") + + val df = sql( + """ + |WITH cte AS (SELECT ta.d, tv.b FROM t2 ta JOIN v tv ON ta.a = tv.a) + |SELECT d, b, count(1) FROM cte GROUP BY 1, 2 + |UNION ALL + |SELECT d, b, count(1) FROM cte GROUP BY 1, 2 + |""".stripMargin) + + val plan = df.queryExecution.executedPlan + val reusedExchanges = plan.collectWithSubqueries { case re: ReusedExchangeExec => re } + assert(reusedExchanges.nonEmpty, + "expected the t1 shuffle to be reused across both UNION branches") + } + } + } }