wangyum commented on code in PR #59636:
URL: https://github.com/apache/spark/pull/59636#discussion_r4378932314


##########
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:
   Fixed in 081b0274e1b: added an assertion that the `BroadcastHashJoinExec` in 
the executed plan has `buildSide == BuildLeft`.
   



##########
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:
   Fixed in 081b0274e1b: dropped `AUTO_BROADCASTJOIN_THRESHOLD`, `CBO_ENABLED`, 
and the `ReorderJoin` exclusion (and its import); confirmed each is a no-op 
given the explicit join hints.
   



##########
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:
   Documented in the PR description's user-facing-change section.
   



##########
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:
   Fixed in 081b0274e1b: reworded the doc comment to point at 
`satisfiesOutputPartitioning` as the actual guard.
   



##########
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:
   Left as-is for now to avoid further churn on the already-reviewed rule file; 
happy to inline this if you'd still like it.
   



##########
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:
   Left as-is for now, same reasoning as the `streamedSide` nit above.
   



##########
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:
   Left as-is for now; agreed this duplicates `newFileSourceScanExec`, can 
follow up if you'd like it deduped.
   



##########
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:
   Left as-is for now; this is correct today, can follow up with the 
`BucketingUtils` version if you'd like.
   



##########
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:
   Left as-is for now; no behavior difference today since the query is 
deterministic, can follow up if you'd like it unified.
   



-- 
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]

Reply via email to