|
17 | 17 |
|
18 | 18 | package org.apache.spark.sql.catalyst.analysis |
19 | 19 |
|
| 20 | +import java.lang.management.ManagementFactory |
20 | 21 | import java.util.{TimeZone, UUID} |
21 | 22 |
|
22 | 23 | import scala.jdk.CollectionConverters._ |
@@ -553,6 +554,108 @@ class AnalysisSuite extends AnalysisTest with Matchers { |
553 | 554 | assertAnalysisSuccess(r2) |
554 | 555 | } |
555 | 556 |
|
| 557 | + test("DeduplicateRelations preserves union branch order and overlapping outputs") { |
| 558 | + case class TestLeaf(label: String, override val output: Seq[Attribute]) extends LeafNode |
| 559 | + |
| 560 | + val a = AttributeReference("a", IntegerType)() |
| 561 | + val b = AttributeReference("b", IntegerType)() |
| 562 | + val c = AttributeReference("c", IntegerType)() |
| 563 | + val d = AttributeReference("d", IntegerType)() |
| 564 | + val first = TestLeaf("first", Seq(a, b)) |
| 565 | + val second = TestLeaf("second", Seq(a, c)) |
| 566 | + val third = TestLeaf("third", Seq(c, d)) |
| 567 | + |
| 568 | + val result = DeduplicateRelations(Union(Seq(first, second, third))).asInstanceOf[Union] |
| 569 | + |
| 570 | + assert(result.children.head eq first) |
| 571 | + assert(result.children(1).asInstanceOf[Project].child eq second) |
| 572 | + assert(result.children(2).asInstanceOf[Project].child eq third) |
| 573 | + assert(result.children.tail.forall(_.getTagValue( |
| 574 | + resolver.ResolverTag.PROJECT_FOR_EXPRESSION_ID_DEDUPLICATION).contains(()))) |
| 575 | + assert(result.children.flatMap(_.output).map(_.exprId).distinct.length == 6) |
| 576 | + } |
| 577 | + |
| 578 | + test("DeduplicateRelations preserves streaming union children under tagged projections") { |
| 579 | + case class StreamingLeaf(override val output: Seq[Attribute]) extends LeafNode { |
| 580 | + override def isStreaming: Boolean = true |
| 581 | + } |
| 582 | + |
| 583 | + val sharedOutput = Seq(AttributeReference("a", IntegerType)()) |
| 584 | + val first = StreamingLeaf(sharedOutput) |
| 585 | + val second = StreamingLeaf(sharedOutput) |
| 586 | + |
| 587 | + val result = DeduplicateRelations(Union(Seq(first, second))).asInstanceOf[Union] |
| 588 | + val project = result.children(1).asInstanceOf[Project] |
| 589 | + |
| 590 | + assert(result.children.head eq first) |
| 591 | + assert(project.child eq second) |
| 592 | + assert(project.getTagValue( |
| 593 | + resolver.ResolverTag.PROJECT_FOR_EXPRESSION_ID_DEDUPLICATION).contains(())) |
| 594 | + } |
| 595 | + |
| 596 | + test("DeduplicateRelations union work scales linearly with branch count") { |
| 597 | + case class TestLeaf(override val output: Seq[Attribute]) extends LeafNode |
| 598 | + |
| 599 | + val bean = ManagementFactory.getThreadMXBean.asInstanceOf[com.sun.management.ThreadMXBean] |
| 600 | + if (!bean.isThreadAllocatedMemoryEnabled) { |
| 601 | + bean.setThreadAllocatedMemoryEnabled(true) |
| 602 | + } |
| 603 | + val threadId = Thread.currentThread().getId |
| 604 | + val sharedOutput = (0 until 26).map(i => AttributeReference(s"c$i", IntegerType)()) |
| 605 | + |
| 606 | + def allocatedBytes(branchCount: Int): Long = { |
| 607 | + val union = Union(Seq.fill(branchCount)(TestLeaf(sharedOutput))) |
| 608 | + val before = bean.getThreadAllocatedBytes(threadId) |
| 609 | + DeduplicateRelations(union) |
| 610 | + bean.getThreadAllocatedBytes(threadId) - before |
| 611 | + } |
| 612 | + |
| 613 | + allocatedBytes(10) |
| 614 | + val small = Seq.fill(3)(allocatedBytes(100)).min |
| 615 | + val large = Seq.fill(3)(allocatedBytes(500)).min |
| 616 | + |
| 617 | + assert(large <= small * 7, |
| 618 | + s"100 branches allocated $small bytes, while 500 branches allocated $large bytes") |
| 619 | + } |
| 620 | + |
| 621 | + test("deduplicateRight matches fake self-join semantics with less scaling overhead") { |
| 622 | + def wideProject(width: Int): LogicalPlan = { |
| 623 | + val relation = LocalRelation(AttributeReference("a", IntegerType)()) |
| 624 | + Project((0 until width).map(i => Alias(relation.output.head, s"c$i")()), relation) |
| 625 | + } |
| 626 | + |
| 627 | + val semanticPlan = wideProject(10) |
| 628 | + val throughJoin = DeduplicateRelations( |
| 629 | + Join(semanticPlan, semanticPlan, Inner, None, JoinHint.NONE)).children(1) |
| 630 | + val direct = DeduplicateRelations.deduplicateRight(semanticPlan, semanticPlan) |
| 631 | + comparePlans(direct, throughJoin, checkAnalysis = false) |
| 632 | + |
| 633 | + val bean = ManagementFactory.getThreadMXBean.asInstanceOf[com.sun.management.ThreadMXBean] |
| 634 | + if (!bean.isThreadAllocatedMemoryEnabled) { |
| 635 | + bean.setThreadAllocatedMemoryEnabled(true) |
| 636 | + } |
| 637 | + val threadId = Thread.currentThread().getId |
| 638 | + |
| 639 | + def allocatedBytes(width: Int, useDirectPath: Boolean): Long = { |
| 640 | + val plan = wideProject(width) |
| 641 | + val before = bean.getThreadAllocatedBytes(threadId) |
| 642 | + if (useDirectPath) { |
| 643 | + DeduplicateRelations.deduplicateRight(plan, plan) |
| 644 | + } else { |
| 645 | + DeduplicateRelations(Join(plan, plan, Inner, None, JoinHint.NONE)).children(1) |
| 646 | + } |
| 647 | + bean.getThreadAllocatedBytes(threadId) - before |
| 648 | + } |
| 649 | + |
| 650 | + allocatedBytes(10, useDirectPath = true) |
| 651 | + allocatedBytes(10, useDirectPath = false) |
| 652 | + val directLarge = Seq.fill(3)(allocatedBytes(500, useDirectPath = true)).min |
| 653 | + val throughJoinLarge = Seq.fill(3)(allocatedBytes(500, useDirectPath = false)).min |
| 654 | + |
| 655 | + assert(directLarge * 3 < throughJoinLarge * 2, |
| 656 | + s"direct path allocated $directLarge bytes, fake self-join allocated $throughJoinLarge bytes") |
| 657 | + } |
| 658 | + |
556 | 659 | test("resolve as with an already existed alias") { |
557 | 660 | checkAnalysis( |
558 | 661 | Project(Seq(UnresolvedAttribute("tbl2.a")), |
|
0 commit comments