diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md
index bbb449360287..df6bb606cb8c 100644
--- a/docs/sql-performance-tuning.md
+++ b/docs/sql-performance-tuning.md
@@ -375,6 +375,30 @@ AQE converts sort-merge join to shuffled hash join when all post shuffle partiti
| 3.2.0 |
+
+ spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled |
+ false |
+
+ When true, Spark converts a sort-merge join to a shuffled hash join during adaptive execution when the build side's materialized per-partition sizes are all within spark.sql.adaptive.maxShuffledHashJoinLocalMapThreshold (which additionally requires spark.sql.adaptive.advisoryPartitionSizeInBytes to not be larger than it), even when non-shuffle operators (such as aggregate, project, filter and window) sit between the join and its input shuffle.
+ |
+ 4.3.0 |
+
+
+ spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.minWideningFactor |
+ 1.0 |
+
+ The lower bound applied to the row-widening factor used by spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled when bounding a build side's shuffled hash map size. The factor scales the input shuffle bytes by the estimated per-row size growth of the operators between the join and its shuffle; a larger lower bound is more conservative and makes the conversion less likely when statistics may under-estimate the build size. Must be positive.
+ |
+ 4.3.0 |
+
+
+ spark.sql.adaptive.costEvaluator.countLocalSort.enabled |
+ (value of spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled) |
+
+ When true, the default AQE cost evaluator also counts the number of local sorts as a lower-priority tiebreaker below the number of shuffles, so that among plans with the same number of shuffles the one with fewer local sorts is preferred. For example, a sort-merge join is replaced by a shuffled hash join only when the conversion does not push extra sorts elsewhere in the plan. Defaults to the value of spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled, so it is enabled together with that conversion.
+ |
+ 4.3.0 |
+
### Optimizing Skew Join
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala
index 66c4a39ce823..71f30ca49d86 100755
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala
@@ -1288,7 +1288,7 @@ case class FindInSet(left: Expression, right: Expression) extends BinaryExpressi
trait String2TrimExpression extends Expression with ImplicitCastInputTypes {
- protected def srcStr: Expression
+ private[sql] def srcStr: Expression
protected def trimStr: Option[Expression]
protected def direction: String
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala
index 13e3cb76805d..833a0e80ab02 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala
@@ -20,8 +20,10 @@ package org.apache.spark.sql.catalyst.optimizer
import scala.annotation.tailrec
import scala.util.control.NonFatal
+import org.apache.spark.MapOutputStatistics
import org.apache.spark.internal.Logging
import org.apache.spark.internal.LogKeys.{HASH_JOIN_KEYS, JOIN_CONDITION}
+import org.apache.spark.sql.catalyst.SQLConfHelper
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
import org.apache.spark.sql.catalyst.planning.{ExtractEquiJoinKeys, ExtractFiltersAndInnerJoins, ExtractSingleColumnNullAwareAntiJoin}
@@ -287,7 +289,18 @@ case object BuildRight extends BuildSide
case object BuildLeft extends BuildSide
-trait JoinSelectionHelper extends Logging {
+trait JoinSelectionHelper extends SQLConfHelper with Logging {
+
+ def preferShuffledHashJoin(
+ mapStats: MapOutputStatistics,
+ sizeInBytesFactor: Double = 1.0): Boolean = {
+ val maxShuffledHashJoinLocalMapThreshold =
+ conf.getConf(SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD)
+ val advisoryPartitionSize = conf.getConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES)
+ advisoryPartitionSize <= maxShuffledHashJoinLocalMapThreshold &&
+ mapStats.bytesByPartitionId.forall(
+ _ * sizeInBytesFactor <= maxShuffledHashJoinLocalMapThreshold)
+ }
def getBroadcastBuildSide(
join: Join,
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 6fe7a957a121..8f44a72b4dec 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -1337,6 +1337,46 @@ object SQLConf {
.bytesConf(ByteUnit.BYTE)
.createWithDefault(0L)
+ val ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED =
+ buildConf("spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled")
+ .doc("When true, Spark converts a sort merge join to a shuffled hash join during adaptive " +
+ "execution when the build side's materialized per-partition sizes are all within " +
+ s"${ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key} (which additionally requires " +
+ s"${ADVISORY_PARTITION_SIZE_IN_BYTES.key} to not be larger than it), even when " +
+ "non-shuffle operators sit between the join and its input shuffle.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .booleanConf
+ .createWithDefault(false)
+
+ val ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_MIN_WIDENING_FACTOR =
+ buildConf("spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.minWideningFactor")
+ .doc("The lower bound applied to the row-widening factor used by " +
+ s"${ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key} when bounding a " +
+ "build side's shuffled hash map size. The factor scales the input shuffle bytes by the " +
+ "estimated per-row size growth of the operators between the join and its shuffle; a " +
+ "larger lower bound is more conservative and makes the conversion less likely when " +
+ "statistics may under-estimate the build size. Must be positive.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .doubleConf
+ .checkValue(_ > 0, "The minimum widening factor must be positive.")
+ .createWithDefault(1.0)
+
+ val ADAPTIVE_COST_EVALUATOR_COUNT_LOCAL_SORT_ENABLED =
+ buildConf("spark.sql.adaptive.costEvaluator.countLocalSort.enabled")
+ .doc("When true, the default AQE cost evaluator also counts the number of local sorts as a " +
+ "lower-priority tiebreaker below the number of shuffles. This lets adaptive execution " +
+ "prefer a plan with fewer local sorts among plans with the same number of shuffles, for " +
+ s"example a shuffled hash join produced by " +
+ s"${ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key} over a " +
+ "sort merge join when the conversion does not push extra sorts elsewhere. Defaults to " +
+ s"${ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key} so that it is " +
+ "enabled together with that conversion.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .fallbackConf(ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED)
+
val ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED =
buildConf("spark.sql.adaptive.optimizeSkewsInRebalancePartitions.enabled")
.doc(s"When true and '${ADAPTIVE_EXECUTION_ENABLED.key}' is true, Spark will optimize the " +
@@ -8090,6 +8130,15 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
def nonEmptyPartitionRatioForBroadcastJoin: Double =
getConf(NON_EMPTY_PARTITION_RATIO_FOR_BROADCAST_JOIN)
+ def convertSortMergeJoinToShuffledHashJoinEnabled: Boolean =
+ getConf(ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED)
+
+ def convertSortMergeJoinToShuffledHashJoinMinWideningFactor: Double =
+ getConf(ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_MIN_WIDENING_FACTOR)
+
+ def costEvaluatorCountLocalSortEnabled: Boolean =
+ getConf(ADAPTIVE_COST_EVALUATOR_COUNT_LOCAL_SORT_ENABLED)
+
def coalesceShufflePartitionsEnabled: Boolean = getConf(COALESCE_PARTITIONS_ENABLED)
def minBatchesToRetain: Int = getConf(MIN_BATCHES_TO_RETAIN)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
index 9a483076ff56..7040ab51cf51 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
@@ -102,7 +102,9 @@ case class AdaptiveSparkPlanExec(
conf.getConf(SQLConf.ADAPTIVE_CUSTOM_COST_EVALUATOR_CLASS) match {
case Some(className) =>
CostEvaluator.instantiate(className, context.session.sparkContext.getConf)
- case _ => SimpleCostEvaluator(conf.getConf(SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN))
+ case _ => SimpleCostEvaluator(
+ conf.getConf(SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN),
+ conf.costEvaluatorCountLocalSortEnabled)
}
// A list of physical plan rules to be applied before creation of query stages. The physical
@@ -125,6 +127,10 @@ case class AdaptiveSparkPlanExec(
InsertSortForLimitAndOffset,
AdjustShuffleExchangePosition,
ValidateSparkPlan,
+ // Must run before `ReplaceHashWithSortAgg`: converting a sort merge join to a shuffled hash
+ // join drops its child ordering, which `ReplaceHashWithSortAgg` would otherwise rely on to
+ // turn a hash aggregate into a sort aggregate.
+ ConvertSortMergeJoinToShuffledHashJoin(ensureRequirements),
ReplaceHashWithSortAgg,
RemoveRedundantSorts,
RemoveRedundantWindowGroupLimits,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ConvertSortMergeJoinToShuffledHashJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ConvertSortMergeJoinToShuffledHashJoin.scala
new file mode 100644
index 000000000000..ea6095638c76
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ConvertSortMergeJoinToShuffledHashJoin.scala
@@ -0,0 +1,229 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.adaptive
+
+import scala.annotation.tailrec
+
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, CaseWhen, Cast, Coalesce, Expression, If, Literal, Lower, String2TrimExpression, Substring, UnsafeRow, Upper}
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, JoinSelectionHelper}
+import org.apache.spark.sql.catalyst.plans.LeftExistence
+import org.apache.spark.sql.catalyst.plans.logical.Join
+import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{CollectMetricsExec, FilterExec, ProjectExec, SortExec, SparkPlan}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
+import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, EnsureRequirements}
+import org.apache.spark.sql.execution.joins.{BaseJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
+import org.apache.spark.sql.execution.window.{WindowExecBase, WindowGroupLimitExec}
+
+/**
+ * Converts a [[SortMergeJoinExec]] into a [[ShuffledHashJoinExec]] during adaptive execution when
+ * a build side's materialized shuffle statistics show it is small enough for a local hash map.
+ * Unlike [[DynamicJoinSelection]], this runs on the physical plan, so it can reach the input
+ * shuffle through operators (aggregate, project, filter, window, etc...) sitting above it.
+ *
+ * The swap is shuffle-free since both joins are `ShuffledJoin`s with the same distribution and
+ * partitioning; only the child sorts become unnecessary. As a shuffled hash join loses the sort
+ * merge join's output ordering, [[EnsureRequirements]] is re-run to restore any ordering an
+ * ancestor still needs, and AQE's [[CostEvaluator]] decides whether to adopt the converted plan.
+ *
+ * A shuffled hash join builds a non-spillable local hash map, so the traversed operators must not
+ * blow up the build size that the input shuffle statistics estimate. Two guards keep that estimate
+ * a valid bound (see [[ExtractShuffleStage]] and [[selectBuildSide]]):
+ * - the traversal only looks through an operator whose output expressions are all size-bounded
+ * (see [[isSizeBoundedExpr]]), so no operator can widen a row in a way the shuffle statistics
+ * cannot see; and
+ * - the build-side estimate is scaled by [[wideningFactor]] to account for the width change the
+ * traversed operators do introduce.
+ */
+case class ConvertSortMergeJoinToShuffledHashJoin(ensureRequirements: EnsureRequirements)
+ extends Rule[SparkPlan] with JoinSelectionHelper {
+
+ /**
+ * Chooses the build side for the shuffled hash join. A side is eligible only if it is allowed
+ * as a build side for this join type and its (widening-adjusted) input shuffle is small enough
+ * to build a local hash map. When both sides are eligible, the smaller one (by widening-adjusted
+ * total shuffle bytes) is chosen.
+ */
+ private def selectBuildSide(
+ smj: SortMergeJoinExec,
+ left: ShuffleQueryStageExec,
+ right: ShuffleQueryStageExec): Option[BuildSide] = {
+ val leftFactor = wideningFactor(smj.left.output, left.output)
+ val rightFactor = wideningFactor(smj.right.output, right.output)
+ val canBuildLeft = canBuildShuffledHashJoinLeft(smj.joinType) &&
+ preferShuffledHashJoin(left.mapStats.get, leftFactor)
+ val canBuildRight = canBuildShuffledHashJoinRight(smj.joinType) &&
+ preferShuffledHashJoin(right.mapStats.get, rightFactor)
+ if (canBuildLeft && canBuildRight) {
+ val leftSize = left.mapStats.get.bytesByPartitionId.sum * leftFactor
+ val rightSize = right.mapStats.get.bytesByPartitionId.sum * rightFactor
+ if (leftSize < rightSize) Some(BuildLeft) else Some(BuildRight)
+ } else if (canBuildLeft) {
+ Some(BuildLeft)
+ } else if (canBuildRight) {
+ Some(BuildRight)
+ } else {
+ None
+ }
+ }
+
+ /**
+ * The estimated per-row byte-size ratio of the build subtree's output to its input shuffle's
+ * output, i.e. how much the traversed operators widen each row. The traversed operators never
+ * increase the row count (`N_build <= N_shuffle`), so scaling the input shuffle bytes by this
+ * ratio keeps them a valid upper bound on the hash-map build size once row width is accounted
+ * for: `buildSize = N_build * buildRowWidth <= shuffleBytes * (buildRowWidth / shuffleRowWidth)`.
+ *
+ * Floored at `spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.minWideningFactor`
+ * (default 1.0). Unlike `SizeInBytesOnlyStatsPlanVisitor`, which computes a best-effort size and
+ * lets a narrowing operator shrink it, the default keeps a conservative bound for a non-spillable
+ * build: `getSizePerRow` under-estimates a variable-width column (it uses `defaultSize`), so a
+ * `factor < 1` could push the scaled bytes below the real build size and reintroduce the
+ * out-of-memory risk, whereas the raw shuffle bytes are always a valid bound when the build side
+ * is no wider than the shuffle row. Raising the floor above 1.0 is more conservative still.
+ */
+ private def wideningFactor(buildOutput: Seq[Attribute], shuffleOutput: Seq[Attribute]): Double = {
+ val buildRowSize = EstimationUtils.getSizePerRow(buildOutput).toDouble
+ val shuffleRowSize = EstimationUtils.getSizePerRow(shuffleOutput).toDouble
+ math.max(conf.convertSortMergeJoinToShuffledHashJoinMinWideningFactor,
+ buildRowSize / shuffleRowSize)
+ }
+
+ private def hasJoinStrategyHint(smj: SortMergeJoinExec): Boolean = smj.logicalLink.exists {
+ case j: Join =>
+ j.hint.leftHint.exists(_.strategy.isDefined) || j.hint.rightHint.exists(_.strategy.isDefined)
+ case _ => false
+ }
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.convertSortMergeJoinToShuffledHashJoinEnabled) {
+ return plan
+ }
+ val optimizedPlan = plan.transformUp {
+ case smj @ SortMergeJoinExec(leftKeys, rightKeys, joinType, condition,
+ ExtractShuffleStage(left), ExtractShuffleStage(right), false)
+ // Do not convert if the join keys are not hash-join-compatible (e.g. collated or other
+ // non-binary-stable string keys), since a hash join matches keys by binary equality and
+ // would return wrong results. This mirrors the guard on the other SHJ-planning paths.
+ if !hasJoinStrategyHint(smj) && hashJoinSupported(leftKeys, rightKeys) =>
+ selectBuildSide(smj, left, right) match {
+ case Some(buildSide) =>
+ ShuffledHashJoinExec(leftKeys, rightKeys, joinType, buildSide, condition,
+ stripSort(smj.left), stripSort(smj.right))
+ case None => smj
+ }
+ }
+ if (optimizedPlan.fastEquals(plan)) {
+ plan
+ } else {
+ // A shuffled hash join does not preserve the sort merge join's output ordering. Re-run
+ // EnsureRequirements so any ordering an ancestor still needs is re-established, keeping the
+ // plan valid. AQE's CostEvaluator then decides between this plan and the current one.
+ ensureRequirements.apply(optimizedPlan)
+ }
+ }
+
+ /**
+ * Drops a top-level [[SortExec]] since a shuffled hash join does not require sorted input;
+ * [[RemoveRedundantSorts]] cleans up any remaining redundant sorts afterwards.
+ */
+ private def stripSort(plan: SparkPlan): SparkPlan = plan match {
+ case s: SortExec if !s.global => s.child
+ case other => other
+ }
+
+ /**
+ * Finds a join child's input shuffle, looking through the [[SortExec]] and other non-shuffle
+ * operators (aggregate, project, filter, window, left-existence join) above it. Descent stops at
+ * the first [[ShuffleQueryStageExec]], which is thus guaranteed to be the join's own input
+ * shuffle whose statistics bound (or, for a reducing aggregate, upper-bound) the build side. The
+ * stage must be materialized with stats and originate from [[EnsureRequirements]], so swapping
+ * the join type does not change the shuffle.
+ *
+ * A [[ProjectExec]], [[BaseAggregateExec]] or [[WindowExecBase]] is only traversed when all of
+ * its output expressions are size-bounded (see [[isSizeBoundedExpr]]); otherwise the shuffle
+ * bytes could badly under-estimate the non-spillable hash-map build size (e.g.
+ * `repeat(max(c2), 10000)` above a small shuffle), so descent stops and the join is left as is.
+ */
+ object ExtractShuffleStage {
+ def unapply(plan: SparkPlan): Option[ShuffleQueryStageExec] = findShuffleStage(plan)
+
+ @tailrec
+ private def findShuffleStage(plan: SparkPlan): Option[ShuffleQueryStageExec] = plan match {
+ case s: ShuffleQueryStageExec if s.isMaterialized && s.mapStats.isDefined &&
+ s.shuffle.shuffleOrigin == ENSURE_REQUIREMENTS => Some(s)
+ case _: FilterExec | _: SortExec | _: WindowGroupLimitExec | _: CollectMetricsExec =>
+ findShuffleStage(plan.children.head)
+ case p: ProjectExec if p.projectList.forall(isSizeBoundedExpr) =>
+ findShuffleStage(p.child)
+ case a: BaseAggregateExec if a.resultExpressions.forall(isSizeBoundedExpr) =>
+ findShuffleStage(a.child)
+ case w: WindowExecBase if w.windowExpression.forall(isSizeBoundedExpr) =>
+ findShuffleStage(w.child)
+ case join: BaseJoinExec =>
+ join.joinType match {
+ case LeftExistence(_) => findShuffleStage(join.left)
+ case _ => None
+ }
+ case _ => None
+ }
+ }
+
+ /**
+ * Whether `expr`'s result byte-size is bounded by the values it reads, so it cannot widen a row.
+ * An operator all of whose outputs are size-bounded keeps the input shuffle bytes a valid bound
+ * on the non-spillable hash-map build size; an unbounded output (e.g. `repeat` or `concat`, which
+ * synthesize a wider value) makes the shuffle bytes an under-estimate and stops the traversal.
+ *
+ * An [[Attribute]] is always bounded: it refers to a value produced by a descendant operator.
+ * The traversal checks every operator down to the input shuffle, so if a descendant synthesized a
+ * wide value (e.g. a lower `ProjectExec` with `repeat(...)`) this rule stops there; by induction
+ * any attribute that survives is grounded in the shuffle output. Note that aggregate functions do
+ * not appear inline here - a physical aggregate exposes them as result attributes - so an
+ * aggregate result (`max`, and equally an accumulating `collect_list` whose bytes are already in
+ * the shuffle below) is bounded through this same [[Attribute]] case.
+ *
+ * A fixed-width result ([[UnsafeRow.isFixedLength]], stored in an 8-byte word) is bounded
+ * regardless of inputs. Beyond that, only a whitelist of length-non-increasing transforms over
+ * bounded children is accepted; anything else (e.g. `repeat`, `concat`, arithmetic on strings) is
+ * treated as potentially widening.
+ */
+ private def isSizeBoundedExpr(expr: Expression): Boolean = {
+ if (UnsafeRow.isFixedLength(expr.dataType)) {
+ return true
+ }
+ expr match {
+ case _: Attribute | _: Literal => true
+ case e: Alias => isSizeBoundedExpr(e.child)
+ // The Cast is a very common expression, it may slightly increase the size in bytes
+ // but should be tolerated.
+ case e: Cast => isSizeBoundedExpr(e.child)
+ case e: Upper => isSizeBoundedExpr(e.child)
+ case e: Lower => isSizeBoundedExpr(e.child)
+ case e: Substring => isSizeBoundedExpr(e.str)
+ case e: String2TrimExpression => isSizeBoundedExpr(e.srcStr)
+ // Conditionals only pick one of their (bounded) branch values.
+ case If(_, t, f) => isSizeBoundedExpr(t) && isSizeBoundedExpr(f)
+ case CaseWhen(branches, elseValue) =>
+ branches.forall(b => isSizeBoundedExpr(b._2)) && elseValue.forall(isSizeBoundedExpr)
+ case Coalesce(children) => children.forall(isSizeBoundedExpr)
+ case _ => false
+ }
+ }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/DynamicJoinSelection.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/DynamicJoinSelection.scala
index 217569ae645c..67a5d2954bdf 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/DynamicJoinSelection.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/DynamicJoinSelection.scala
@@ -23,7 +23,6 @@ import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys
import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftOuter, RightOuter}
import org.apache.spark.sql.catalyst.plans.logical.{HintInfo, Join, JoinStrategyHint, LogicalPlan, NO_BROADCAST_HASH, PREFER_SHUFFLE_HASH, SHUFFLE_HASH}
import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.internal.SQLConf
/**
* This optimization rule includes three join selection:
@@ -44,14 +43,6 @@ object DynamicJoinSelection extends Rule[LogicalPlan] with JoinSelectionHelper {
(nonZeroCnt * 1.0 / partitionCnt) < conf.nonEmptyPartitionRatioForBroadcastJoin
}
- private def preferShuffledHashJoin(mapStats: MapOutputStatistics): Boolean = {
- val maxShuffledHashJoinLocalMapThreshold =
- conf.getConf(SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD)
- val advisoryPartitionSize = conf.getConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES)
- advisoryPartitionSize <= maxShuffledHashJoinLocalMapThreshold &&
- mapStats.bytesByPartitionId.forall(_ <= maxShuffledHashJoinLocalMapThreshold)
- }
-
private def selectJoinStrategy(
join: Join,
isLeft: Boolean): Option[JoinStrategyHint] = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/simpleCosting.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/simpleCosting.scala
index 28b757114ebe..84da42ff68a6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/simpleCosting.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/simpleCosting.scala
@@ -18,18 +18,33 @@
package org.apache.spark.sql.execution.adaptive
import org.apache.spark.sql.errors.QueryExecutionErrors
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{SortExec, SparkPlan}
import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledJoin}
/**
- * A simple implementation of [[Cost]], which takes a number of [[Long]] as the cost value.
+ * A simple implementation of [[Cost]] produced by [[SimpleCostEvaluator]]. Its three components are
+ * compared lexicographically in priority order:
+ * 1. `numSkewJoins`: more skew joins means lower cost, so it is compared descending and first;
+ * 2. `numShuffles`: fewer shuffles means lower cost;
+ * 3. `numLocalSorts`: the lowest-priority tiebreaker, so among plans with the same number of skew
+ * joins and shuffles the one with fewer local sorts is preferred (e.g. a shuffled hash join
+ * over a sort merge join when the conversion does not push extra sorts elsewhere).
+ *
+ * `numSkewJoins` and `numLocalSorts` are `0` when the corresponding feature is disabled in the
+ * evaluator, so they do not affect the comparison in that case.
*/
-case class SimpleCost(value: Long) extends Cost {
+case class SimpleCost(numSkewJoins: Int, numShuffles: Int, numLocalSorts: Int) extends Cost {
override def compare(that: Cost): Int = that match {
- case SimpleCost(thatValue) =>
- if (value < thatValue) -1 else if (value > thatValue) 1 else 0
+ case SimpleCost(thatSkewJoins, thatShuffles, thatLocalSorts) =>
+ val bySkewJoins = Integer.compare(thatSkewJoins, numSkewJoins)
+ if (bySkewJoins != 0) {
+ bySkewJoins
+ } else {
+ val byShuffles = Integer.compare(numShuffles, thatShuffles)
+ if (byShuffles != 0) byShuffles else Integer.compare(numLocalSorts, thatLocalSorts)
+ }
case _ =>
throw QueryExecutionErrors.cannotCompareCostWithTargetCostError(that.toString)
}
@@ -37,24 +52,23 @@ case class SimpleCost(value: Long) extends Cost {
/**
* A skew join aware implementation of [[CostEvaluator]], which counts the number of
- * [[ShuffleExchangeLike]] nodes and skew join nodes in the plan.
+ * [[ShuffleExchangeLike]] nodes, skew join nodes and (optionally) local [[SortExec]] nodes in the
+ * plan. See [[SimpleCost]] for how the components are compared.
*/
-case class SimpleCostEvaluator(forceOptimizeSkewedJoin: Boolean) extends CostEvaluator {
- override def evaluateCost(plan: SparkPlan): Cost = {
- val numShuffles = plan.collect {
- case s: ShuffleExchangeLike => s
- }.size
+case class SimpleCostEvaluator(forceOptimizeSkewedJoin: Boolean, countLocalSort: Boolean)
+ extends CostEvaluator {
- if (forceOptimizeSkewedJoin) {
- val numSkewJoins = plan.collect {
- case j: ShuffledJoin if j.isSkewJoin => j
- case j: BroadcastHashJoinExec if j.isSkewJoin => j
- }.size
- // We put `-numSkewJoins` in the first 32 bits of the long value, so that it's compared first
- // when comparing the cost, and larger `numSkewJoins` means lower cost.
- SimpleCost(-numSkewJoins.toLong << 32 | numShuffles)
- } else {
- SimpleCost(numShuffles)
+ override def evaluateCost(plan: SparkPlan): Cost = {
+ var numSkewJoins = 0
+ var numShuffles = 0
+ var numLocalSorts = 0
+ plan.foreach {
+ case j: ShuffledJoin if forceOptimizeSkewedJoin && j.isSkewJoin => numSkewJoins += 1
+ case j: BroadcastHashJoinExec if forceOptimizeSkewedJoin && j.isSkewJoin => numSkewJoins += 1
+ case _: ShuffleExchangeLike => numShuffles += 1
+ case s: SortExec if countLocalSort && !s.global => numLocalSorts += 1
+ case _ =>
}
+ SimpleCost(numSkewJoins, numShuffles, numLocalSorts)
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
index 8cf6fbf921da..357a5937e423 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -29,11 +29,11 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListe
import org.apache.spark.shuffle.sort.SortShuffleManager
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, EqualTo, IsNull, Or}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, EqualTo, IsNull, Or, SortOrder}
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
import org.apache.spark.sql.catalyst.plans.{Inner, LeftAnti}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Join, JoinHint, LocalRelation, LogicalPlan}
-import org.apache.spark.sql.catalyst.plans.physical.CoalescedNullAwareHashPartitioning
+import org.apache.spark.sql.catalyst.plans.physical.{CoalescedNullAwareHashPartitioning, SinglePartition}
import org.apache.spark.sql.classic.Strategy
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
@@ -2461,6 +2461,358 @@ class AdaptiveQueryExecSuite
}
}
+ test("SPARK-58084: Replace sort merge join to shuffled hash join through operators") {
+ withTempView("t1", "t2", "t3") {
+ spark.sparkContext.parallelize(
+ (1 to 100).map(i => TestData(i, i.toString)), 10)
+ .toDF("c1", "c2").createOrReplaceTempView("t1")
+ spark.sparkContext.parallelize(
+ (1 to 10).map(i => TestData(i, i.toString)), 5)
+ .toDF("c1", "c2").createOrReplaceTempView("t2")
+
+ // The t2 side has a non-shuffle operator (aggregate, optionally with a filter) between the
+ // join and its input shuffle, so DynamicJoinSelection's hint path does not fire. The new
+ // physical rule looks through those operators to the materialized shuffle stage.
+ val queries = Seq(
+ "SELECT t1.c1, x.cnt FROM t1 JOIN " +
+ "(SELECT c1, count(*) AS cnt FROM t2 GROUP BY c1) x ON t1.c1 = x.c1",
+ "SELECT t1.c1, x.cnt FROM t1 JOIN " +
+ "(SELECT c1, count(*) AS cnt FROM t2 GROUP BY c1 HAVING count(*) >= 0) x " +
+ "ON t1.c1 = x.c1")
+
+ // t1 partition size: [926, 729, 731]; t2 (aggregated) side: [372, 126, 0]. With a small
+ // advisory partition size and a local map threshold of 500, only the t2 side has all
+ // partitions within the threshold, so the join is converted with the t2 side as build side.
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "500") {
+ queries.foreach { query =>
+ // Enabled (default): the sort merge join is converted to a shuffled hash join.
+ withSQLConf(
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key -> "true") {
+ val (origin, adaptive) = runAdaptiveAndVerifyResult(query)
+ assert(findTopLevelSortMergeJoin(origin).size === 1)
+ val shj = findTopLevelShuffledHashJoin(adaptive)
+ assert(shj.size === 1, s"expected a shuffled hash join for query: $query")
+ assert(shj.head.buildSide == BuildRight)
+ assert(findTopLevelSortMergeJoin(adaptive).isEmpty)
+ }
+ // Disabled: the join stays a sort merge join (previous behavior).
+ withSQLConf(
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key -> "false") {
+ val (_, adaptive) = runAdaptiveAndVerifyResult(query)
+ assert(findTopLevelShuffledHashJoin(adaptive).isEmpty,
+ s"expected no shuffled hash join for query: $query")
+ assert(findTopLevelSortMergeJoin(adaptive).size === 1,
+ s"expected a sort merge join for query: $query")
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-58084: do not convert when an operator adds a variable-width column") {
+ withTempView("t1", "t2") {
+ spark.sparkContext.parallelize(
+ (1 to 100).map(i => TestData(i, i.toString)), 10)
+ .toDF("c1", "c2").createOrReplaceTempView("t1")
+ spark.sparkContext.parallelize(
+ (1 to 10).map(i => TestData(i, i.toString)), 5)
+ .toDF("c1", "c2").createOrReplaceTempView("t2")
+
+ // The shuffle below the aggregate is tiny, but the aggregate widens each build row with a
+ // large variable-width string (`repeat(max(c2), 500)`), so the shuffle bytes badly
+ // under-estimate the non-spillable hash-map build size. The traversal must stop at that
+ // widening operator and leave the join as a sort merge join, even though the shuffle looks
+ // small enough for a local hash map. The wide column is selected in the output so column
+ // pruning cannot drop it before the join.
+ val query =
+ "SELECT t1.c1, x.wide FROM t1 JOIN " +
+ "(SELECT c1, repeat(max(c2), 500) AS wide FROM t2 GROUP BY c1) x ON t1.c1 = x.c1"
+
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.execution.adaptive.DynamicJoinSelection",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "500",
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key -> "true") {
+ val (_, adaptive) = runAdaptiveAndVerifyResult(query)
+ assert(findTopLevelShuffledHashJoin(adaptive).isEmpty,
+ "a widening aggregate above the shuffle must keep the join as a sort merge join")
+ assert(findTopLevelSortMergeJoin(adaptive).size === 1)
+ }
+ }
+ }
+
+ test("SPARK-58084: convert through size-bounded (non-widening) operators") {
+ withTempView("t1", "t2") {
+ spark.sparkContext.parallelize(
+ (1 to 100).map(i => TestData(i, i.toString)), 10)
+ .toDF("c1", "c2").createOrReplaceTempView("t1")
+ spark.sparkContext.parallelize(
+ (1 to 10).map(i => TestData(i, i.toString)), 5)
+ .toDF("c1", "c2").createOrReplaceTempView("t2")
+
+ // The aggregate emits a variable-width string column, but only through size-bounded
+ // expressions: `max` selects an existing value, `cast` and `substring` cannot widen it. The
+ // traversal must look through them and still convert the join, unlike the `repeat(...)` case.
+ val query =
+ "SELECT t1.c1, x.m, x.s FROM t1 JOIN " +
+ "(SELECT c1, substring(max(c2), 1, 1) AS m, cast(count(*) AS string) AS s " +
+ "FROM t2 GROUP BY c1) x ON t1.c1 = x.c1"
+
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.execution.adaptive.DynamicJoinSelection",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "100000",
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key -> "true") {
+ val (_, adaptive) = runAdaptiveAndVerifyResult(query)
+ val shj = findTopLevelShuffledHashJoin(adaptive)
+ assert(shj.size === 1,
+ "size-bounded operators above the shuffle must not block the conversion")
+ assert(shj.head.buildSide == BuildRight)
+ assert(findTopLevelSortMergeJoin(adaptive).isEmpty)
+ }
+ }
+ }
+
+ test("SPARK-58084: minWideningFactor makes the size bound more conservative") {
+ withTempView("t1", "t2") {
+ spark.sparkContext.parallelize(
+ (1 to 100).map(i => TestData(i, i.toString)), 10)
+ .toDF("c1", "c2").createOrReplaceTempView("t1")
+ spark.sparkContext.parallelize(
+ (1 to 10).map(i => TestData(i, i.toString)), 5)
+ .toDF("c1", "c2").createOrReplaceTempView("t2")
+
+ // The t2 (aggregated) build side fits the local map threshold at the default widening factor,
+ // so the join converts. A large minWideningFactor scales the estimated build size past the
+ // threshold, so the conversion is rejected and the join stays a sort merge join.
+ val query =
+ "SELECT t1.c1, x.cnt FROM t1 JOIN " +
+ "(SELECT c1, count(*) AS cnt FROM t2 GROUP BY c1) x ON t1.c1 = x.c1"
+
+ def convertsWith(minWideningFactor: String): Boolean = {
+ var converted = false
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.execution.adaptive.DynamicJoinSelection",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "500",
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_MIN_WIDENING_FACTOR.key ->
+ minWideningFactor) {
+ val (_, adaptive) = runAdaptiveAndVerifyResult(query)
+ converted = findTopLevelShuffledHashJoin(adaptive).nonEmpty
+ }
+ converted
+ }
+
+ // Default factor: the build side fits, so the join converts.
+ assert(convertsWith("1.0"), "the join should convert at the default widening factor")
+ // A large factor scales the estimated build size past the threshold, rejecting the conversion.
+ assert(!convertsWith("1000.0"), "a large minWideningFactor should reject the conversion")
+ }
+ }
+
+ test("SPARK-58084: Replace sort merge join keeps required ordering valid") {
+ withTempView("small1", "small2", "big") {
+ spark.sparkContext.parallelize(
+ (1 to 20).map(i => TestData(i, i.toString)), 4)
+ .toDF("c1", "c2").createOrReplaceTempView("small1")
+ spark.sparkContext.parallelize(
+ (1 to 20).map(i => TestData(i, i.toString)), 4)
+ .toDF("c1", "c2").createOrReplaceTempView("small2")
+ spark.sparkContext.parallelize(
+ (1 to 4000).map(i => TestData(i % 20 + 1, i.toString)), 4)
+ .toDF("c1", "c2").createOrReplaceTempView("big")
+
+ // The inner join over the two small tables is convertible, but the outer sort merge join
+ // requires its (left) child ordered on the join key. When the inner join is converted to a
+ // shuffled hash join (ordering Nil), EnsureRequirements must re-insert the sort above it so
+ // the outer sort merge join's required ordering is still satisfied and the result is correct.
+ val query = "SELECT small1.c1 FROM small1 JOIN small2 ON small1.c1 = small2.c1 " +
+ "JOIN big ON small1.c1 = big.c1"
+
+ // Exclude DynamicJoinSelection so only the new physical rule can convert, isolating it.
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.execution.adaptive.DynamicJoinSelection",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "100000",
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key -> "true") {
+ val (_, adaptive) = runAdaptiveAndVerifyResult(query)
+ // The inner join is converted; the outer join stays a sort merge join whose (left) child
+ // ordering is re-established by EnsureRequirements, so the plan remains valid.
+ val smj = findTopLevelSortMergeJoin(adaptive)
+ assert(smj.size === 1)
+ assert(smj.head.left.outputOrdering.nonEmpty,
+ "outer sort merge join must keep its left child ordered on the join key")
+ assert(findTopLevelShuffledHashJoin(adaptive).size === 1)
+ }
+ }
+ }
+
+ test("SPARK-58084: SimpleCostEvaluator counts local sorts as a lower-priority tiebreaker") {
+ def leaf: SparkPlan = CostTestLeafExec()
+ def shuffle(child: SparkPlan): SparkPlan = ShuffleExchangeExec(SinglePartition, child)
+ def localSort(child: SparkPlan): SparkPlan =
+ SortExec(SortOrder(child.output.head, Ascending) :: Nil, global = false, child)
+
+ val evaluator = SimpleCostEvaluator(forceOptimizeSkewedJoin = false, countLocalSort = true)
+ def cost(plan: SparkPlan): Cost = evaluator.evaluateCost(plan)
+
+ // Same number of shuffles: fewer local sorts is cheaper.
+ val oneShuffleTwoSorts = localSort(localSort(shuffle(leaf)))
+ val oneShuffleOneSort = localSort(shuffle(leaf))
+ val oneShuffleNoSort = shuffle(leaf)
+ assert(cost(oneShuffleOneSort).compare(cost(oneShuffleTwoSorts)) < 0)
+ assert(cost(oneShuffleNoSort).compare(cost(oneShuffleOneSort)) < 0)
+
+ // The number of shuffles dominates: a plan with more shuffles is costlier even with no sorts.
+ val twoShufflesNoSort = shuffle(shuffle(leaf))
+ assert(cost(oneShuffleTwoSorts).compare(cost(twoShufflesNoSort)) < 0)
+
+ // When countLocalSort is disabled, local sorts do not affect the cost.
+ val noSortEvaluator = SimpleCostEvaluator(
+ forceOptimizeSkewedJoin = false, countLocalSort = false)
+ assert(noSortEvaluator.evaluateCost(oneShuffleTwoSorts)
+ .compare(noSortEvaluator.evaluateCost(oneShuffleNoSort)) === 0)
+
+ // Skew join dominates, ahead of shuffles and sorts: with forceOptimizeSkewedJoin, a plan with
+ // a skew join is cheaper than one without, even if the skew-join plan has more shuffles and
+ // local sorts.
+ def join(l: SparkPlan, r: SparkPlan, isSkew: Boolean): SparkPlan =
+ SortMergeJoinExec(l.output.take(1), r.output.take(1), Inner, None, l, r, isSkewJoin = isSkew)
+ val skewEvaluator = SimpleCostEvaluator(forceOptimizeSkewedJoin = true, countLocalSort = true)
+ // Skew-join plan: 1 skew join, 3 shuffles, 2 local sorts.
+ val withSkewJoin = skewEvaluator.evaluateCost(
+ join(localSort(shuffle(shuffle(leaf))), localSort(shuffle(leaf)), isSkew = true))
+ // Non-skew plan: 0 skew joins, 2 shuffles, 0 local sorts.
+ val withoutSkewJoin = skewEvaluator.evaluateCost(
+ join(shuffle(leaf), shuffle(leaf), isSkew = false))
+ assert(withSkewJoin.compare(withoutSkewJoin) < 0)
+ }
+
+ test("SPARK-58084: do not convert sort merge join when it adds local sorts") {
+ withTempView("big", "small") {
+ spark.sparkContext.parallelize(
+ (1 to 2000).map(i => TestData(i % 20 + 1, i.toString)), 4)
+ .toDF("k", "v").createOrReplaceTempView("big")
+ spark.sparkContext.parallelize(
+ (1 to 20).map(i => TestData(i, i.toString)), 4)
+ .toDF("k", "v").createOrReplaceTempView("small")
+
+ // Both join sides are sort aggregates grouped by the join key, so each child is already
+ // ordered on the key for free and the sort merge join needs no explicit child sort. A parent
+ // window partitions by the right join key. A sort merge inner join keeps both sides' key
+ // orderings, satisfying the window; a shuffled hash join with build-right keeps only the left
+ // ordering (see HashJoin.outputOrdering), so converting it forces an extra local sort above
+ // the window. The conversion is therefore only beneficial without counting local sorts.
+ val query =
+ "SELECT l.k, count(*) OVER (PARTITION BY r.k) c " +
+ "FROM (SELECT k, count(*) c FROM big GROUP BY k) l " +
+ "JOIN (SELECT k, count(*) c FROM small GROUP BY k) r ON l.k = r.k"
+
+ def countLocalSorts(plan: SparkPlan): Int = collect(plan) {
+ case s: SortExec if !s.global => s
+ }.size
+
+ // Force sort aggregate and exclude DynamicJoinSelection so only the new rule can convert.
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.USE_HASH_AGG.key -> "false",
+ SQLConf.USE_OBJECT_HASH_AGG.key -> "false",
+ SQLConf.ADAPTIVE_OPTIMIZER_EXCLUDED_RULES.key ->
+ "org.apache.spark.sql.execution.adaptive.DynamicJoinSelection",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "100000",
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key -> "true") {
+ // Not counting local sorts: the conversion is adopted even though it adds a local sort.
+ withSQLConf(SQLConf.ADAPTIVE_COST_EVALUATOR_COUNT_LOCAL_SORT_ENABLED.key -> "false") {
+ val (_, adaptive) = runAdaptiveAndVerifyResult(query)
+ assert(findTopLevelShuffledHashJoin(adaptive).size === 1)
+ assert(findTopLevelSortMergeJoin(adaptive).isEmpty)
+ assert(countLocalSorts(adaptive) == 5)
+ }
+ // Counting local sorts: the converted plan has more local sorts, so it is rejected and the
+ // sort merge join is kept.
+ withSQLConf(SQLConf.ADAPTIVE_COST_EVALUATOR_COUNT_LOCAL_SORT_ENABLED.key -> "true") {
+ val (_, adaptive) = runAdaptiveAndVerifyResult(query)
+ assert(findTopLevelShuffledHashJoin(adaptive).isEmpty)
+ assert(findTopLevelSortMergeJoin(adaptive).size === 1)
+ assert(countLocalSorts(adaptive) == 4)
+ }
+ }
+ }
+ }
+
+ test("SPARK-58084: do not convert sort merge join with non-binary-stable (collated) keys") {
+ withTempView("t1", "t2") {
+ spark.sparkContext.parallelize(
+ (1 to 100).map(i => TestData(i, s"v$i")), 10)
+ .toDF("c1", "c2").createOrReplaceTempView("t1")
+ spark.sparkContext.parallelize(
+ (1 to 10).map(i => TestData(i, s"v$i")), 5)
+ .toDF("c1", "c2").createOrReplaceTempView("t2")
+
+ // A UTF8_LCASE key is orderable (so a sort merge join is planned) but not binary-stable. When
+ // the equi-condition wraps the key (here `concat(...)`), `RewriteCollationJoin` does not
+ // inject a `CollationKey`, so the physical join keys stay non-binary-stable. A shuffled hash
+ // join matches keys by `UnsafeRow` binary equality, which would return wrong results, so the
+ // conversion must skip such joins even with the config enabled - mirroring the
+ // `hashJoinSupported` guard on the other SHJ-planning paths.
+ val query =
+ "SELECT t1.c2 FROM t1 JOIN t2 ON " +
+ "concat(cast(t1.c2 AS STRING COLLATE UTF8_LCASE), 'x') = " +
+ "concat(cast(t2.c2 AS STRING COLLATE UTF8_LCASE), 'x')"
+
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "100000",
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key -> "true") {
+ val (_, adaptive) = runAdaptiveAndVerifyResult(query)
+ assert(findTopLevelShuffledHashJoin(adaptive).isEmpty,
+ "non-binary-stable collated keys must keep the join as a sort merge join")
+ assert(findTopLevelSortMergeJoin(adaptive).size === 1)
+ }
+ }
+ }
+
+ test("SPARK-58084: do not convert sort merge join requested with an explicit MERGE hint") {
+ withTempView("t1", "t2") {
+ spark.sparkContext.parallelize(
+ (1 to 100).map(i => TestData(i, i.toString)), 10)
+ .toDF("c1", "c2").createOrReplaceTempView("t1")
+ spark.sparkContext.parallelize(
+ (1 to 10).map(i => TestData(i, i.toString)), 5)
+ .toDF("c1", "c2").createOrReplaceTempView("t2")
+
+ // The join is convertible by size, but the user explicitly asked for a sort merge join with
+ // a MERGE hint. The conversion must respect the hint and keep the sort merge join, matching
+ // DynamicJoinSelection which never overrides an existing join strategy hint.
+ val query = "SELECT /*+ MERGE(t1, t2) */ t1.c1, t2.c2 FROM t1 JOIN t2 ON t1.c1 = t2.c1"
+
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key -> "100000",
+ SQLConf.ADAPTIVE_CONVERT_SORT_MERGE_JOIN_TO_SHUFFLED_HASH_JOIN_ENABLED.key -> "true") {
+ val (_, adaptive) = runAdaptiveAndVerifyResult(query)
+ assert(findTopLevelShuffledHashJoin(adaptive).isEmpty,
+ "an explicit MERGE hint must keep the join as a sort merge join")
+ assert(findTopLevelSortMergeJoin(adaptive).size === 1)
+ }
+ }
+ }
+
test("SPARK-35650: Coalesce number of partitions by AEQ") {
withSQLConf(SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1") {
Seq("REPARTITION", "REBALANCE(key)")
@@ -3709,6 +4061,16 @@ class AdaptiveQueryExecSuite
}
}
+/**
+ * A minimal leaf plan with a single output attribute, used to build tiny plans for cost tests.
+ */
+private case class CostTestLeafExec() extends LeafExecNode {
+ override protected def doExecute(): RDD[InternalRow] =
+ throw SparkException.internalError("should not be executed")
+ override def output: Seq[Attribute] =
+ AttributeReference("a", org.apache.spark.sql.types.IntegerType)() :: Nil
+}
+
/**
* Invalid implementation class for [[CostEvaluator]].
*/
@@ -3723,7 +4085,7 @@ private case class SimpleShuffleSortCostEvaluator() extends CostEvaluator {
case s: ShuffleExchangeLike => s
case s: SortExec => s
}.size
- SimpleCost(cost)
+ SimpleCost(numSkewJoins = 0, numShuffles = cost, numLocalSorts = 0)
}
}