perf: keep the sort-merge join when the forced hash join's build side is too large

benchmarks/pyspark/run_all_benchmarks.sh

@@ -67,6 +67,7 @@ $SPARK_HOME/bin/spark-submit \
 # Run Comet JVM shuffle
 echo ""
 echo ">>> Running COMET JVM shuffle benchmark..."
+# maxBuildSize=-1 measures the hash join alone: no build-side size limit on the rewrite.
 $SPARK_HOME/bin/spark-submit \
   --master "$SPARK_MASTER" \
   --executor-memory "$EXECUTOR_MEMORY" \
@@ -86,6 +87,7 @@ $SPARK_HOME/bin/spark-submit \
   --conf spark.comet.shuffle.mode=jvm \
   --conf spark.comet.shuffle.mode=jvm \
   --conf spark.comet.exec.replaceSortMergeJoin=true \
+  --conf spark.comet.exec.forceShuffledHashJoin.maxBuildSize=-1 \
   --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \
   --conf spark.sql.extensions=org.apache.comet.CometSparkSessionExtensions \
   --conf spark.comet.cast.allowIncompatible=true \
@@ -96,6 +98,7 @@ $SPARK_HOME/bin/spark-submit \
 # Run Comet Native shuffle
 echo ""
 echo ">>> Running COMET NATIVE shuffle benchmark..."
+# maxBuildSize=-1 measures the hash join alone: no build-side size limit on the rewrite.
 $SPARK_HOME/bin/spark-submit \
   --master "$SPARK_MASTER" \
   --executor-memory "$EXECUTOR_MEMORY" \
@@ -114,6 +117,7 @@ $SPARK_HOME/bin/spark-submit \
   --conf spark.comet.explain.fallback.enabled=true \
   --conf spark.comet.shuffle.mode=native \
   --conf spark.comet.exec.replaceSortMergeJoin=true \
+  --conf spark.comet.exec.forceShuffledHashJoin.maxBuildSize=-1 \
   --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \
   --conf spark.sql.extensions=org.apache.comet.CometSparkSessionExtensions \
   --conf spark.comet.cast.allowIncompatible=true \

benchmarks/tpc/engines/comet-hashjoin.toml

@@ -31,4 +31,6 @@ driver_class_path = ["$COMET_JAR"]
 "spark.plugins" = "org.apache.spark.CometPlugin"
 "spark.shuffle.manager" = "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager"
 "spark.comet.exec.replaceSortMergeJoin" = "true"
+# Measure the hash join alone: no build-side size limit on the rewrite.
+"spark.comet.exec.forceShuffledHashJoin.maxBuildSize" = "-1"
 "spark.comet.expression.Cast.allowIncompatible" = "true"

benchmarks/tpc/engines/comet-iceberg-hashjoin.toml

@@ -34,6 +34,8 @@ driver_class_path = ["$COMET_JAR", "$ICEBERG_JAR"]
 "spark.plugins" = "org.apache.spark.CometPlugin"
 "spark.shuffle.manager" = "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager"
 "spark.comet.exec.replaceSortMergeJoin" = "true"
+# Measure the hash join alone: no build-side size limit on the rewrite.
+"spark.comet.exec.forceShuffledHashJoin.maxBuildSize" = "-1"
 "spark.comet.expression.Cast.allowIncompatible" = "true"
 "spark.comet.enabled" = "true"
 "spark.comet.exec.enabled" = "true"

docs/source/contributor-guide/benchmarking_macos.md

@@ -154,6 +154,7 @@ $SPARK_HOME/bin/spark-submit \
     --conf spark.comet.enabled=true \
     --conf spark.comet.shuffle.enabled=true \
     --conf spark.comet.exec.forceShuffledHashJoin=true \
+    --conf spark.comet.exec.forceShuffledHashJoin.maxBuildSize=-1 \
     $DF_BENCH/runners/datafusion-comet/tpcbench.py \
     --benchmark tpch \
     --data $BENCH_DATA/ \

docs/source/contributor-guide/memory_management.md

@@ -631,4 +631,4 @@ A checklist for triaging an executor OOM kill:
    `batch_size * columns`, and wide or deeply nested schemas amplify it.
 4. Check whether the operators involved can spill at all. `ShuffledHashJoin` cannot, so
    `spark.comet.exec.forceShuffledHashJoin=true` converts a spillable sort-merge join into one that
