wangyum commented on code in PR #59636:
URL: https://github.com/apache/spark/pull/59636#discussion_r4331967810
##########
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:
fixed
--
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]
