diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DeduplicateRelations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DeduplicateRelations.scala index ce29706501f28..4a06fab691df1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DeduplicateRelations.scala +++ b/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) => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala index 1b36ba04cfd79..e989489a730e6 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala +++ b/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 diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala index 5454c6d88b001..c4bc51f020b62 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala +++ b/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) } } 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 7949d6f3f6b06..3daa5ba6c37d1 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 @@ -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) } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala index 4a7adcc050a00..a07163c8b90a2 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala +++ b/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")),