-   is not.
+   is not, for build sides under `spark.comet.exec.forceShuffledHashJoin.maxBuildSize`.

docs/source/user-guide/latest/tuning/operators.md

@@ -30,6 +30,14 @@ to configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`. Comet does
 to test with both for your specific workloads.
 
 To configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`, set `spark.comet.exec.forceShuffledHashJoin=true`.
+The conversion only happens when the build side is under a size limit, and a join whose build side has no statistics
+is left as a `SortMergeJoin`. The size is Spark's planning estimate of the build side, or under AQE the materialized
+shuffle size of the build side. By default the limit is Spark's own rule for choosing a `ShuffledHashJoin`:
+`spark.sql.autoBroadcastJoinThreshold` times the initial shuffle partition count
+(`spark.sql.adaptive.coalescePartitions.initialPartitionNum` when AQE and partition coalescing are both on and it is
+set, else `spark.sql.shuffle.partitions`). When broadcasts are disabled with a non-positive threshold, Spark's default
+threshold of 10 MB is used instead. Set `spark.comet.exec.forceShuffledHashJoin.maxBuildSize` to a size in bytes to use a fixed
+limit, or to a non-positive value to convert every eligible join regardless of size.
 
 ### Join Runtime Filters
 

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -559,6 +559,22 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_FORCE_SHJ_MAX_BUILD_SIZE: OptionalConfigEntry[Long] =
+    conf(s"$COMET_EXEC_CONFIG_PREFIX.forceShuffledHashJoin.maxBuildSize")
+      .category(CATEGORY_EXEC)
+      .doc(s"The build side size below which `${COMET_FORCE_SHJ.key}` converts a " +
+        "SortMergeJoin to ShuffledHashJoin. The size is Spark's planning estimate of the build " +
+        "child, or under AQE the materialized shuffle size of the build side. A build side at " +
+        "or over this size, or one with no statistics, keeps the SortMergeJoin. When unset, " +
+        "Spark's own rule applies: `spark.sql.autoBroadcastJoinThreshold` times the initial " +
+        "shuffle partition count (`spark.sql.adaptive.coalescePartitions.initialPartitionNum` " +
+        "when AQE and partition coalescing are both on and it is set, else " +
+        "`spark.sql.shuffle.partitions`). When broadcasts are disabled with a non-positive " +
+        "threshold, Spark's default threshold of 10 MB is used instead. A " +
+        s"non-positive value removes the limit. $TUNING_GUIDE.")
+      .bytesConf(ByteUnit.BYTE)
+      .createOptional
+
   val COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: ConfigEntry[Boolean] =
     conf(s"$COMET_EXEC_CONFIG_PREFIX.aggregate.skipPartial.enabled")
       .category(CATEGORY_EXEC)

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -994,10 +994,13 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
     } else {
       val normalizedPlan = normalizeDivisors(plan)
 
+      // The rewrite removes the sorts of each join it converts, so put back any sort a kept
+      // operator above still needs. This runs before transform(), which converts the sorts it
+      // adds like any other SortExec.
       val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) {
-        normalizedPlan.transformUp { case p =>
-          RewriteJoin.rewrite(p)
-        }
+        RewriteJoin.restoreRequiredOrdering(normalizedPlan.transformUp { case p =>
+          RewriteJoin.rewrite(p, conf)
+        })
       } else {
         normalizedPlan
       }

spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala

@@ -19,13 +19,16 @@
 
 package org.apache.comet.rules
 
+import org.apache.spark.sql.catalyst.expressions.SortOrder
 import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, JoinSelectionHelper}
 import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, LeftSemi}
 import org.apache.spark.sql.catalyst.plans.logical.Join
 import org.apache.spark.sql.execution.{SortExec, SparkPlan}
 import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec}
+import org.apache.spark.sql.internal.SQLConf
 
