[SQL] Preserve incoming relations when deduplicating unions

Author: bhollis-dbx

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

@@ -243,12 +243,20 @@ object DeduplicateRelations extends Rule[LogicalPlan] {
   private def deduplicate(
       existingRelations: ExprIdMap,
       plan: LogicalPlan): (LogicalPlan, Boolean) = {
+    def copyExistingRelations(): ExprIdMap = {
+      val copied = mutable.HashMap.empty[Class[_], mutable.HashSet[Long]]
+      existingRelations.foreach { case (relationClass, exprIds) =>
+        copied.put(relationClass, exprIds.clone())
+      }
+      copied
+    }
+
     var planChanged = false
     val newPlan = if (plan.children.nonEmpty) {
       val newChildren = mutable.ArrayBuffer.empty[LogicalPlan]
       for (c <- plan.children) {
         val relations = plan match {
-          case _: Union => mutable.HashMap.empty[Class[_], mutable.HashSet[Long]]
+          case _: Union => copyExistingRelations()
           case _ => existingRelations
         }
         val (renewed, changed) = renewDuplicatedRelations(relations, c)

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala

@@ -728,6 +728,20 @@ class AnalysisSuite extends AnalysisTest with Matchers {
       s"direct path allocated $directLarge bytes, fake self-join allocated $throughJoinLarge bytes")
   }
 
+  test("deduplicateRight renews union branches against the left side") {
+    val conflicting = LocalRelation(AttributeReference("a", IntegerType)())
+    val shared = LocalRelation(AttributeReference("b", IntegerType)())
+    val right = Union(Seq(
+      Join(conflicting, shared, Inner, None, JoinHint.NONE),
+      Join(conflicting, shared, Inner, None, JoinHint.NONE)))
+
+    val result = DeduplicateRelations.deduplicateRight(conflicting, right).asInstanceOf[Union]
+    val relations = result.children.map(_.collect { case r: LocalRelation => r })
+
+    assert(relations.forall(_.head.output != conflicting.output))
+    assert(relations.flatMap(_(1).output).map(_.exprId).distinct.length == shared.output.length)
+  }
+
   test("resolve as with an already existed alias") {
     checkAnalysis(
       Project(Seq(UnresolvedAttribute("tbl2.a")),