cloud-fan commented on code in PR #59636:
URL: https://github.com/apache/spark/pull/59636#discussion_r4327793513
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala:
##########
@@ -262,6 +262,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 ->
+ "org.apache.spark.sql.catalyst.optimizer.ReorderJoin") {
+ withTable("dim16", "fact100", "other8") {
+ 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
Review Comment:
**Nit (P3):** `picked` is created here, but only the three tables are scoped
by `withTable`, so the view stays registered in the suite's shared session
after this test, whether it passes or fails. `SharedSparkSession` resets the
catalog only in `afterAll`. Nothing else in the suite uses the name today, but
a later test that does would silently pick up this definition. Wrapping the
body in `withTempView("picked")` next to the existing `withTable(...)` drops
the view in a `finally` block, which is how the test already handles its tables.
**Verification:**
- **Behavior:** The end-to-end bucket-coalescing join test still passes: the
matching row is returned and the 100-bucket scan is not coalesced to 8.
- **Inspection:** The temp view the end-to-end test creates is created and
used only inside a temp-view scoping helper, so the test leaves no session
catalog state behind.
--
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]