-import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
+import org.apache.comet.CometConf
+import org.apache.comet.CometSparkSessionExtensions.{withFallbackReason, withInfo}
 
 /**
  * Adapted from equivalent rule in Apache Gluten.
@@ -64,7 +67,69 @@ object RewriteJoin extends JoinSelectionHelper {
     case _ => plan
   }
 
-  def rewrite(plan: SparkPlan): SparkPlan = plan match {
+  /**
+   * Adds a local sort under any operator whose child no longer has the ordering the operator
+   * requires, using the same check as EnsureRequirements. EnsureRequirements puts no sort above a
+   * sort-merge join whose output ordering already satisfies its parent, so when rewrite turns
+   * that join into a hash join and removes its sorts, a kept parent such as another sort-merge
+   * join on the same key would read unsorted input. Plans that already satisfy their required
+   * orderings are returned unchanged.
+   */
+  def restoreRequiredOrdering(plan: SparkPlan): SparkPlan = plan.transformUp { case p =>
+    p.withNewChildren(p.children.zip(p.requiredChildOrdering).map {
+      case (child, required) if !SortOrder.orderingSatisfies(child.outputOrdering, required) =>
+        SortExec(required, global = false, child = child)
+      case (child, _) => child
+    })
+  }
+
+  /**
+   * The largest build side the rewrite converts, or None when the limit is switched off. Unset
+   * means Spark's own rule for choosing a shuffled hash join by size: the broadcast threshold
+   * times the initial shuffle partition count. A non-positive threshold only disables broadcast
+   * joins and says nothing about the hash table an executor can hold, so Spark's default
+   * threshold stands in for it.
+   */
+  private def maxBuildSize(conf: SQLConf): Option[BigInt] =
+    CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.get(conf) match {
+      case Some(limit) if limit <= 0 => None
+      case Some(limit) => Some(BigInt(limit))
+      case None =>
+        val threshold = conf.autoBroadcastJoinThreshold
+        val perPartition =
+          if (threshold > 0) threshold else SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get
+        Some(BigInt(perPartition) * conf.numShufflePartitions)
+    }
+
+  /**
+   * Why the build side is too large to rewrite, or None when it fits. Like Spark's
+   * canBuildLocalHashMapBySize, the size must be strictly under the limit. The hash join cannot
+   * spill, so a build side without statistics is refused as well. Spark reports an unknown size
+   * as Long.MaxValue, which fails the comparison on its own.
+   */
+  private def buildSideOverLimit(
+      smj: SortMergeJoinExec,
+      buildSide: BuildSide,
+      conf: SQLConf): Option[String] =
+    maxBuildSize(conf).flatMap { limit =>
+      smj.logicalLink match {
+        case Some(join: Join) =>
+          val buildSize = buildSide match {
+            case BuildLeft => join.left.stats.sizeInBytes
+            case BuildRight => join.right.stats.sizeInBytes
+          }
+          if (buildSize >= limit) {
+            Some(
+              s"build side size estimate of $buildSize bytes is not under the limit of " +
+                s"$limit bytes")
+          } else {
+            None
+          }
+        case _ => Some("no statistics are available for the build side")
+      }
+    }
+
+  def rewrite(plan: SparkPlan, conf: SQLConf): SparkPlan = plan match {
     case smj: SortMergeJoinExec =>
       getSmjBuildSide(smj) match {
         case Some(BuildRight)
@@ -77,15 +142,23 @@ object RewriteJoin extends JoinSelectionHelper {
               s"BuildRight with ${smj.joinType} is not supported")
           plan
         case Some(buildSide) =>
-          ShuffledHashJoinExec(
-            smj.leftKeys,
-            smj.rightKeys,
-            smj.joinType,
-            buildSide,
-            smj.condition,
-            removeSort(smj.left),
-            removeSort(smj.right),
-            smj.isSkewJoin)
+          buildSideOverLimit(smj, buildSide, conf) match {
+            case Some(reason) =>
+              // The sort-merge join still runs natively, so this is information for extended
+              // explain rather than a fallback reason.
+              withInfo(smj, s"Cannot rewrite SortMergeJoin to HashJoin: $reason")
+              plan
+            case None =>
+              ShuffledHashJoinExec(
+                smj.leftKeys,
+                smj.rightKeys,
+                smj.joinType,
+                buildSide,
+                smj.condition,
+                removeSort(smj.left),
+                removeSort(smj.right),
+                smj.isSkewJoin)
+          }
         case _ => plan
       }
     case _ => plan