dongjoon-hyun commented on code in PR #59636:
URL: https://github.com/apache/spark/pull/59636#discussion_r4371926421
##########
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 {
- case f: FileSourceScanExec if f.relation.bucketSpec.nonEmpty =>
+ case f: FileSourceScanExec if f eq targetScan =>
Review Comment:
**This narrowing also fixes a second wrong-result shape that needs no
`BuildLeft`, but the description and the JIRA scope the bug to `BuildLeft`
children, and no test covers it.** The base stamped every bucketed scan of the
coalesced side (`case f: FileSourceScanExec if
f.relation.bucketSpec.nonEmpty`), including scans inside a broadcast build
side. For example, take `SMJ(s.i = o.i)(BHJ(BuildRight, s16, SMJ(a.i =
c.i)(a100, c8)), o8)`:
- `collectFirst` correctly finds the streamed s16, so the left side is
coalesced from 16 to 8.
- The base `transformUp` also stamps a100 and c8 inside the build side with
`Some(8)`.
- When the top-down pass reaches the inner join, the old `getBucketSpec`
skips the stamped scans, so that join makes no decision of its own.
- a100 then reports `HashPartitioning(i, 8)` while it groups its files by
`bucketId % 8`, and 100 % 8 != 0. EnsureRequirements adds no shuffle for the
inner join, so the broadcast relation loses the rows whose keys land in
different partitions.
With this PR only s16 is stamped, and the inner join is judged on its own
(100 vs 8 is not divisible). So a bucketed join under a `BuildRight` broadcast
is affected too, which matters for the JIRA and the release note. I traced this
by reading the code and did not run it. Could the description and the JIRA also
cover bucketed scans inside a broadcast build side, on either build side, and
could the plan-level test add a case like this?
```scala
val s = cols()
val a = cols()
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 p = CoalesceBucketsInJoin(SortMergeJoinExec(
Seq(s.head), Seq(o.head), Inner, None, bhj, scan(o, 8)))
assert(coalesced(p) == Seq((16, Some(8)), (100, None), (8, None), (8, None)))
```
The base returns `(100, Some(8))` and `(8, Some(8))` for the inner pair.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +263,111 @@ 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(
Review Comment:
Minor: this is the only plan-level case that rewrites a scan, and it uses a
broadcast hash join, while the PR description says the divisible 16-to-8 case
is checked "for both broadcast hash and broadcast nested-loop joins". The
nested-loop join at L302-306 covers only the 100-vs-8 case, where nothing is
rewritten. A nested-loop rewrite through `BuildRight` is still pinned by
BucketedReadSuite.scala:1067-1073, so this is a coverage gap rather than an
untested branch. Could this add the nested-loop variant, e.g.
`BroadcastNestedLoopJoinExec(scan(buildDiv, 100), scan(streamDiv, 16),
BuildLeft, Inner, None)` under the same sort-merge join with the expected
`Seq((100, None), (16, Some(8)), (8, None))`, or could the description sentence
be narrowed?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +263,111 @@ 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)))
+ }
+ }
+
+ test("SPARK-60382: broadcast hash join bucket coalescing keeps the
sort-merge match") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
Review Comment:
**This pins AQE off, so the default AQE path has no end-to-end coverage,
although the base bug reproduces there too.** `CoalesceBucketsInJoin` is also
the first rule of `queryStagePreparationRules`
(AdaptiveSparkPlanExec.scala:145). With AQE on, the base code stamps fact100
with 8 in the initial plan, so EnsureRequirements adds no shuffle under the
sort-merge join. After the dim16 broadcast stage materializes, the re-planned
join has a `BroadcastQueryStageExec` leaf on its build side, so the base
`collectFirst` finds 100 and the new plan needs a shuffle. AQE keeps the
current plan unless the new one is cheaper (AdaptiveSparkPlanExec.scala:451),
so the 0-shuffle wrong plan runs and the query returns no rows. I traced this
by reading the code and did not run it. Could this test loop over `Seq(true,
false)` for AQE as BucketedReadSuite.scala:1031 does, and mix in
`AdaptiveSparkPlanHelper` so that `collect(executed)` sees the scans inside
`AdaptiveSparkPlanExec`?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +263,111 @@ 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(
Review Comment:
Minor: this test does not carry over to branch-3.5 as it is, although the
bug has been there since 3.5.0 (SPARK-43107). On branch-3.5:
- `FileSourceScanExec` has no `stream` parameter, so
`FileSourceScanExec(relation, None, cols, ...)` does not compile.
- `satisfiesOutputPartitioning` has no `PartitioningCollection` case,
because SPARK-43228 is only in 4.0.0 and later, while an inner
`BroadcastHashJoinExec` always reports a `PartitioningCollection` there
(BroadcastHashJoinExec.scala:76 and :115 on branch-3.5). With the arity fixed,
`planDiv` at L319 fails even with the fix, and the first case at L300 passes
even without it.
- The nested-loop case and the end-to-end test still tell the base and the
fix apart, because a `ProjectExec` unwraps a one-element collection there
(AliasAwareOutputExpression.scala:58 on branch-3.5).
I checked this by reading branch-3.5 and did not run it. The rule file is
unchanged on branch-4.0 and later, so only a 3.5 backport is affected. Could
the broadcast hash join cases put a `ProjectExec` over the broadcast join, or
set `spark.sql.execution.broadcastHashJoin.outputPartitioningExpandLimit` to 0,
so that the same test works on branch-3.5?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoin.scala:
##########
@@ -111,26 +116,34 @@ object CoalesceBucketsInJoin extends Rule[SparkPlan] {
* and numbers of buckets are not equal but divisible.
*/
object ExtractJoinWithBuckets {
- @tailrec
- private def hasScanOperation(plan: SparkPlan): Boolean = plan match {
- case f: FilterExec => hasScanOperation(f.child)
- case p: ProjectExec => hasScanOperation(p.child)
- case j: BroadcastHashJoinExec =>
- if (j.buildSide == BuildLeft) hasScanOperation(j.right) else
hasScanOperation(j.left)
- case j: BroadcastNestedLoopJoinExec =>
- if (j.buildSide == BuildLeft) hasScanOperation(j.right) else
hasScanOperation(j.left)
- case f: FileSourceScanExec => f.relation.bucketSpec.nonEmpty
- case _ => false
+ private def streamedSide(
Review Comment:
Nit: the base already used a one-line form for these two cases (`if
(j.buildSide == BuildLeft) hasScanOperation(j.right) else
hasScanOperation(j.left)`). The helper does not let the two cases merge, since
no common trait exposes `buildSide`, and both branches of the inline form are
tail calls, so `@tailrec` still holds. Could L135 and L137 be `if (j.buildSide
== BuildLeft) findBucketedScan(j.right) else findBucketedScan(j.left)`,
dropping `streamedSide` and the `BuildSide` import at L24?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +263,111 @@ 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 = {
Review Comment:
Nit: this copies `newFileSourceScanExec` (L69-79) line for line except for
the bucket columns (`Seq("i")` instead of all columns). That constructor call
has already needed two edits in the helper (SPARK-44475 and SPARK-50144), and
the next one would have to be made in both places. Could
`newFileSourceScanExec` take an optional list of bucket columns, or could these
cases use single-column `RelationSetting(n, None)` scans with the broadcast
join keyed on `i`? In the latter form the base code still fails all three
assertions, because `collectFirst` reaches the build scan first whatever the
join keys are.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoin.scala:
##########
@@ -85,18 +90,18 @@ object CoalesceBucketsInJoin extends Rule[SparkPlan] {
}
plan transform {
- case ExtractJoinWithBuckets(join, numLeftBuckets, numRightBuckets)
+ case ExtractJoinWithBuckets(join, leftScan, rightScan, numLeftBuckets,
numRightBuckets)
if math.max(numLeftBuckets, numRightBuckets) /
math.min(numLeftBuckets, numRightBuckets) <=
Review Comment:
Minor: this ratio check now uses the streamed scan's count, so some plans
that the base coalesced correctly now get a shuffle, and the "user-facing
change" section does not say so. With the default `maxBucketRatio` of 4, take
`SMJ(BHJ(BuildLeft, dim16, fact64), other8)` with every scan bucketed on the
join keys. The base took 16 from the build scan (ratio 2) and stamped fact64
with 8, which was valid because 64 % 8 == 0, so the join ran without a shuffle.
Now the ratio is 64 / 8 = 8, the rule declines, and EnsureRequirements shuffles
other8. That follows the documented meaning of
`spark.sql.bucketing.coalesceBucketsInJoin.maxBucketRatio`, so the new behavior
looks right to me. Likewise, a bucketed build scan is no longer stamped, so its
broadcast read uses its own bucket count instead of the coalesced one. I traced
this by reading the code and did not run it. Could the description list these
next to the newly enabled 16-to-8 case?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +263,111 @@ 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)))
+ }
+ }
+
+ test("SPARK-60382: broadcast hash join bucket coalescing keeps the
sort-merge match") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.CBO_ENABLED.key -> "false",
+ SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true",
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ReorderJoin.ruleName) {
+ 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 executed = spark.sql(query).queryExecution.executedPlan
Review Comment:
Nit: the plan check builds its own `Dataset` here, and `checkAnswer` at L365
builds and plans another one, so the asserted plan is not the one that runs.
They are the same today with AQE off. Could this be `val df =
spark.sql(query)`, with `df.queryExecution.executedPlan` here and
`checkAnswer(df, Seq(Row(key, "fact", "other")))` at L365?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +263,111 @@ 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)))
+ }
+ }
+
+ test("SPARK-60382: broadcast hash join bucket coalescing keeps the
sort-merge match") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.CBO_ENABLED.key -> "false",
+ SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true",
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ReorderJoin.ruleName) {
+ 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 executed = spark.sql(query).queryExecution.executedPlan
+ val scans = executed.collect { case f: FileSourceScanExec =>
+ (f.relation.bucketSpec.map(_.numBuckets),
f.optionalNumCoalescedBuckets)
+ }
+ assert(scans == Seq((Some(16), None), (Some(100), None), (Some(8),
None)), executed)
Review Comment:
Minor: these two checks would also pass on the base code if the inner join
stopped being a `BuildLeft` broadcast hash join. If it were planned as a
sort-merge join, or with dim16 as the streamed side, the base rule would not
coalesce anything either: all three scans would stay `(Some(n), None)` in the
same order, and the answer would be right. That is unlikely today, since
`BROADCAST` takes precedence over `MERGE` on the same join, and the plan-level
test above still guards the rule. Could this also assert the shape that the
test relies on, e.g. `assert(executed.collectFirst { case b:
BroadcastHashJoinExec => b.buildSide }.contains(BuildLeft), executed)`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoin.scala:
##########
@@ -111,26 +116,34 @@ object CoalesceBucketsInJoin extends Rule[SparkPlan] {
* and numbers of buckets are not equal but divisible.
*/
object ExtractJoinWithBuckets {
- @tailrec
- private def hasScanOperation(plan: SparkPlan): Boolean = plan match {
- case f: FilterExec => hasScanOperation(f.child)
- case p: ProjectExec => hasScanOperation(p.child)
- case j: BroadcastHashJoinExec =>
- if (j.buildSide == BuildLeft) hasScanOperation(j.right) else
hasScanOperation(j.left)
- case j: BroadcastNestedLoopJoinExec =>
- if (j.buildSide == BuildLeft) hasScanOperation(j.right) else
hasScanOperation(j.left)
- case f: FileSourceScanExec => f.relation.bucketSpec.nonEmpty
- case _ => false
+ private def streamedSide(
+ buildSide: BuildSide, left: SparkPlan, right: SparkPlan): SparkPlan = {
+ if (buildSide == BuildLeft) right else left
}
- private def getBucketSpec(plan: SparkPlan): Option[BucketSpec] = {
- plan.collectFirst {
- case f: FileSourceScanExec if f.relation.bucketSpec.nonEmpty &&
- f.optionalNumCoalescedBuckets.isEmpty =>
- f.relation.bucketSpec.get
+ /**
+ * Bucketed scan whose partitioning this side exposes. Broadcast joins
Review Comment:
Minor: this doc says more than `findBucketedScan` checks. L138 returns any
scan with `relation.bucketSpec.nonEmpty`, including one whose `bucketedScan` is
false (bucketing disabled or a bucket column pruned) and one behind a
nested-loop join type that reports `UnknownPartitioning(0)`, so the returned
scan does not always expose its partitioning to this side. The real guard is
`satisfiesOutputPartitioning` at L177-178. Could this say something like "The
only bucketed scan whose partitioning can reach the parent join through this
side; the caller still checks the side's `outputPartitioning`."?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoin.scala:
##########
@@ -111,26 +116,34 @@ object CoalesceBucketsInJoin extends Rule[SparkPlan] {
* and numbers of buckets are not equal but divisible.
*/
object ExtractJoinWithBuckets {
- @tailrec
- private def hasScanOperation(plan: SparkPlan): Boolean = plan match {
- case f: FilterExec => hasScanOperation(f.child)
- case p: ProjectExec => hasScanOperation(p.child)
- case j: BroadcastHashJoinExec =>
- if (j.buildSide == BuildLeft) hasScanOperation(j.right) else
hasScanOperation(j.left)
- case j: BroadcastNestedLoopJoinExec =>
- if (j.buildSide == BuildLeft) hasScanOperation(j.right) else
hasScanOperation(j.left)
- case f: FileSourceScanExec => f.relation.bucketSpec.nonEmpty
- case _ => false
+ private def streamedSide(
+ buildSide: BuildSide, left: SparkPlan, right: SparkPlan): SparkPlan = {
+ if (buildSide == BuildLeft) right else left
}
- private def getBucketSpec(plan: SparkPlan): Option[BucketSpec] = {
- plan.collectFirst {
- case f: FileSourceScanExec if f.relation.bucketSpec.nonEmpty &&
- f.optionalNumCoalescedBuckets.isEmpty =>
- f.relation.bucketSpec.get
+ /**
+ * Bucketed scan whose partitioning this side exposes. Broadcast joins
+ * expose the streamed child, so a bucketed build scan is not the count
+ * the parent sort-merge or shuffled-hash join coalesces against.
+ */
+ @tailrec
+ private def findBucketedScan(plan: SparkPlan): Option[FileSourceScanExec] = {
+ plan match {
+ case f: FilterExec => findBucketedScan(f.child)
+ case p: ProjectExec => findBucketedScan(p.child)
+ case j: BroadcastHashJoinExec =>
+ findBucketedScan(streamedSide(j.buildSide, j.left, j.right))
+ case j: BroadcastNestedLoopJoinExec =>
+ findBucketedScan(streamedSide(j.buildSide, j.left, j.right))
+ case f: FileSourceScanExec if f.relation.bucketSpec.nonEmpty => Some(f)
+ case _ => None
}
}
+ private def getBucketSpec(scan: FileSourceScanExec): Option[BucketSpec] = {
Review Comment:
Nit: `findBucketedScan` already guarantees `relation.bucketSpec.nonEmpty`
(L138), so this helper only rejects a scan that is already coalesced. Could
that check move into the walk's leaf case, `case f: FileSourceScanExec if
f.relation.bucketSpec.nonEmpty && f.optionalNumCoalescedBuckets.isEmpty =>
Some(f)`, with L179-180 reading `leftScan.relation.bucketSpec` and
`rightScan.relation.bucketSpec`? A coalesced leaf then falls to `case _ =>
None`, so the extractor still returns `None`, and the `BucketSpec` import at
L22 goes away.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +263,111 @@ 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)))
+ }
+ }
+
+ test("SPARK-60382: broadcast hash join bucket coalescing keeps the
sort-merge match") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.CBO_ENABLED.key -> "false",
+ SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true",
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ReorderJoin.ruleName) {
+ 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)
Review Comment:
Nit: this view re-writes Spark's bucket-id function in SQL and is evaluated
three times, by the two `INSERT ... SELECT` statements and by the read-back at
L352. `BucketingUtils.getBucketIdFromValue` already computes the writer's
bucket id, and BucketedReadSuite.scala:140 uses it that way. SQL `hash` and the
writer's hash agree for an INT key, so this is about clarity and cost, not
correctness today. Could the key be computed once, e.g. `val iAttr =
AttributeReference("i", IntegerType)()` and `val key = (0 until 4000).find(v =>
BucketingUtils.getBucketIdFromValue(iAttr, 100, v) % 8 !=
BucketingUtils.getBucketIdFromValue(iAttr, 8, v)).get`, and inserted with
`VALUES`? That also removes the view and its `withTempView`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +263,111 @@ 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)))
+ }
+ }
+
+ test("SPARK-60382: broadcast hash join bucket coalescing keeps the
sort-merge match") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.CBO_ENABLED.key -> "false",
+ SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true",
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ReorderJoin.ruleName) {
Review Comment:
Nit: three of these five confs have no effect on this query, and this line
also replaces the suite-wide
`spark.sql.optimizer.excludedRules=ConvertToLocalRelation`
(SparkSessionBinder.scala:90) for the test, because `withSQLConf` overwrites
the value instead of appending to it.
- `ReorderJoin` does not touch hinted joins: `ExtractFiltersAndInnerJoins`
matches only a `Join` whose `hint == JoinHint.NONE` (patterns.scala:270 and
:272), and both joins here carry hints after `EliminateResolvedHint`.
- `spark.sql.cbo.enabled` (L327) is false by default.
- `spark.sql.autoBroadcastJoinThreshold` (L326) is not consulted, because
`JoinSelection` takes both strategies from the hints.
Re-enabling `ConvertToLocalRelation` changes nothing here, but the extra
pins suggest that each one is needed for the repro. Could L326, L327 and L329
be dropped, together with the `ReorderJoin` import at L23?
--
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]
