fix

Author: mihailotim-db

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DeduplicateRelations.scala

@@ -59,7 +59,7 @@ object DeduplicateRelations extends Rule[LogicalPlan] {
       case e @ Except(left, right, _) if !e.duplicateResolved && noMissingInput(right) =>
         e.copy(right = dedupRight(left, right))
       // Only after we finish by-name resolution for Union
-      case u: Union if !u.byName && !u.duplicateResolved =>
+      case u: Union if !u.byName && !u.duplicatesResolvedBetweenBranches =>
         val unionWithChildOutputsDeduplicated =
           DeduplicateUnionChildOutput.deduplicateOutputPerChild(u)
         // Use projection-based de-duplication for Union to avoid breaking the checkpoint sharing

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
+import java.util.HashSet
+
 import scala.collection.mutable
 
 import org.apache.spark.SparkException
@@ -895,6 +897,24 @@ object LimitPushDown extends Rule[LogicalPlan] {
  */
 object PushProjectionThroughUnion extends Rule[LogicalPlan] {
 
+  /**
+   * When pushing a [[Project]] through [[Union]] we need to maintain the invariant that
+   * [[Union]] children must have unique [[ExprId]]s per branch.
+   */
+  private def deduplicateProjectList(projectList: Seq[NamedExpression]) = {
+    val existingExprIds = new HashSet[ExprId]
+    projectList.map(attr => if (existingExprIds.contains(attr.exprId)) {
+      val newMetadata = new MetadataBuilder()
+        .withMetadata(attr.metadata)
+        .putNull("__is_duplicate")
+        .build()
+      Alias(attr, attr.name)(explicitMetadata = Some(newMetadata))
+    } else {
+      existingExprIds.add(attr.exprId)
+      attr
+    })
+  }
+
   /**
    * Maps Attributes from the left side to the corresponding Attribute on the right side.
    */
@@ -923,10 +943,11 @@ object PushProjectionThroughUnion extends Rule[LogicalPlan] {
   }
 
   def pushProjectionThroughUnion(projectList: Seq[NamedExpression], u: Union): Seq[LogicalPlan] = {
-    val newFirstChild = Project(projectList, u.children.head)
+    val deduplicatedProjectList = deduplicateProjectList(projectList)
+    val newFirstChild = Project(deduplicatedProjectList, u.children.head)
     val newOtherChildren = u.children.tail.map { child =>
       val rewrites = buildRewrites(u.children.head, child)
-      Project(projectList.map(pushToRight(_, rewrites)), child)
+      Project(deduplicatedProjectList.map(pushToRight(_, rewrites)), child)
     }
     newFirstChild +: newOtherChildren
   }

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

@@ -614,13 +614,20 @@ case class Union(
     Some(sum.toLong)
   }
 
-  def duplicateResolved: Boolean = {
+  private def duplicatesResolvedPerBranch: Boolean =
+    children.forall(child => child.outputSet.size == child.output.size)
+
+  def duplicatesResolvedBetweenBranches: Boolean = {
     children.map(_.outputSet.size).sum ==
       AttributeSet.fromAttributeSets(children.map(_.outputSet)).size
   }
 
   override lazy val resolved: Boolean = {
-    children.length > 1 && !(byName || allowMissingCol) && childrenResolved && allChildrenCompatible
+    children.length > 1 &&
+    !(byName || allowMissingCol) &&
+    childrenResolved &&
+    allChildrenCompatible &&
+    duplicatesResolvedPerBranch
   }
 
   override protected def withNewChildrenInternal(newChildren: IndexedSeq[LogicalPlan]): Union =

sql/connect/common/src/test/resources/query-tests/explain-results/union.explain

@@ -1,3 +1,3 @@
-Union false, false
+'Union false, false
 :- LocalRelation <empty>, [id#0L, a#0, b#0]
 +- LocalRelation <empty>, [id#0L, a#0, b#0]

sql/connect/common/src/test/resources/query-tests/explain-results/unionAll.explain

@@ -1,3 +1,3 @@
-Union false, false
+'Union false, false
 :- LocalRelation <empty>, [id#0L, a#0, b#0]
 +- LocalRelation <empty>, [id#0L, a#0, b#0]

sql/connect/common/src/test/resources/query-tests/explain-results/unionByName.explain

@@ -1,4 +1,4 @@
-Union false, false
+'Union false, false
 :- Project [id#0L, a#0]
 :  +- LocalRelation <empty>, [id#0L, a#0, b#0]
 +- Project [id#0L, a#0]

sql/connect/common/src/test/resources/query-tests/explain-results/unionByName_allowMissingColumns.explain

@@ -1,4 +1,4 @@
-Union false, false
+'Union false, false
 :- Project [id#0L, a#0, b#0, null AS payload#0]
 :  +- LocalRelation <empty>, [id#0L, a#0, b#0]
 +- Project [id#0L, a#0, null AS b#0, payload#0]

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -4963,6 +4963,23 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
     )
   }
 
+  test("SPARK-52686: Union should be resolved only if there are no duplicates") {
+    withTable("t1") {
+      withView("v1") {
+        sql("CREATE TABLE t1(col1 STRING)")
+        sql("CREATE VIEW v1 AS SELECT * FROM t1")
+        val analyzedPlan = sql(
+          "SELECT * FROM (SELECT col1, col1 FROM v1 UNION SELECT col1, col1 FROM v1)"
+        ).queryExecution.analyzed
+
+        // Resolving * should wait for Union output to be deduplicated, otherwise we are
+        // left with duplicate ExprIds in the result.
+        val exprIds = analyzedPlan.asInstanceOf[Project].projectList.map(_.exprId)
+        assert(exprIds.size == exprIds.distinct.size)
+      }
+    }
+  }
+
   Seq(true, false).foreach { codegenEnabled =>
     test(s"SPARK-52060: one row relation with codegen enabled - $codegenEnabled") {
       withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegenEnabled.toString) {