You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/08 14:29:23 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6384-7d294535e87367d819c4e2a23422cd1158ac31b8 created (now 4b160dc09a)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6384-7d294535e87367d819c4e2a23422cd1158ac31b8
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 4b160dc09a perf: keep the sort-merge join when the forced hash join's build side is too large (#6384)
This branch includes the following new commits:
new 4b160dc09a perf: keep the sort-merge join when the forced hash join's build side is too large (#6384)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: perf: keep the sort-merge join when the forced hash join's build side is too large (#6384)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6384-7d294535e87367d819c4e2a23422cd1158ac31b8
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 4b160dc09a22418d29f8a49d37a816894fad69f9
Author: dustin <dw...@users.noreply.github.com>
AuthorDate: Thu Oct 8 14:29:10 2026 +0000
perf: keep the sort-merge join when the forced hash join's build side is too large (#6384)
* perf: keep the sort-merge join when the forced hash join's build side is too large
With spark.comet.exec.forceShuffledHashJoin on, RewriteJoin turned every
eligible SortMergeJoinExec into a shuffled hash join and read the build side's
size estimate only to pick the build side. DataFusion's hash join has to hold
that side in memory, so one large join failed the task while the rest of the
query would have gained from the rewrite.
The rule now rewrites only when the build side's estimate is under a limit and
otherwise keeps the SortMergeJoinExec, which still runs natively, recording
the reason as plan info. The limit is Spark's own rule for choosing a shuffled
hash join, spark.sql.autoBroadcastJoinThreshold times the initial shuffle
partition count, with Spark's 10 MB default threshold standing in when
broadcasts are disabled. spark.comet.exec.forceShuffledHashJoin.maxBuildSize
sets a fixed limit, and a non-positive value removes it, which is what the
hash join benchmark configurations and the macOS benchmarking guide now pin so
they keep measuring every join.
* fix: restore the sort a kept sort-merge join needs after a hash join rewrite
When RewriteJoin keeps a sort-merge join above a join it has already
turned into a hash join on the same key, the kept join reads unsorted
input. EnsureRequirements placed no sort between the two because the
inner sort-merge join's ordering satisfied the outer one, and the
rewrite then removed the inner join's sorts. The native sort-merge join
does not check its input ordering, so matches were silently dropped.
After the rewrite, add a local SortExec wherever a child no longer
satisfies its parent's required ordering, the same check
EnsureRequirements makes. This covers joins kept by the build-size
limit and the LeftSemi and ExistenceJoin joins the rewrite already
keeps, which lost rows the same way.
* test: cover a window over a join the hash join rewrite reordered
A window partitioned by the join key gets no sort of its own, because the
sort-merge join's output ordering already satisfies it. Once RewriteJoin
turns that join into a hash join, the window needs the ordering restored,
or the native window fails its partition-ordering assertion. The test runs
with AQE on and off and checks the join became a hash join.
---
benchmarks/pyspark/run_all_benchmarks.sh | 4 +
benchmarks/tpc/engines/comet-hashjoin.toml | 2 +
benchmarks/tpc/engines/comet-iceberg-hashjoin.toml | 2 +
.../source/contributor-guide/benchmarking_macos.md | 1 +
docs/source/contributor-guide/memory_management.md | 2 +-
docs/source/user-guide/latest/tuning/operators.md | 8 +
.../main/scala/org/apache/comet/CometConf.scala | 16 +
.../org/apache/comet/rules/CometExecRule.scala | 9 +-
.../scala/org/apache/comet/rules/RewriteJoin.scala | 95 +++++-
.../org/apache/comet/exec/CometJoinSuite.scala | 339 ++++++++++++++++++++-
10 files changed, 458 insertions(+), 20 deletions(-)
diff --git a/benchmarks/pyspark/run_all_benchmarks.sh b/benchmarks/pyspark/run_all_benchmarks.sh
index 452ea8d1ff..4e43641874 100755
--- a/benchmarks/pyspark/run_all_benchmarks.sh
+++ b/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 \
diff --git a/benchmarks/tpc/engines/comet-hashjoin.toml b/benchmarks/tpc/engines/comet-hashjoin.toml
index 202dcad914..50e184b049 100644
--- a/benchmarks/tpc/engines/comet-hashjoin.toml
+++ b/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"
diff --git a/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml b/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml
index 421d320131..4a56361310 100644
--- a/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml
+++ b/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"
diff --git a/docs/source/contributor-guide/benchmarking_macos.md b/docs/source/contributor-guide/benchmarking_macos.md
index 13ed2c26a2..4f3541811c 100644
--- a/docs/source/contributor-guide/benchmarking_macos.md
+++ b/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/ \
diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index 7bd16d880f..e3283fe26c 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/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`.
diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md
index 644f35483b..97baabfcca 100644
--- a/docs/source/user-guide/latest/tuning/operators.md
+++ b/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
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 248772fb6c..a187f8fa5f 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/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)
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index fdf37b4e2b..a47f1c2a6b 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/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
}
diff --git a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala b/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala
index 4ffd34451b..5f543cc84a 100644
--- a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala
+++ b/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
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
index eec16205a1..369c251763 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala
@@ -31,16 +31,20 @@ import org.apache.spark.sql.catalyst.TableIdentifier
import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull}
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
-import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec}
-import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan}
-import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec}
+import org.apache.spark.sql.catalyst.plans.Inner
+import org.apache.spark.sql.catalyst.plans.logical.Join
+import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec}
+import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec}
import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException
import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike}
+import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType}
-import org.apache.comet.CometConf
-import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus
+import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo}
+import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus}
+import org.apache.comet.rules.RewriteJoin
class CometJoinSuite extends CometTestBase {
@@ -890,6 +894,331 @@ class CometJoinSuite extends CometTestBase {
}
}
+ // The MERGE hint makes Spark plan a SortMergeJoin, which forceShuffledHashJoin then rewrites.
+ private val forcedSmjQuery =
+ "SELECT /*+ MERGE(tbl_a) */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1"
+
+ // tbl_b is the smaller side, so the rewrite builds on the right.
+ private def withForcedSmjTables(f: => Unit): Unit = {
+ withParquetTable((0 until 1000).map(i => (i, i % 5)), "tbl_a") {
+ withParquetTable((0 until 10).map(i => (i % 10, i + 2)), "tbl_b") {
+ f
+ }
+ }
+ }
+
+ private def rewriteReason(reason: String): String =
+ s"Cannot rewrite SortMergeJoin to HashJoin: $reason"
+
+ // Checks that the SortMergeJoin and its sorts survived and that the kept join, which still
+ // runs natively, is not reported as a fallback. Returns Spark's size estimate of the build
+ // side, read through the same logical link the rewrite used.
+ private def assertSortMergeJoinKept(plan: SparkPlan): BigInt = {
+ val joins = collect(plan) { case j: CometSortMergeJoinExec => j }
+ assert(joins.size == 1, plan)
+ assert(collect(plan) { case s: CometSortExec => s }.size == 2, plan)
+ assert(collect(plan) { case j: CometHashJoinExec => j }.isEmpty, plan)
+ val kept = joins.head
+ assert(!hasFallbackReason(kept) && !hasFallbackReason(kept.originalPlan), plan)
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(plan)
+ assert(!reasons.exists(_.contains("SortMergeJoin")), reasons)
+ kept.originalPlan.logicalLink match {
+ case Some(join: Join) => join.right.stats.sizeInBytes
+ case other => fail(s"Expected a Join logical link, got $other")
+ }
+ }
+
+ // The refusal is recorded as information on the kept join and rendered as COMET-INFO.
+ private def assertKeptWithInfo(plan: SparkPlan, expected: String): Unit = {
+ val join = collect(plan) { case j: CometSortMergeJoinExec => j }.head
+ assert(join.getTagValue(CometExplainInfo.EXTENSION_INFO).exists(_.contains(expected)), plan)
+ val verbose = new ExtendedExplainInfo().generateVerboseInfo(plan)
+ assert(verbose.contains("[COMET-INFO:") && verbose.contains(expected), verbose)
+ }
+
+ private def assertRewrittenToHashJoin(plan: SparkPlan): Unit = {
+ val joins = collect(plan) { case j: CometHashJoinExec => j }
+ assert(joins.map(_.buildSide) == Seq(BuildRight), plan)
+ assert(collect(plan) { case s: CometSortExec => s }.isEmpty, plan)
+ assert(collect(plan) { case j: CometSortMergeJoinExec => j }.isEmpty, plan)
+ }
+
+ private def assertKeptWithReason(plan: SparkPlan, limit: Long): Unit = {
+ val estimate = assertSortMergeJoinKept(plan)
+ assert(estimate >= limit, estimate)
+ assertKeptWithInfo(
+ plan,
+ rewriteReason(
+ s"build side size estimate of $estimate bytes is not under the limit of $limit bytes"))
+ }
+
+ for (adaptive <- Seq(false, true)) {
+ test(s"forceShuffledHashJoin keeps SortMergeJoin over maxBuildSize, AQE=$adaptive") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") {
+ withForcedSmjTables {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ assertKeptWithReason(cometPlan, 1)
+ }
+ }
+ }
+
+ test(s"forceShuffledHashJoin rewrites SortMergeJoin under maxBuildSize, AQE=$adaptive") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1g") {
+ withForcedSmjTables {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ assertRewrittenToHashJoin(cometPlan)
+ }
+ }
+ }
+ }
+
+ test("forceShuffledHashJoin matches Spark's strict size rule at the limit") {
+ // AQE off, so the estimate read from the first plan is the one the later runs compare.
+ withForcedSmjTables {
+ var estimate = BigInt(0)
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ estimate = assertSortMergeJoinKept(cometPlan)
+ }
+ // A build side exactly at the limit is not under it, as in canBuildLocalHashMapBySize.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> estimate.toString) {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ assertKeptWithReason(cometPlan, estimate.toLong)
+ }
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> (estimate + 1).toString) {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ assertRewrittenToHashJoin(cometPlan)
+ }
+ }
+ }
+
+ test("forceShuffledHashJoin with a non-positive maxBuildSize rewrites regardless of size") {
+ withSQLConf(
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "0",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") {
+ withForcedSmjTables {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ assertRewrittenToHashJoin(cometPlan)
+ }
+ }
+ }
+
+ test("forceShuffledHashJoin without maxBuildSize uses Spark's shuffled hash join size rule") {
+ // The limit is autoBroadcastJoinThreshold times the shuffle partition count.
+ withSQLConf(
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") {
+ withForcedSmjTables {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ assertKeptWithReason(cometPlan, 2)
+ }
+ }
+ withSQLConf(
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1g") {
+ withForcedSmjTables {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ assertRewrittenToHashJoin(cometPlan)
+ }
+ }
+ }
+
+ test("forceShuffledHashJoin uses Spark's default threshold when broadcasts are disabled") {
+ // Disabling broadcasts says nothing about the hash table an executor can hold, so the limit
+ // falls back to Spark's default threshold times the partition count. AQE is off so the
+ // decision uses the planning estimate, which fileCompressionFactor scales past 10 MB without
+ // writing a 10 MB table.
+ val defaultThreshold = SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get
+ assert(defaultThreshold == 10L * 1024 * 1024)
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ withForcedSmjTables {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ assertRewrittenToHashJoin(cometPlan)
+ }
+ withSQLConf(SQLConf.FILE_COMPRESSION_FACTOR.key -> "1000000") {
+ withForcedSmjTables {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery))
+ assertKeptWithReason(cometPlan, defaultThreshold)
+ }
+ }
+ }
+ }
+
+ test("forceShuffledHashJoin keeps a SortMergeJoin that has no statistics") {
+ val left = spark.range(1).toDF("a").queryExecution.sparkPlan
+ val right = spark.range(1).toDF("b").queryExecution.sparkPlan
+ def smj: SortMergeJoinExec =
+ SortMergeJoinExec(left.output, right.output, Inner, None, left, right)
+ assert(smj.logicalLink.isEmpty)
+
+ val kept = smj
+ assert(RewriteJoin.rewrite(kept, spark.sessionState.conf) eq kept)
+ assert(!hasFallbackReason(kept))
+ val info = kept.getTagValue(CometExplainInfo.EXTENSION_INFO)
+ assert(
+ info.exists(_.contains(rewriteReason("no statistics are available for the build side"))),
+ info)
+
+ withSQLConf(CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") {
+ val rewritten = RewriteJoin.rewrite(smj, spark.sessionState.conf)
+ assert(rewritten.isInstanceOf[ShuffledHashJoinExec], rewritten)
+ }
+ }
+
+ // Both joins are on the same key, so EnsureRequirements puts no sort between them: the inner
+ // sort-merge join's output ordering satisfies the outer one. When the inner join is rewritten
+ // to a hash join and the outer one is kept, the outer join needs that ordering restored.
+ private def withChainedJoinTables(midRows: Int, midKeys: Int)(f: => Unit): Unit = {
+ withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") {
+ withParquetTable((0 until 10).map(i => (i * 10, i)), "small") {
+ withParquetTable((0 until midRows).map(i => (i % midKeys, i)), "mid") {
+ f
+ }
+ }
+ }
+ }
+
+ // Whether a hash join is reachable without leaving the stage.
+ private def reachesHashJoinInStage(plan: SparkPlan): Boolean = plan match {
+ case _: CometHashJoinExec => true
+ case _: ShuffleExchangeLike | _: QueryStageExec => false
+ case p => p.children.exists(reachesHashJoinInStage)
+ }
+
+ // Whether a sort-merge join input is a sort over a hash join in the same stage. A Spark
+ // sort-merge join reads the sort through codegen wrappers and a columnar to row transition.
+ private def isSortOverHashJoin(input: SparkPlan): Boolean = input match {
+ case s: CometSortExec => reachesHashJoinInStage(s.child)
+ case s: SortExec => reachesHashJoinInStage(s.child)
+ case t: ColumnarToRowTransition => isSortOverHashJoin(t.child)
+ case w: WholeStageCodegenExec => isSortOverHashJoin(w.child)
+ case a: InputAdapter => isSortOverHashJoin(a.child)
+ case _ => false
+ }
+
+ // Checks the answer first, then that a kept sort-merge join reads the rewritten hash join
+ // through a sort.
+ private def checkSortRestoredOverHashJoin(query: String): Unit = {
+ val (_, cometPlan) = checkSparkAnswer(query)
+ val inputs = collect(cometPlan) {
+ case j: CometSortMergeJoinExec => j.children
+ case j: SortMergeJoinExec => j.children
+ }.flatten
+ assert(inputs.exists(isSortOverHashJoin), cometPlan)
+ }
+
+ private val chainedJoinQuery = "SELECT big._1, big._2, small._2, mid._2 FROM big " +
+ "JOIN small ON big._1 = small._1 JOIN mid ON big._1 = mid._1"
+
+ test("forceShuffledHashJoin restores ordering for a kept SortMergeJoin, AQE=false") {
+ // The default limit, 3000 bytes times two partitions, falls between the two build sides.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "3000") {
+ withChainedJoinTables(midRows = 3000, midKeys = 100) {
+ checkSortRestoredOverHashJoin(chainedJoinQuery)
+ }
+ }
+ }
+
+ test("forceShuffledHashJoin restores ordering for a kept SortMergeJoin, AQE=true") {
+ // Under AQE the decision uses the materialized shuffle sizes, so mid is made large enough
+ // to stay over the limit while small stays under it.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1000000",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "3000") {
+ withChainedJoinTables(midRows = 300000, midKeys = 30000) {
+ checkSortRestoredOverHashJoin(chainedJoinQuery)
+ }
+ }
+ }
+
+ // The rewrite never converts a LeftSemi or ExistenceJoin that builds on the right, whatever
+ // the build size, so these joins are kept above the rewritten inner join with no size limit.
+ // Each predicate also refers to small, so neither the filter nor the semi join is pushed below
+ // the inner join.
+ for (adaptive <- Seq(false, true)) {
+ def withKeptSemiJoinConf(f: => Unit): Unit = withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") {
+ withChainedJoinTables(midRows = 3000, midKeys = 100)(f)
+ }
+
+ test(s"forceShuffledHashJoin restores ordering for a kept ExistenceJoin, AQE=$adaptive") {
+ withKeptSemiJoinConf {
+ checkSortRestoredOverHashJoin(
+ "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " +
+ "WHERE small._2 < 0 OR EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1)")
+ }
+ }
+
+ test(s"forceShuffledHashJoin restores ordering for a kept LeftSemi join, AQE=$adaptive") {
+ withSQLConf(
+ SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.catalyst.optimizer.PushLeftSemiLeftAntiThroughJoin") {
+ withKeptSemiJoinConf {
+ checkSortRestoredOverHashJoin(
+ "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " +
+ "WHERE EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1 AND mid._2 > small._2)")
+ }
+ }
+ }
+ }
+
+ // A window partitioned by the join key gets no sort of its own, since the sort-merge join's
+ // output ordering satisfies it, so the window needs that ordering restored too.
+ for (adaptive <- Seq(false, true)) {
+ test(s"forceShuffledHashJoin restores ordering for a window over the join, AQE=$adaptive") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ CometConf.COMET_FORCE_SHJ.key -> "true",
+ CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") {
+ withChainedJoinTables(midRows = 3000, midKeys = 100) {
+ val (_, cometPlan) = checkSparkAnswer(
+ "SELECT big._1, big._2, small._2, count(*) OVER (PARTITION BY big._1) " +
+ "FROM big JOIN small ON big._1 = small._1")
+ assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan)
+ assert(collect(cometPlan) { case w: CometWindowExec => w }.nonEmpty, cometPlan)
+ }
+ }
+ }
+ }
+
test("BroadcastHashJoin with LeftAnti and NOT IN subquery (null-aware)") {
withSQLConf(
SQLConf.PREFER_SORTMERGEJOIN.key -> "false",
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org