wangyum commented on code in PR #59636:
URL: https://github.com/apache/spark/pull/59636#discussion_r4481621468
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +264,161 @@ class CoalesceBucketsInJoinSuite extends
SharedSparkSession {
}
}
+ test("SPARK-60382: broadcast join coalescing uses the streamed side bucket
count") {
+ withSQLConf(SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true") {
+ def scan(cols: Seq[Attribute], numBuckets: Int): FileSourceScanExec = {
+ val relation = HadoopFsRelation(
+ location = new InMemoryFileIndex(spark, Nil, Map.empty, None),
+ partitionSchema = PartitionSpec.emptySpec.partitionColumns,
+ dataSchema = DataTypeUtils.fromAttributes(cols),
+ bucketSpec = Some(BucketSpec(numBuckets, Seq("i"), Nil)),
+ fileFormat = new ParquetFileFormat(),
+ options = Map.empty)(spark)
+ FileSourceScanExec(
+ relation, None, cols, relation.dataSchema, Nil, None, None, Nil,
None)
+ }
+
+ def cols(): Seq[Attribute] = Seq(
+ AttributeReference("i", IntegerType)(),
+ AttributeReference("k", IntegerType)())
+
+ def coalesced(plan: SparkPlan): Seq[(Int, Option[Int])] = {
+ plan.collect { case f: FileSourceScanExec =>
+ (f.relation.bucketSpec.get.numBuckets, f.optionalNumCoalescedBuckets)
+ }
+ }
+
+ // Streamed side is 100 buckets. 100 % 8 != 0, so nothing is coalesced,
+ // including the 16-bucket build scan that sits on the left.
+ val build = cols()
+ val stream = cols()
+ val other = Seq(AttributeReference("i", IntegerType)())
+ val hashJoin = BroadcastHashJoinExec(
+ Seq(build(1)), Seq(stream(1)), Inner, BuildLeft, None,
+ scan(build, 16), scan(stream, 100))
+ val plan = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(stream.head), Seq(other.head), Inner, None, hashJoin, scan(other,
8)))
+ assert(coalesced(plan) == Seq((16, None), (100, None), (8, None)))
+
+ val nestedLoop = BroadcastNestedLoopJoinExec(
+ scan(build, 16), scan(stream, 100), BuildLeft, Inner, None)
+ val nestedPlan = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(stream.head), Seq(other.head), Inner, None, nestedLoop,
scan(other, 8)))
+ assert(coalesced(nestedPlan) == Seq((16, None), (100, None), (8, None)))
+
+ // Streamed side is 16 buckets and the other input is 8. Only the
+ // streamed scan is rewritten to 8. The 100-bucket build scan is not.
+ val buildDiv = cols()
+ val streamDiv = cols()
+ val otherDiv = Seq(AttributeReference("i", IntegerType)())
+ val hashJoinDiv = BroadcastHashJoinExec(
+ Seq(buildDiv(1)), Seq(streamDiv(1)), Inner, BuildLeft, None,
+ scan(buildDiv, 100), scan(streamDiv, 16))
+ val planDiv = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(streamDiv.head), Seq(otherDiv.head), Inner, None,
+ hashJoinDiv, scan(otherDiv, 8)))
+ assert(coalesced(planDiv) == Seq((100, None), (16, Some(8)), (8, None)))
+
+ val nestedLoopDiv = BroadcastNestedLoopJoinExec(
+ scan(buildDiv, 100), scan(streamDiv, 16), BuildLeft, Inner, None)
+ val nestedPlanDiv = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(streamDiv.head), Seq(otherDiv.head), Inner, None,
+ nestedLoopDiv, scan(otherDiv, 8)))
+ assert(coalesced(nestedPlanDiv) == Seq((100, None), (16, Some(8)), (8,
None)))
+ }
+ }
+
+ test("SPARK-60382: a bucketed join inside a broadcast build side keeps its
own bucket count") {
+ withSQLConf(SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true") {
+ def scan(cols: Seq[Attribute], numBuckets: Int): FileSourceScanExec = {
+ val relation = HadoopFsRelation(
+ location = new InMemoryFileIndex(spark, Nil, Map.empty, None),
+ partitionSchema = PartitionSpec.emptySpec.partitionColumns,
+ dataSchema = DataTypeUtils.fromAttributes(cols),
+ bucketSpec = Some(BucketSpec(numBuckets, Seq("i"), Nil)),
+ fileFormat = new ParquetFileFormat(),
+ options = Map.empty)(spark)
+ FileSourceScanExec(
+ relation, None, cols, relation.dataSchema, Nil, None, None, Nil,
None)
+ }
+
+ def coalesced(plan: SparkPlan): Seq[(Int, Option[Int])] = {
+ plan.collect { case f: FileSourceScanExec =>
+ (f.relation.bucketSpec.get.numBuckets, f.optionalNumCoalescedBuckets)
+ }
+ }
+
+ // SMJ(s.i = o.i)(BHJ(BuildRight, s16, SMJ(a.i = c.i)(a100, c8)), o8).
The outer join
+ // coalesces s16 (the streamed side of the inner broadcast join) to 8.
The inner
+ // sort-merge join is then judged on its own: 100 is not divisible by 8,
so a100 and
+ // c8 are left untouched even though they sit inside the outer join's
coalesced side.
+ val s = Seq(
+ AttributeReference("i", IntegerType)(), AttributeReference("k",
IntegerType)())
+ val a = Seq(
+ AttributeReference("i", IntegerType)(), AttributeReference("k",
IntegerType)())
+ val c = Seq(AttributeReference("i", IntegerType)())
+ val o = Seq(AttributeReference("i", IntegerType)())
+
+ val inner = SortMergeJoinExec(
+ Seq(a.head), Seq(c.head), Inner, None, scan(a, 100), scan(c, 8))
+ val bhj = BroadcastHashJoinExec(
+ Seq(s(1)), Seq(a(1)), Inner, BuildRight, None, scan(s, 16), inner)
Review Comment:
Fixed in dd986166df8: added the BuildLeft variant
(innerLeft/bhjLeft/planLeft) to this test, asserting that only the streamed
scan s16 is coalesced to 8 while the nested a16/c12 sort-merge join on the
build side is left untouched.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +264,161 @@ class CoalesceBucketsInJoinSuite extends
SharedSparkSession {
}
}
+ test("SPARK-60382: broadcast join coalescing uses the streamed side bucket
count") {
+ withSQLConf(SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true") {
+ def scan(cols: Seq[Attribute], numBuckets: Int): FileSourceScanExec = {
+ val relation = HadoopFsRelation(
+ location = new InMemoryFileIndex(spark, Nil, Map.empty, None),
+ partitionSchema = PartitionSpec.emptySpec.partitionColumns,
+ dataSchema = DataTypeUtils.fromAttributes(cols),
+ bucketSpec = Some(BucketSpec(numBuckets, Seq("i"), Nil)),
+ fileFormat = new ParquetFileFormat(),
+ options = Map.empty)(spark)
+ FileSourceScanExec(
+ relation, None, cols, relation.dataSchema, Nil, None, None, Nil,
None)
+ }
+
+ def cols(): Seq[Attribute] = Seq(
+ AttributeReference("i", IntegerType)(),
+ AttributeReference("k", IntegerType)())
+
+ def coalesced(plan: SparkPlan): Seq[(Int, Option[Int])] = {
+ plan.collect { case f: FileSourceScanExec =>
+ (f.relation.bucketSpec.get.numBuckets, f.optionalNumCoalescedBuckets)
+ }
+ }
+
+ // Streamed side is 100 buckets. 100 % 8 != 0, so nothing is coalesced,
+ // including the 16-bucket build scan that sits on the left.
+ val build = cols()
+ val stream = cols()
+ val other = Seq(AttributeReference("i", IntegerType)())
+ val hashJoin = BroadcastHashJoinExec(
+ Seq(build(1)), Seq(stream(1)), Inner, BuildLeft, None,
+ scan(build, 16), scan(stream, 100))
+ val plan = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(stream.head), Seq(other.head), Inner, None, hashJoin, scan(other,
8)))
+ assert(coalesced(plan) == Seq((16, None), (100, None), (8, None)))
+
+ val nestedLoop = BroadcastNestedLoopJoinExec(
+ scan(build, 16), scan(stream, 100), BuildLeft, Inner, None)
+ val nestedPlan = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(stream.head), Seq(other.head), Inner, None, nestedLoop,
scan(other, 8)))
+ assert(coalesced(nestedPlan) == Seq((16, None), (100, None), (8, None)))
+
+ // Streamed side is 16 buckets and the other input is 8. Only the
+ // streamed scan is rewritten to 8. The 100-bucket build scan is not.
+ val buildDiv = cols()
+ val streamDiv = cols()
+ val otherDiv = Seq(AttributeReference("i", IntegerType)())
+ val hashJoinDiv = BroadcastHashJoinExec(
+ Seq(buildDiv(1)), Seq(streamDiv(1)), Inner, BuildLeft, None,
+ scan(buildDiv, 100), scan(streamDiv, 16))
+ val planDiv = CoalesceBucketsInJoin(SortMergeJoinExec(
Review Comment:
Fixed in dd986166df8: added ShuffledHashJoinExec parent cases
(shjRight/shjLeft) over the broadcast join, covering both the coalescing
outcome (BuildRight) and the no-op outcome (BuildLeft).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoin.scala:
##########
@@ -111,26 +116,35 @@ object CoalesceBucketsInJoin extends Rule[SparkPlan] {
* and numbers of buckets are not equal but divisible.
Review Comment:
Fixed in dd986166df8: reworded the scaladoc to describe the walk through
filters, projects, and the streamed child of a broadcast join, and to say that
unapply returns the bucketed scan of each side, which is the only one the rule
may stamp and whose partitioning satisfiesOutputPartitioning still checks.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +264,161 @@ class CoalesceBucketsInJoinSuite extends
SharedSparkSession {
}
}
+ test("SPARK-60382: broadcast join coalescing uses the streamed side bucket
count") {
+ withSQLConf(SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true") {
+ def scan(cols: Seq[Attribute], numBuckets: Int): FileSourceScanExec = {
+ val relation = HadoopFsRelation(
+ location = new InMemoryFileIndex(spark, Nil, Map.empty, None),
+ partitionSchema = PartitionSpec.emptySpec.partitionColumns,
+ dataSchema = DataTypeUtils.fromAttributes(cols),
+ bucketSpec = Some(BucketSpec(numBuckets, Seq("i"), Nil)),
+ fileFormat = new ParquetFileFormat(),
+ options = Map.empty)(spark)
+ FileSourceScanExec(
+ relation, None, cols, relation.dataSchema, Nil, None, None, Nil,
None)
+ }
+
+ def cols(): Seq[Attribute] = Seq(
+ AttributeReference("i", IntegerType)(),
+ AttributeReference("k", IntegerType)())
+
+ def coalesced(plan: SparkPlan): Seq[(Int, Option[Int])] = {
+ plan.collect { case f: FileSourceScanExec =>
+ (f.relation.bucketSpec.get.numBuckets, f.optionalNumCoalescedBuckets)
+ }
+ }
+
+ // Streamed side is 100 buckets. 100 % 8 != 0, so nothing is coalesced,
+ // including the 16-bucket build scan that sits on the left.
+ val build = cols()
+ val stream = cols()
+ val other = Seq(AttributeReference("i", IntegerType)())
+ val hashJoin = BroadcastHashJoinExec(
+ Seq(build(1)), Seq(stream(1)), Inner, BuildLeft, None,
+ scan(build, 16), scan(stream, 100))
+ val plan = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(stream.head), Seq(other.head), Inner, None, hashJoin, scan(other,
8)))
+ assert(coalesced(plan) == Seq((16, None), (100, None), (8, None)))
+
+ val nestedLoop = BroadcastNestedLoopJoinExec(
+ scan(build, 16), scan(stream, 100), BuildLeft, Inner, None)
+ val nestedPlan = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(stream.head), Seq(other.head), Inner, None, nestedLoop,
scan(other, 8)))
+ assert(coalesced(nestedPlan) == Seq((16, None), (100, None), (8, None)))
+
+ // Streamed side is 16 buckets and the other input is 8. Only the
+ // streamed scan is rewritten to 8. The 100-bucket build scan is not.
+ val buildDiv = cols()
+ val streamDiv = cols()
+ val otherDiv = Seq(AttributeReference("i", IntegerType)())
+ val hashJoinDiv = BroadcastHashJoinExec(
+ Seq(buildDiv(1)), Seq(streamDiv(1)), Inner, BuildLeft, None,
+ scan(buildDiv, 100), scan(streamDiv, 16))
+ val planDiv = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(streamDiv.head), Seq(otherDiv.head), Inner, None,
+ hashJoinDiv, scan(otherDiv, 8)))
+ assert(coalesced(planDiv) == Seq((100, None), (16, Some(8)), (8, None)))
+
+ val nestedLoopDiv = BroadcastNestedLoopJoinExec(
+ scan(buildDiv, 100), scan(streamDiv, 16), BuildLeft, Inner, None)
+ val nestedPlanDiv = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(streamDiv.head), Seq(otherDiv.head), Inner, None,
+ nestedLoopDiv, scan(otherDiv, 8)))
+ assert(coalesced(nestedPlanDiv) == Seq((100, None), (16, Some(8)), (8,
None)))
+ }
+ }
+
+ test("SPARK-60382: a bucketed join inside a broadcast build side keeps its
own bucket count") {
+ withSQLConf(SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true") {
+ def scan(cols: Seq[Attribute], numBuckets: Int): FileSourceScanExec = {
Review Comment:
Fixed in dd986166df8: deduped the scan/coalesced helpers from both new tests
into shared suite-level testScan/coalescedBuckets helpers. While doing this I
noticed testScan itself duplicated newFileSourceScanExec's relation-building
code and hardcoded the bucket column to "i" instead of deriving it, so I also
had it delegate to newFileSourceScanExec with an explicit bucketCols override.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
