You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ch...@apache.org on 2026/03/18 13:09:09 UTC
(spark) branch master updated: [SPARK-44065][SQL] Optimize BroadcastHashJoin skew in OptimizeSkewedJoin
This is an automated email from the ASF dual-hosted git repository.
chengpan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 99044a856bbb [SPARK-44065][SQL] Optimize BroadcastHashJoin skew in OptimizeSkewedJoin
99044a856bbb is described below
commit 99044a856bbb138d1ae7b50bc8bdef9f53e8a25d
Author: wforget <64...@qq.com>
AuthorDate: Wed Mar 18 21:08:52 2026 +0800
[SPARK-44065][SQL] Optimize BroadcastHashJoin skew in OptimizeSkewedJoin
### What changes were proposed in this pull request?
Split the skewed partitions of BroadcastHashJoin streamPlan when localShuffleReader is disabled.
+ Add `isSkewJoin` attribute in BroadcastHashJoinExec.
+ Handle skewed BroadcastHashJoinExec in `OptimizeSkewedJoin`.
+ Calculate skewed BroadcastHashJoinExec in `SimpleCostEvaluator.evaluateCost`.
### Why are the changes needed?
In RemoteShuffleService services such as uniffle and celeborn, it is recommended to disable localShuffleReader by default for better performance, but it may make BroadcastHashJoin skew again.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
new unit test
Closes #41609 from wForget/SPARK-44065.
Authored-by: wforget <64...@qq.com>
Signed-off-by: Cheng Pan <ch...@apache.org>
---
.../OptimizeSkewInRebalancePartitions.scala | 37 ++---------
.../execution/adaptive/OptimizeSkewedJoin.scala | 72 +++++++++++++++++-----
.../PlanAdaptiveDynamicPruningFilters.scala | 4 +-
.../execution/adaptive/ShufflePartitionsUtil.scala | 29 +++++++++
.../sql/execution/adaptive/simpleCosting.scala | 3 +-
.../dynamicpruning/PlanDynamicPruningFilters.scala | 4 +-
.../execution/joins/BroadcastHashJoinExec.scala | 13 +++-
.../spark/sql/DynamicPartitionPruningSuite.scala | 2 +-
.../sql/execution/WholeStageCodegenSuite.scala | 8 +--
.../adaptive/AdaptiveQueryExecSuite.scala | 62 +++++++++++++++++++
10 files changed, 172 insertions(+), 62 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewInRebalancePartitions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewInRebalancePartitions.scala
index abd096b9c7c3..de7cfca158c3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewInRebalancePartitions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewInRebalancePartitions.scala
@@ -17,7 +17,7 @@
package org.apache.spark.sql.execution.adaptive
-import org.apache.spark.sql.execution.{CoalescedPartitionSpec, ShufflePartitionSpec, SparkPlan}
+import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.exchange.{REBALANCE_PARTITIONS_BY_COL, REBALANCE_PARTITIONS_BY_NONE, ShuffleOrigin}
import org.apache.spark.sql.internal.SQLConf
@@ -39,35 +39,6 @@ object OptimizeSkewInRebalancePartitions extends AQEShuffleReadRule {
override val supportedShuffleOrigins: Seq[ShuffleOrigin] =
Seq(REBALANCE_PARTITIONS_BY_NONE, REBALANCE_PARTITIONS_BY_COL)
- /**
- * Splits the skewed partition based on the map size and the target partition size
- * after split. Create a list of `PartialReducerPartitionSpec` for skewed partition and
- * create `CoalescedPartition` for normal partition.
- */
- private def optimizeSkewedPartitions(
- shuffleId: Int,
- bytesByPartitionId: Array[Long],
- targetSize: Long): Seq[ShufflePartitionSpec] = {
- val smallPartitionFactor =
- conf.getConf(SQLConf.ADAPTIVE_REBALANCE_PARTITIONS_SMALL_PARTITION_FACTOR)
- bytesByPartitionId.indices.flatMap { reduceIndex =>
- val bytes = bytesByPartitionId(reduceIndex)
- if (bytes > targetSize) {
- val newPartitionSpec = ShufflePartitionsUtil.createSkewPartitionSpecs(
- shuffleId, reduceIndex, targetSize, smallPartitionFactor)
- if (newPartitionSpec.isEmpty) {
- CoalescedPartitionSpec(reduceIndex, reduceIndex + 1, bytes) :: Nil
- } else {
- logDebug(s"For shuffle $shuffleId, partition $reduceIndex is skew, " +
- s"split it into ${newPartitionSpec.get.size} parts.")
- newPartitionSpec.get
- }
- } else {
- CoalescedPartitionSpec(reduceIndex, reduceIndex + 1, bytes) :: Nil
- }
- }
- }
-
private def tryOptimizeSkewedPartitions(shuffle: ShuffleQueryStageExec): SparkPlan = {
val defaultAdvisorySize = conf.getConf(SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES)
val advisorySize = shuffle.advisoryPartitionSize.getOrElse(defaultAdvisorySize)
@@ -77,8 +48,10 @@ object OptimizeSkewInRebalancePartitions extends AQEShuffleReadRule {
return shuffle
}
- val newPartitionsSpec = optimizeSkewedPartitions(
- mapStats.get.shuffleId, mapStats.get.bytesByPartitionId, advisorySize)
+ val smallPartitionFactor =
+ conf.getConf(SQLConf.ADAPTIVE_REBALANCE_PARTITIONS_SMALL_PARTITION_FACTOR)
+ val newPartitionsSpec = ShufflePartitionsUtil.optimizeSkewedPartitions(mapStats.get.shuffleId,
+ mapStats.get.bytesByPartitionId, advisorySize, advisorySize, smallPartitionFactor)
// return origin plan if we can not optimize partitions
if (newPartitionsSpec.length == mapStats.get.bytesByPartitionId.length) {
shuffle
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala
index c256b3fcb6bc..7f55487407a7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala
@@ -23,12 +23,13 @@ import org.apache.spark.SparkUnsupportedOperationException
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder}
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
import org.apache.spark.sql.catalyst.plans._
import org.apache.spark.sql.catalyst.plans.physical.Partitioning
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.exchange.{ENSURE_REQUIREMENTS, EnsureRequirements, ValidateRequirements}
-import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec}
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.util.Utils
@@ -197,23 +198,60 @@ case class OptimizeSkewedJoin(ensureRequirements: EnsureRequirements)
}
}
- def optimizeSkewJoin(plan: SparkPlan): SparkPlan = plan.transformUp {
- case smj @ SortMergeJoinExec(_, _, joinType, _,
- s1 @ SortExec(_, _, ShuffleStage(left: ShuffleQueryStageExec), _),
- s2 @ SortExec(_, _, ShuffleStage(right: ShuffleQueryStageExec), _), false) =>
- tryOptimizeJoinChildren(left, right, joinType).map {
- case (newLeft, newRight) =>
- smj.copy(
- left = s1.copy(child = newLeft), right = s2.copy(child = newRight), isSkewJoin = true)
- }.getOrElse(smj)
+ /**
+ * Split the skewed partition of the BroadcastHashJoin streamPlan.
+ */
+ private def tryOptimizeBroadcastHashJoinStreamedPlan(
+ shuffle: ShuffleQueryStageExec): Option[SparkPlan] = {
+ val mapStats = shuffle.mapStats.get
+ val bytesByPartitionId = mapStats.bytesByPartitionId
+ val medSize = Utils.median(bytesByPartitionId, false)
+ val skewThreshold = getSkewThreshold(medSize)
+ val streamedPlanTargetSize = targetSize(bytesByPartitionId, skewThreshold)
+
+ val newPartitionsSpec = ShufflePartitionsUtil.optimizeSkewedPartitions(mapStats.shuffleId,
+ bytesByPartitionId, skewThreshold, streamedPlanTargetSize)
+ // return origin plan if we can not optimize partitions
+ if (newPartitionsSpec.length == bytesByPartitionId.length) {
+ None
+ } else {
+ Some(SkewJoinChildWrapper(AQEShuffleReadExec(shuffle, newPartitionsSpec)))
+ }
+ }
+
+ def optimizeSkewJoin(plan: SparkPlan): SparkPlan = {
+ val localShuffleReaderEnabled = conf.getConf(SQLConf.LOCAL_SHUFFLE_READER_ENABLED)
+ plan.transformUp {
+ case smj @ SortMergeJoinExec(_, _, joinType, _,
+ s1 @ SortExec(_, _, ShuffleStage(left: ShuffleQueryStageExec), _),
+ s2 @ SortExec(_, _, ShuffleStage(right: ShuffleQueryStageExec), _), false) =>
+ tryOptimizeJoinChildren(left, right, joinType).map {
+ case (newLeft, newRight) =>
+ smj.copy(
+ left = s1.copy(child = newLeft), right = s2.copy(child = newRight), isSkewJoin = true)
+ }.getOrElse(smj)
- case shj @ ShuffledHashJoinExec(_, _, joinType, _, _,
- ShuffleStage(left: ShuffleQueryStageExec),
- ShuffleStage(right: ShuffleQueryStageExec), false) =>
- tryOptimizeJoinChildren(left, right, joinType).map {
- case (newLeft, newRight) =>
- shj.copy(left = newLeft, right = newRight, isSkewJoin = true)
- }.getOrElse(shj)
+ case shj @ ShuffledHashJoinExec(_, _, joinType, _, _,
+ ShuffleStage(left: ShuffleQueryStageExec),
+ ShuffleStage(right: ShuffleQueryStageExec), false) =>
+ tryOptimizeJoinChildren(left, right, joinType).map {
+ case (newLeft, newRight) =>
+ shj.copy(left = newLeft, right = newRight, isSkewJoin = true)
+ }.getOrElse(shj)
+
+ // When localShuffleReader is disabled, the streamedPlan of BroadcastHashJoin may be skewed,
+ // so we try to optimize it.
+ case bhj @ BroadcastHashJoinExec(_, _, _, BuildRight, _,
+ ShuffleStage(left: ShuffleQueryStageExec), _, _, false) if !localShuffleReaderEnabled =>
+ tryOptimizeBroadcastHashJoinStreamedPlan(left).map {
+ case newLeft => bhj.copy(left = newLeft, isSkewJoin = true)
+ }.getOrElse(bhj)
+ case bhj @ BroadcastHashJoinExec(_, _, _, BuildLeft, _, _,
+ ShuffleStage(right: ShuffleQueryStageExec), _, false) if !localShuffleReaderEnabled =>
+ tryOptimizeBroadcastHashJoinStreamedPlan(right).map {
+ case newRight => bhj.copy(right = newRight, isSkewJoin = true)
+ }.getOrElse(bhj)
+ }
}
override def apply(plan: SparkPlan): SparkPlan = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala
index 751cfe5b7bb6..6b721d5442e2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala
@@ -53,9 +53,9 @@ case class PlanAdaptiveDynamicPruningFilters(
val canReuseExchange = conf.exchangeReuseEnabled && buildKeys.nonEmpty &&
find(rootPlan) {
- case BroadcastHashJoinExec(_, _, _, BuildLeft, _, left, _, _) =>
+ case BroadcastHashJoinExec(_, _, _, BuildLeft, _, left, _, _, _) =>
left.sameResult(exchange)
- case BroadcastHashJoinExec(_, _, _, BuildRight, _, _, right, _) =>
+ case BroadcastHashJoinExec(_, _, _, BuildRight, _, _, right, _, _) =>
right.sameResult(exchange)
case _ => false
}.isDefined
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
index fc7d556a6bf7..7b161bd17f42 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
@@ -425,4 +425,33 @@ object ShufflePartitionsUtil extends Logging {
None
}
}
+
+ /**
+ * Splits the skewed partition based on the map size and the target partition size
+ * after split. Create a list of `PartialReducerPartitionSpec` for skewed partition and
+ * create `CoalescedPartition` for normal partition.
+ */
+ def optimizeSkewedPartitions(
+ shuffleId: Int,
+ bytesByPartitionId: Array[Long],
+ skewThreshold: Long,
+ targetSize: Long,
+ smallPartitionFactor: Double = SMALL_PARTITION_FACTOR): Seq[ShufflePartitionSpec] = {
+ bytesByPartitionId.indices.flatMap { reduceIndex =>
+ val bytes = bytesByPartitionId(reduceIndex)
+ if (bytes > skewThreshold) {
+ val newPartitionSpec = createSkewPartitionSpecs(
+ shuffleId, reduceIndex, targetSize, smallPartitionFactor)
+ if (newPartitionSpec.isEmpty) {
+ CoalescedPartitionSpec(reduceIndex, reduceIndex + 1, bytes) :: Nil
+ } else {
+ logDebug(s"For shuffle $shuffleId, partition $reduceIndex is skew, " +
+ s"split it into ${newPartitionSpec.get.size} parts.")
+ newPartitionSpec.get
+ }
+ } else {
+ CoalescedPartitionSpec(reduceIndex, reduceIndex + 1, bytes) :: Nil
+ }
+ }
+ }
}
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 9c9c8e13d2d5..28b757114ebe 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
@@ -20,7 +20,7 @@ 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.exchange.ShuffleExchangeLike
-import org.apache.spark.sql.execution.joins.ShuffledJoin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledJoin}
/**
* A simple implementation of [[Cost]], which takes a number of [[Long]] as the cost value.
@@ -48,6 +48,7 @@ case class SimpleCostEvaluator(forceOptimizeSkewedJoin: Boolean) extends CostEva
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.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala
index fbd341b6e7b8..fdcb78fdb55a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala
@@ -60,9 +60,9 @@ case class PlanDynamicPruningFilters(sparkSession: SparkSession) extends Rule[Sp
// the first to be applied (apart from `InsertAdaptiveSparkPlan`).
val canReuseExchange = conf.exchangeReuseEnabled && buildKeys.nonEmpty &&
plan.exists {
- case BroadcastHashJoinExec(_, _, _, BuildLeft, _, left, _, _) =>
+ case BroadcastHashJoinExec(_, _, _, BuildLeft, _, left, _, _, _) =>
left.sameResult(sparkPlan)
- case BroadcastHashJoinExec(_, _, _, BuildRight, _, _, right, _) =>
+ case BroadcastHashJoinExec(_, _, _, BuildRight, _, _, right, _, _) =>
right.sameResult(sparkPlan)
case _ => false
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala
index 4fd5eb50dbfc..e4f18c9144dd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala
@@ -45,9 +45,14 @@ case class BroadcastHashJoinExec private(
condition: Option[Expression],
left: SparkPlan,
right: SparkPlan,
- isNullAwareAntiJoin: Boolean = false)
+ isNullAwareAntiJoin: Boolean = false,
+ isSkewJoin: Boolean = false)
extends HashJoin {
+ override def nodeName: String = {
+ if (isSkewJoin) super.nodeName + "(skew=true)" else super.nodeName
+ }
+
if (isNullAwareAntiJoin) {
require(leftKeys.length == 1, "leftKeys length should be 1")
require(rightKeys.length == 1, "rightKeys length should be 1")
@@ -255,7 +260,8 @@ object BroadcastHashJoinExec extends JoinSelectionHelper {
condition: Option[Expression],
left: SparkPlan,
right: SparkPlan,
- isNullAwareAntiJoin: Boolean = false): BroadcastHashJoinExec = {
+ isNullAwareAntiJoin: Boolean = false,
+ isSkewJoin: Boolean = false): BroadcastHashJoinExec = {
val (normalizedLeftKeys, normalizedRightKeys) = HashJoin.normalizeJoinKeys(leftKeys, rightKeys)
new BroadcastHashJoinExec(
@@ -266,6 +272,7 @@ object BroadcastHashJoinExec extends JoinSelectionHelper {
condition,
left,
right,
- isNullAwareAntiJoin)
+ isNullAwareAntiJoin,
+ isSkewJoin)
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
index e1a2fd33c7c9..83d7d2695652 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
@@ -345,7 +345,7 @@ abstract class DynamicPartitionPruningSuiteBase
""".stripMargin)
val found = df.queryExecution.executedPlan.exists {
- case BroadcastHashJoinExec(_, _, p: ExistenceJoin, _, _, _, _, _) => true
+ case BroadcastHashJoinExec(_, _, p: ExistenceJoin, _, _, _, _, _, _) => true
case _ => false
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
index a3cfdc5a240a..6004f6f76394 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala
@@ -813,8 +813,8 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
val distinctWithId = baseTable.distinct().withColumn("id", monotonically_increasing_id())
.join(baseTable, "idx")
assert(distinctWithId.queryExecution.executedPlan.exists {
- case WholeStageCodegenExec(
- ProjectExec(_, BroadcastHashJoinExec(_, _, _, _, _, _: HashAggregateExec, _, _))) => true
+ case WholeStageCodegenExec(ProjectExec(_,
+ BroadcastHashJoinExec(_, _, _, _, _, _: HashAggregateExec, _, _, _))) => true
case _ => false
})
checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0)))
@@ -825,8 +825,8 @@ class WholeStageCodegenSuite extends QueryTest with SharedSparkSession
baseTable.groupBy("idx").sum().withColumn("id", monotonically_increasing_id())
.join(baseTable, "idx")
assert(groupByWithId.queryExecution.executedPlan.exists {
- case WholeStageCodegenExec(
- ProjectExec(_, BroadcastHashJoinExec(_, _, _, _, _, _: HashAggregateExec, _, _))) => true
+ case WholeStageCodegenExec(ProjectExec(_,
+ BroadcastHashJoinExec(_, _, _, _, _, _: HashAggregateExec, _, _, _))) => true
case _ => false
})
checkAnswer(groupByWithId, Seq(Row(1, 2, 0), Row(1, 2, 0)))
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 06ed672d5815..a3e6e65d5027 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
@@ -3452,6 +3452,68 @@ class AdaptiveQueryExecSuite
}
}
}
+
+ test("SPARK-44065: Optimize BroadcastHashJoin skew") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "1000",
+ SQLConf.LOCAL_SHUFFLE_READER_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "10",
+ SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "600",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "600") {
+ withTempView("skewData", "smallData") {
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("if(id >= 5, 5, id) as key1", "id as value1")
+ .createOrReplaceTempView("skewData1")
+ spark
+ .range(0, 5, 1, 10)
+ .selectExpr("id key2", "id as value2")
+ .createOrReplaceTempView("smallData")
+
+ Seq(true, false).foreach { localShuffleReader =>
+ withSQLConf(SQLConf.LOCAL_SHUFFLE_READER_ENABLED.key -> localShuffleReader.toString) {
+ val sqlText =
+ s"""
+ |select * from skewData1 a join smallData b on a.key1 = b.key2
+ |""".stripMargin
+ val (_, plan) = runAdaptiveAndVerifyResult(sqlText)
+ val bhjs = findTopLevelBroadcastHashJoin(plan)
+ assert(bhjs.nonEmpty)
+
+ if (localShuffleReader) {
+ val localShuffleReaders = collect(plan) {
+ case c: AQEShuffleReadExec if c.isLocalRead => c
+ }
+ assert(localShuffleReaders.nonEmpty)
+ } else {
+ val skewedShuffleReaders = collect(plan) {
+ case c: AQEShuffleReadExec if c.hasSkewedPartition => c
+ }
+ assert(skewedShuffleReaders.nonEmpty)
+ }
+ }
+ }
+
+ withSQLConf(SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN.key -> "true") {
+ val sqlText =
+ s"""
+ |select a.key1, count(*) from skewData1 a join smallData b
+ | on a.key1 = b.key2 group by a.key1
+ |""".stripMargin
+ val (_, plan) = runAdaptiveAndVerifyResult(sqlText)
+ val bhjs = findTopLevelBroadcastHashJoin(plan)
+ assert(bhjs.nonEmpty)
+
+ val skewedBroadcastHashJoins = collect(plan) {
+ case c: BroadcastHashJoinExec if c.isSkewJoin => c
+ }
+ assert(skewedBroadcastHashJoins.nonEmpty)
+ }
+ }
+ }
+ }
}
/**
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org