dongjoon-hyun commented on code in PR #59636:
URL: https://github.com/apache/spark/pull/59636#discussion_r4464269545
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoin.scala:
##########
@@ -40,27 +40,32 @@ import
org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNes
object CoalesceBucketsInJoin extends Rule[SparkPlan] {
private def updateNumCoalescedBucketsInScan(
plan: SparkPlan,
+ targetScan: FileSourceScanExec,
numCoalescedBuckets: Int): SparkPlan = {
+ // Only the scan whose partitioning reaches this join. A broadcast build
+ // scan in the same subtree keeps its own bucket count.
plan transformUp {
Review Comment:
Minor: this still runs `transformUp` over the whole coalesced side and finds
the target by reference (`f eq targetScan`), instead of rewriting along the
path that `findBucketedScan` walked, so every occurrence of that instance is
stamped. If the same `FileSourceScanExec` instance appeared twice in the
coalesced side, e.g. once on the streamed path and once inside a bucketed
sort-merge join under a broadcast build side, both copies would get
`Some(numCoalescedBuckets)`, which brings back the nested-join corruption this
PR fixes. The planner creates a new scan instance for each logical node, so I
do not think this happens today. It would take a hand-built plan or a later
rule that shares subtrees by reference. Could the rewrite follow the same
`FilterExec`, `ProjectExec` and streamed-child path as `findBucketedScan`,
rebuilding each node with `withNewChildren`, so that it stamps exactly the
walked leaf? If that is too much code for this PR, could the comment at L45-46
at least say t
hat it relies on the scan instance not being shared?
##########
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:
Minor: this scaladoc still describes the old extractor. It does not say that
`unapply` now also returns the bucketed scan of each side, which is the only
scan the rule may stamp and the invariant this PR relies on. Also, "are
consisted of only the scan operation" did not match the walk through filters,
projects and broadcast joins even before this PR. Could L112-117 say something
like this?
```scala
/**
* An extractor that extracts `SortMergeJoinExec` and `ShuffledHashJoin`
whose sides each
* reach a bucketed scan through filters, projects, and the streamed child
of a broadcast
* join, and whose numbers of buckets are not equal but divisible. It
returns the join, the
* bucketed scan of each side, which is the only scan the caller may
coalesce, and the two
* numbers of buckets.
*/
```
##########
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:
Minor, following up on my item 1: this covers a nested bucketed join only
under a `BuildRight` broadcast, while the description now says the corruption
happens "on either build side". Under `BuildLeft`, the base reads the count
from the nested join, because `collectFirst` reaches its left scan first. For
example, in `SMJ(BHJ(BuildLeft, SMJ(a16, c12), s16), o8)` the base reads 16
from a16, accepts 16 vs 8, and stamps a16, c12 and s16 with `Some(8)`. Since 12
% 8 != 0, the inner join then runs without a shuffle and drops rows, while this
PR stamps only s16. I traced this by reading the code and did not run it. Could
this test add that case?
```scala
val innerLeft = SortMergeJoinExec(
Seq(a.head), Seq(c.head), Inner, None, scan(a, 16), scan(c, 12))
val bhjLeft = BroadcastHashJoinExec(
Seq(a(1)), Seq(s(1)), Inner, BuildLeft, None, innerLeft, scan(s, 16))
val planLeft = CoalesceBucketsInJoin(SortMergeJoinExec(
Seq(s.head), Seq(o.head), Inner, None, bhjLeft, scan(o, 8)))
assert(coalesced(planLeft) == Seq((16, None), (12, None), (16, Some(8)), (8,
None)))
```
##########
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:
Nit: every new plan-level case puts a `SortMergeJoinExec` over the broadcast
join, so no case covers a `ShuffledHashJoinExec` parent, where
`isCoalesceSHJStreamSide` (CoalesceBucketsInJoin.scala:103) now gets the
streamed scan's count. For example, over `hashJoinDiv` (a `BuildLeft` join of a
100-bucket build scan and a 16-bucket streamed scan) and an 8-bucket scan, a
shuffled hash join should coalesce the streamed scan to 8 when it builds on the
right, and change nothing when it builds on the left. The base coalesces
nothing in either case, because it reads 100. Could this test add those two
cases?
```scala
val shj = CoalesceBucketsInJoin(ShuffledHashJoinExec(
Seq(streamDiv.head), Seq(otherDiv.head), Inner, BuildRight, None,
hashJoinDiv, scan(otherDiv, 8)))
assert(coalesced(shj) == Seq((100, None), (16, Some(8)), (8, None)))
```
With `BuildLeft`, the expected value is `Seq((100, None), (16, None), (8,
None))`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -17,20 +17,22 @@
package org.apache.spark.sql.execution.bucketing
+import org.apache.spark.sql.Row
import org.apache.spark.sql.catalyst.catalog.BucketSpec
import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference}
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight,
BuildSide}
import org.apache.spark.sql.catalyst.plans.Inner
import org.apache.spark.sql.catalyst.types.DataTypeUtils
import org.apache.spark.sql.execution.{BinaryExecNode, FileSourceScanExec,
SparkPlan}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.datasources.{HadoopFsRelation,
InMemoryFileIndex, PartitionSpec}
import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
-import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec,
ShuffledHashJoinExec, SortMergeJoinExec}
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec,
BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.IntegerType
-class CoalesceBucketsInJoinSuite extends SharedSparkSession {
+class CoalesceBucketsInJoinSuite extends SharedSparkSession with
AdaptiveSparkPlanHelper {
Review Comment:
Nit: since this line changes anyway and the suite now calls `checkAnswer`,
could it mix in `QueryTest` explicitly, with `import
org.apache.spark.sql.{QueryTest, Row}` at L20? `SharedSparkSession` already
extends `QueryTest`, so nothing changes at runtime, but AGENTS.md:61 asks for
the explicit mix-in. The one-line form would be 105 characters, so it needs a
break:
```scala
class CoalesceBucketsInJoinSuite extends QueryTest with SharedSparkSession
with AdaptiveSparkPlanHelper {
```
##########
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:
Nit: this `scan` and the `coalesced` helper at L345-349 are line-for-line
copies of the ones in the previous test (L269-279 and L285-289), so the file
now has three copies of this `FileSourceScanExec` construction, counting
`newFileSourceScanExec` at L70-80. This adds to my item 12, which you left for
a follow-up, but this second copy is new in 081b027. Could the two helpers at
least move to private suite-level methods that both tests share?
##########
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)
+ val plan = CoalesceBucketsInJoin(SortMergeJoinExec(
+ Seq(s.head), Seq(o.head), Inner, None, bhj, scan(o, 8)))
+ assert(coalesced(plan) == Seq((16, Some(8)), (100, None), (8, None), (8,
None)))
+ }
+ }
+
+ test("SPARK-60382: broadcast hash join bucket coalescing keeps the
sort-merge match") {
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString,
+ SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true") {
+ withTable("dim16", "fact100", "other8") {
+ withTempView("picked") {
+ spark.sql(
+ """CREATE TABLE dim16 (i INT, k INT) USING parquet
+ |CLUSTERED BY (i) INTO 16 BUCKETS""".stripMargin)
+ spark.sql(
+ """CREATE TABLE fact100 (i INT, k INT, v STRING) USING parquet
+ |CLUSTERED BY (i) INTO 100 BUCKETS""".stripMargin)
+ spark.sql(
+ """CREATE TABLE other8 (i INT, v STRING) USING parquet
+ |CLUSTERED BY (i) INTO 8 BUCKETS""".stripMargin)
+ spark.sql("INSERT INTO dim16 VALUES (0, 1)")
+ spark.sql(
+ """CREATE OR REPLACE TEMP VIEW picked AS
+ |SELECT CAST(id AS INT) AS i
+ |FROM range(0, 4000)
+ |WHERE pmod(pmod(hash(CAST(id AS INT)), 100), 8)
+ | <> pmod(hash(CAST(id AS INT)), 8)
+ |ORDER BY id
+ |LIMIT 1""".stripMargin)
+ spark.sql("INSERT INTO fact100 SELECT i, 1, 'fact' FROM picked")
+ spark.sql("INSERT INTO other8 SELECT i, 'other' FROM picked")
+ val key = spark.sql("SELECT i FROM picked").head().getInt(0)
+
+ val query =
+ """SELECT /*+ BROADCAST(dim16), MERGE(fact100, other8) */
+ | fact100.i, fact100.v, other8.v
+ |FROM dim16
+ |JOIN fact100 ON dim16.k = fact100.k
+ |JOIN other8 ON fact100.i = other8.i""".stripMargin
+ val df = spark.sql(query)
+ val executed = df.queryExecution.executedPlan
Review Comment:
Minor: with AQE on, `executed` is an `AdaptiveSparkPlanExec` that has not
run yet, so `collect(executed)` walks its initial plan (`executedPlan` returns
`currentPhysicalPlan`, AdaptiveSparkPlanExec.scala:302), and these two asserts
check only the rule's first decision. `CoalesceBucketsInJoin` runs again in
`reOptimize` (AdaptiveSparkPlanExec.scala:1150) after the dim16 broadcast stage
materializes, which is where AQE kept the base code's wrong 0-shuffle plan
(item 2), but nothing here pins the plan that AQE finally runs. `checkAnswer`
at L415 runs this same `df`, since `df.collect()` uses `df.queryExecution`, so
afterwards `executedPlan` holds the final plan. Could `checkAnswer` move above
L408, so that these asserts check the plan that produced the answer? I traced
this by reading the code and did not run it.
--
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]
