[SPARK-59452][SQL] Optimize DeduplicateRelations for wide plans

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

@@ -29,6 +29,13 @@ import org.apache.spark.sql.internal.SQLConf
 object DeduplicateRelations extends Rule[LogicalPlan] {
   type ExprIdMap = mutable.HashMap[Class[_], mutable.HashSet[Long]]
 
+  /** Renews `right` against expression IDs collected from `left`. */
+  private[sql] def deduplicateRight(left: LogicalPlan, right: LogicalPlan): LogicalPlan = {
+    val existingRelations = mutable.HashMap.empty[Class[_], mutable.HashSet[Long]]
+    renewDuplicatedRelations(existingRelations, left)
+    renewDuplicatedRelations(existingRelations, right)._1
+  }
+
   override def apply(plan: LogicalPlan): LogicalPlan = {
     val newPlan = renewDuplicatedRelations(mutable.HashMap.empty, plan)._1
 
@@ -67,24 +74,25 @@ object DeduplicateRelations extends Rule[LogicalPlan] {
           DeduplicateUnionChildOutput.deduplicateOutputPerChild(u)
         // Use projection-based de-duplication for Union to avoid breaking the checkpoint sharing
         // feature in streaming.
-        val newChildren =
-          unionWithChildOutputsDeduplicated.children.foldRight(Seq.empty[LogicalPlan]) {
-            (head, tail) =>
-              head +: tail.map {
-                case child if head.outputSet.intersect(child.outputSet).isEmpty =>
-                  child
-                case child =>
-                  val projectList = child.output.map { attr =>
-                    Alias(attr, attr.name)()
-                  }
-                  val project = Project(projectList, child)
-                  project.setTagValue(
-                    ResolverTag.PROJECT_FOR_EXPRESSION_ID_DEDUPLICATION,
-                    ()
-                  )
-                  project
-              }
+        val seenExprIds = mutable.HashSet.empty[Long]
+        val newChildren = unionWithChildOutputsDeduplicated.children.map { child =>
+          val childOutput = child.output
+          val hasConflictingExprId = childOutput.exists(attr => seenExprIds(attr.exprId.id))
+          childOutput.foreach(attr => seenExprIds += attr.exprId.id)
+          if (hasConflictingExprId) {
+            val projectList = childOutput.map { attr =>
+              Alias(attr, attr.name)()
+            }
+            val project = Project(projectList, child)
+            project.setTagValue(
+              ResolverTag.PROJECT_FOR_EXPRESSION_ID_DEDUPLICATION,
+              ()
+            )
+            project
+          } else {
+            child
           }
+        }
         unionWithChildOutputsDeduplicated.copy(children = newChildren)
       case merge: MergeIntoTable
           if !merge.duplicateResolved && noMissingInput(merge.sourceTable) =>

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

@@ -22,8 +22,7 @@ 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, OuterReference, OuterScopeReference, SubqueryExpression}
-import org.apache.spark.sql.catalyst.plans.Inner
-import org.apache.spark.sql.catalyst.plans.logical.{CTERelationDef, CTERelationRef, Join, JoinHint, LogicalPlan, Project, Subquery, UnionLoop, WithCTE}
+import org.apache.spark.sql.catalyst.plans.logical.{CTERelationDef, CTERelationRef, LogicalPlan, Project, Subquery, UnionLoop, WithCTE}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreePattern.{CTE, PLAN_EXPRESSION}
 
@@ -278,15 +277,7 @@ case class InlineCTE(
             if (ref.outputSet == refInfo.cteDef.outputSet) {
               cteBody
             } else {
-              val ctePlan = DeduplicateRelations(
-                Join(
-                  cteBody,
-                  cteBody,
-                  Inner,
-                  None,
-                  JoinHint(None, None)
-                )
-              ).children(1)
+              val ctePlan = DeduplicateRelations.deduplicateRight(cteBody, cteBody)
               val projectList = ref.output.zip(ctePlan.output).map { case (tgtAttr, srcAttr) =>
                 if (srcAttr.semanticEquals(tgtAttr)) {
                   tgtAttr

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

@@ -17,7 +17,6 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
-import org.apache.spark.SparkException
 import org.apache.spark.sql.catalyst.analysis.DeduplicateRelations
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans._
@@ -141,17 +140,9 @@ case class PushDownJoinThroughUnion(override val conf: SQLConf)
   }
 
   /**
-   * Creates a copy of `plan` with fresh ExprIds on all output attributes,
-   * using the same "fake self-join + DeduplicateRelations" pattern as InlineCTE.
+   * Creates a copy of `plan` with fresh ExprIds on all output attributes.
    */
   private def dedupRight(plan: LogicalPlan): LogicalPlan = {
-    DeduplicateRelations(
-      Join(plan, plan, Inner, None, JoinHint.NONE)
-    ) match {
-      case Join(_, deduped, _, _, _) => deduped
-      case other =>
-        throw SparkException.internalError(
-          s"Unexpected plan shape after DeduplicateRelations: ${other.getClass.getName}")
-    }
+    DeduplicateRelations.deduplicateRight(plan, plan)
   }
 }

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

@@ -22,7 +22,6 @@ 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
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreePattern.{CTE, PLAN_EXPRESSION}
@@ -74,8 +73,7 @@ object ReplaceCTERefWithRepartition extends Rule[LogicalPlan] {
       if (ref.outputSet == cteDefPlan.outputSet) {
         cteDefPlan
       } else {
-        val ctePlan = DeduplicateRelations(
-          Join(cteDefPlan, cteDefPlan, Inner, None, JoinHint(None, None))).children(1)
+        val ctePlan = DeduplicateRelations.deduplicateRight(cteDefPlan, cteDefPlan)
         val projectList = ref.output.zip(ctePlan.output).map { case (tgtAttr, srcAttr) =>
           Alias(srcAttr, tgtAttr.name)(exprId = tgtAttr.exprId)
         }

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

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
+import java.lang.management.ManagementFactory
 import java.util.{TimeZone, UUID}
 
 import scala.jdk.CollectionConverters._
@@ -553,6 +554,108 @@ class AnalysisSuite extends AnalysisTest with Matchers {
     assertAnalysisSuccess(r2)
   }
 
+  test("DeduplicateRelations preserves union branch order and overlapping outputs") {
+    case class TestLeaf(label: String, override val output: Seq[Attribute]) extends LeafNode
+
+    val a = AttributeReference("a", IntegerType)()
+    val b = AttributeReference("b", IntegerType)()
+    val c = AttributeReference("c", IntegerType)()
+    val d = AttributeReference("d", IntegerType)()
+    val first = TestLeaf("first", Seq(a, b))
+    val second = TestLeaf("second", Seq(a, c))
+    val third = TestLeaf("third", Seq(c, d))
+
+    val result = DeduplicateRelations(Union(Seq(first, second, third))).asInstanceOf[Union]
+
+    assert(result.children.head eq first)
+    assert(result.children(1).asInstanceOf[Project].child eq second)
+    assert(result.children(2).asInstanceOf[Project].child eq third)
+    assert(result.children.tail.forall(_.getTagValue(
+      resolver.ResolverTag.PROJECT_FOR_EXPRESSION_ID_DEDUPLICATION).contains(())))
+    assert(result.children.flatMap(_.output).map(_.exprId).distinct.length == 6)
+  }
+
+  test("DeduplicateRelations preserves streaming union children under tagged projections") {
+    case class StreamingLeaf(override val output: Seq[Attribute]) extends LeafNode {
+      override def isStreaming: Boolean = true
+    }
+
+    val sharedOutput = Seq(AttributeReference("a", IntegerType)())
+    val first = StreamingLeaf(sharedOutput)
+    val second = StreamingLeaf(sharedOutput)
+
+    val result = DeduplicateRelations(Union(Seq(first, second))).asInstanceOf[Union]
+    val project = result.children(1).asInstanceOf[Project]
+
+    assert(result.children.head eq first)
+    assert(project.child eq second)
+    assert(project.getTagValue(
+      resolver.ResolverTag.PROJECT_FOR_EXPRESSION_ID_DEDUPLICATION).contains(()))
+  }
+
+  test("DeduplicateRelations union work scales linearly with branch count") {
+    case class TestLeaf(override val output: Seq[Attribute]) extends LeafNode
+
+    val bean = ManagementFactory.getThreadMXBean.asInstanceOf[com.sun.management.ThreadMXBean]
+    if (!bean.isThreadAllocatedMemoryEnabled) {
+      bean.setThreadAllocatedMemoryEnabled(true)
+    }
+    val threadId = Thread.currentThread().getId
+    val sharedOutput = (0 until 26).map(i => AttributeReference(s"c$i", IntegerType)())
+
+    def allocatedBytes(branchCount: Int): Long = {
+      val union = Union(Seq.fill(branchCount)(TestLeaf(sharedOutput)))
+      val before = bean.getThreadAllocatedBytes(threadId)
+      DeduplicateRelations(union)
+      bean.getThreadAllocatedBytes(threadId) - before
+    }
+
+    allocatedBytes(10)
+    val small = Seq.fill(3)(allocatedBytes(100)).min
+    val large = Seq.fill(3)(allocatedBytes(500)).min
+
+    assert(large <= small * 7,
+      s"100 branches allocated $small bytes, while 500 branches allocated $large bytes")
+  }
+
+  test("deduplicateRight matches fake self-join semantics with less scaling overhead") {
+    def wideProject(width: Int): LogicalPlan = {
+      val relation = LocalRelation(AttributeReference("a", IntegerType)())
+      Project((0 until width).map(i => Alias(relation.output.head, s"c$i")()), relation)
+    }
+
+    val semanticPlan = wideProject(10)
+    val throughJoin = DeduplicateRelations(
+      Join(semanticPlan, semanticPlan, Inner, None, JoinHint.NONE)).children(1)
+    val direct = DeduplicateRelations.deduplicateRight(semanticPlan, semanticPlan)
+    comparePlans(direct, throughJoin, checkAnalysis = false)
+
+    val bean = ManagementFactory.getThreadMXBean.asInstanceOf[com.sun.management.ThreadMXBean]
+    if (!bean.isThreadAllocatedMemoryEnabled) {
+      bean.setThreadAllocatedMemoryEnabled(true)
+    }
+    val threadId = Thread.currentThread().getId
+
+    def allocatedBytes(width: Int, useDirectPath: Boolean): Long = {
+      val plan = wideProject(width)
+      val before = bean.getThreadAllocatedBytes(threadId)
+      if (useDirectPath) {
+        DeduplicateRelations.deduplicateRight(plan, plan)
+      } else {
+        DeduplicateRelations(Join(plan, plan, Inner, None, JoinHint.NONE)).children(1)
+      }
+      bean.getThreadAllocatedBytes(threadId) - before
+    }
+
+    allocatedBytes(10, useDirectPath = true)
+    allocatedBytes(10, useDirectPath = false)
+    val directLarge = Seq.fill(3)(allocatedBytes(500, useDirectPath = true)).min
+    val throughJoinLarge = Seq.fill(3)(allocatedBytes(500, useDirectPath = false)).min
+
+    assert(directLarge * 3 < throughJoinLarge * 2,
+      s"direct path allocated $directLarge bytes, fake self-join allocated $throughJoinLarge bytes")
+  }
+
   test("resolve as with an already existed alias") {
     checkAnalysis(
       Project(Seq(UnresolvedAttribute("tbl2.a")),