Separate max bucket ratio for SMJ and SHJ and add OOM related documentation

Author: c21

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -2656,8 +2656,8 @@ object SQLConf {
       .booleanConf
       .createWithDefault(false)
 
-  val COALESCE_BUCKETS_IN_JOIN_MAX_BUCKET_RATIO =
-    buildConf("spark.sql.bucketing.coalesceBucketsInJoin.maxBucketRatio")
+  val COALESCE_BUCKETS_IN_SORT_MERGE_JOIN_MAX_BUCKET_RATIO =
+    buildConf("spark.sql.bucketing.coalesceBucketsInSortMergeJoin.maxBucketRatio")
       .doc("The ratio of the number of two buckets being coalesced should be less than or " +
         "equal to this value for bucket coalescing to be applied. This configuration only " +
         s"has an effect when '${COALESCE_BUCKETS_IN_JOIN_ENABLED.key}' is set to true.")
@@ -2666,6 +2666,18 @@ object SQLConf {
       .checkValue(_ > 0, "The difference must be positive.")
       .createWithDefault(4)
 
+  val COALESCE_BUCKETS_IN_SHUFFLED_HASH_JOIN_MAX_BUCKET_RATIO =
+    buildConf("spark.sql.bucketing.coalesceBucketsInShuffledHashJoin.maxBucketRatio")
+      .doc("The ratio of the number of two buckets being coalesced should be less than or " +
+        "equal to this value for bucket coalescing to be applied. This configuration only " +
+        s"has an effect when '${COALESCE_BUCKETS_IN_JOIN_ENABLED.key}' is set to true. " +
+        "Note as coalescing reduces parallelism, there might be a higher risk for " +
+        "out of memory error at shuffled hash join build side.")
+      .version("3.1.0")
+      .intConf
+      .checkValue(_ > 0, "The difference must be positive.")
+      .createWithDefault(4)
+
   /**
    * Holds information about keys that have been deprecated.
    *
@@ -3263,8 +3275,11 @@ class SQLConf extends Serializable with Logging {
 
   def coalesceBucketsInJoinEnabled: Boolean = getConf(SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED)
 
-  def coalesceBucketsInJoinMaxBucketRatio: Int =
-    getConf(SQLConf.COALESCE_BUCKETS_IN_JOIN_MAX_BUCKET_RATIO)
+  def coalesceBucketsInSortMergeJoinMaxBucketRatio: Int =
+    getConf(SQLConf.COALESCE_BUCKETS_IN_SORT_MERGE_JOIN_MAX_BUCKET_RATIO)
+
+  def coalesceBucketsInShuffledHashJoinMaxBucketRatio: Int =
+    getConf(SQLConf.COALESCE_BUCKETS_IN_SHUFFLED_HASH_JOIN_MAX_BUCKET_RATIO)
 
   /** ********************** SQLConf functionality methods ************ */
 

sql/core/src/main/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoin.scala

@@ -36,7 +36,8 @@ import org.apache.spark.sql.internal.SQLConf
  *   - The larger bucket number is divisible by the smaller bucket number.
  *   - COALESCE_BUCKETS_IN_JOIN_ENABLED is set to true.
  *   - The ratio of the number of buckets is less than the value set in
- *     COALESCE_BUCKETS_IN_JOIN_MAX_BUCKET_RATIO.
+ *     COALESCE_BUCKETS_IN_SORT_MERGE_JOIN_MAX_BUCKET_RATIO (`SortMergeJoin`) or,
+ *     COALESCE_BUCKETS_IN_SHUFFLED_HASH_JOIN_MAX_BUCKET_RATIO (`ShuffledHashJoin`).
  */
 case class CoalesceBucketsInJoin(conf: SQLConf) extends Rule[SparkPlan] {
   private def updateNumCoalescedBuckets(
@@ -83,17 +84,19 @@ case class CoalesceBucketsInJoin(conf: SQLConf) extends Rule[SparkPlan] {
     }
 
     plan transform {
-      case ExtractJoinWithBuckets(join, numLeftBuckets, numRightBuckets)
-        if math.max(numLeftBuckets, numRightBuckets) / math.min(numLeftBuckets, numRightBuckets) <=
-          conf.coalesceBucketsInJoinMaxBucketRatio =>
+      case ExtractJoinWithBuckets(join, numLeftBuckets, numRightBuckets) =>
+        val bucketRatio = math.max(numLeftBuckets, numRightBuckets) /
+          math.min(numLeftBuckets, numRightBuckets)
         val numCoalescedBuckets = math.min(numLeftBuckets, numRightBuckets)
         join match {
-          case j: SortMergeJoinExec =>
+          case j: SortMergeJoinExec
+            if bucketRatio <= conf.coalesceBucketsInSortMergeJoinMaxBucketRatio =>
             updateNumCoalescedBuckets(j, numLeftBuckets, numRightBuckets, numCoalescedBuckets)
           case j: ShuffledHashJoinExec
             // Only coalesce the buckets for shuffled hash join stream side,
             // to avoid OOM for build side.
-            if isCoalesceSHJStreamSide(j, numLeftBuckets, numRightBuckets, numCoalescedBuckets) =>
+            if bucketRatio <= conf.coalesceBucketsInShuffledHashJoinMaxBucketRatio &&
+              isCoalesceSHJStreamSide(j, numLeftBuckets, numRightBuckets, numCoalescedBuckets) =>
             updateNumCoalescedBuckets(j, numLeftBuckets, numRightBuckets, numCoalescedBuckets)
           case other => other
         }

sql/core/src/test/scala/org/apache/spark/sql/execution/bucketing/CoalesceBucketsInJoinSuite.scala

@@ -167,14 +167,16 @@ class CoalesceBucketsInJoinSuite extends SQLTestUtils with SharedSparkSession {
   }
 
   test("the ratio of the number of buckets is greater than max allowed") {
-    withSQLConf(
-        SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true",
-        SQLConf.COALESCE_BUCKETS_IN_JOIN_MAX_BUCKET_RATIO.key -> "2") {
-      run(JoinSetting(
-        RelationSetting(4, None), RelationSetting(16, None), joinOperator = sortMergeJoin))
-      run(JoinSetting(
-        RelationSetting(4, None), RelationSetting(16, None), joinOperator = shuffledHashJoin,
-        shjBuildSide = Some(BuildLeft)))
+    withSQLConf(SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true") {
+      withSQLConf(SQLConf.COALESCE_BUCKETS_IN_SORT_MERGE_JOIN_MAX_BUCKET_RATIO.key -> "2") {
+        run(JoinSetting(
+          RelationSetting(4, None), RelationSetting(16, None), joinOperator = sortMergeJoin))
+      }
+      withSQLConf(SQLConf.COALESCE_BUCKETS_IN_SHUFFLED_HASH_JOIN_MAX_BUCKET_RATIO.key -> "2") {
+        run(JoinSetting(
+          RelationSetting(4, None), RelationSetting(16, None), joinOperator = shuffledHashJoin,
+          shjBuildSide = Some(BuildLeft)))
+      }
     }
   }
 

sql/core/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala

@@ -922,7 +922,7 @@ abstract class BucketedReadSuite extends QueryTest with SQLTestUtils {
 
     withSQLConf(
       SQLConf.COALESCE_BUCKETS_IN_JOIN_ENABLED.key -> "true",
-      SQLConf.COALESCE_BUCKETS_IN_JOIN_MAX_BUCKET_RATIO.key -> "2") {
+      SQLConf.COALESCE_BUCKETS_IN_SORT_MERGE_JOIN_MAX_BUCKET_RATIO.key -> "2") {
       // Coalescing buckets is not applied because the ratio of the number of buckets (3)
       // is greater than max allowed (2).
       run(