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 fb43c8f9941ea..56ea12918448c 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 @@ -87,9 +87,9 @@ case class AdaptiveSparkPlanExec( // optimizations should be stage-independent. @transient private val queryStageOptimizerRules: Seq[Rule[SparkPlan]] = Seq( ReuseAdaptiveSubquery(conf, context.subqueryCache), - // Here the 'OptimizeSkewedPartitions' rule should be executed + // Here the 'OptimizeSkewedJoin' rule should be executed // before 'ReduceNumShufflePartitions', as the skewed partition handled - // in 'OptimizeSkewedPartitions' rule, should be omitted in 'ReduceNumShufflePartitions'. + // in 'OptimizeSkewedJoin' rule, should be omitted in 'ReduceNumShufflePartitions'. OptimizeSkewedJoin(conf), ReduceNumShufflePartitions(conf), // The rule of 'OptimizeLocalShuffleReader' need to make use of the 'partitionStartIndices' 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 75d4184a2c14e..ad828006e3315 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 @@ -28,12 +28,14 @@ import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.exchange.{EnsureRequirements, ShuffleExchangeExec} import org.apache.spark.sql.execution.joins.SortMergeJoinExec import org.apache.spark.sql.internal.SQLConf case class OptimizeSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] { + private val ensureRequirements = EnsureRequirements(conf) + private val supportedJoinTypes = Inner :: Cross :: LeftSemi :: LeftAnti :: LeftOuter :: RightOuter :: Nil @@ -54,7 +56,11 @@ case class OptimizeSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] { private def medianSize(stats: MapOutputStatistics): Long = { val numPartitions = stats.bytesByPartitionId.length val bytes = stats.bytesByPartitionId.sorted - if (bytes(numPartitions / 2) > 0) bytes(numPartitions / 2) else 1 + numPartitions match { + case _ if (numPartitions % 2 == 0) => + math.max((bytes(numPartitions / 2) + bytes(numPartitions / 2 - 1)) / 2, 1) + case _ => math.max(bytes(numPartitions / 2), 1) + } } /** @@ -76,18 +82,19 @@ case class OptimizeSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] { val avgPartitionSize = mapPartitionSizes.sum / maxSplits val advisoryPartitionSize = math.max(avgPartitionSize, conf.getConf(SQLConf.ADAPTIVE_EXECUTION_SKEWED_PARTITION_SIZE_THRESHOLD)) - val partitionIndices = mapPartitionSizes.indices val partitionStartIndices = ArrayBuffer[Int]() - var postMapPartitionSize = mapPartitionSizes(0) partitionStartIndices += 0 - partitionIndices.drop(1).foreach { nextPartitionIndex => - val nextMapPartitionSize = mapPartitionSizes(nextPartitionIndex) - if (postMapPartitionSize + nextMapPartitionSize > advisoryPartitionSize) { - partitionStartIndices += nextPartitionIndex + var i = 0 + var postMapPartitionSize = 0L + while (i < mapPartitionSizes.length) { + val nextMapPartitionSize = mapPartitionSizes(i) + if (i > 0 && postMapPartitionSize + nextMapPartitionSize > advisoryPartitionSize) { + partitionStartIndices += i postMapPartitionSize = nextMapPartitionSize } else { postMapPartitionSize += nextMapPartitionSize } + i += 1 } if (partitionStartIndices.size > maxSplits) { @@ -114,25 +121,56 @@ case class OptimizeSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] { stage.shuffle.shuffleDependency.rdd.partitions.length } - def handleSkewJoin(plan: SparkPlan): SparkPlan = plan.transformUp { - case smj @ SortMergeJoinExec(leftKeys, rightKeys, joinType, condition, - s1 @ SortExec(_, _, left: ShuffleQueryStageExec, _), - s2 @ SortExec(_, _, right: ShuffleQueryStageExec, _)) - if supportedJoinTypes.contains(joinType) => + private def getShuffleQueryStage(plan : SparkPlan): Option[ShuffleQueryStageExec] = + plan match { + case stage: ShuffleQueryStageExec => Some(stage) + case SortExec(_, _, s: ShuffleQueryStageExec, _) => + Some(s) + case _ => None + } + + private def reOptimizeChild( + skewedReader: SkewedPartitionReaderExec, + child: SparkPlan): SparkPlan = child match { + case sort @ SortExec(_, _, s: ShuffleQueryStageExec, _) => + sort.copy(child = skewedReader) + case _: ShuffleQueryStageExec => skewedReader + } + + private def getSizeInfo(medianSize: Long, maxSize: Long): String = { + s"median size: $medianSize, max size: ${maxSize}" + } + + /* + * This method aim to optimize the skewed join with the following steps: + * 1. Check whether the shuffle partition is skewed based on the median size + * and the skewed partition threshold in origin smj. + * 2. Assuming partition0 is skewed in left side, and it has 5 mappers (Map0, Map1...Map4). + * And we will split the 5 Mappers into 3 mapper ranges [(Map0, Map1), (Map2, Map3), (Map4)] + * based on the map size and the max split number. + * 3. Create the 3 smjs with separately reading the above mapper ranges and then join with + * the Partition0 in right side. + * 4. Finally union the above 3 split smjs and the origin smj. + */ + def optimizeSkewJoin(plan: SparkPlan): SparkPlan = plan.transformUp { + case smj @ SortMergeJoinExec(leftKeys, rightKeys, joinType, condition, leftPlan, rightPlan) + if (getShuffleQueryStage(leftPlan).nonEmpty && getShuffleQueryStage(rightPlan).nonEmpty) && + supportedJoinTypes.contains(joinType) => + val left = getShuffleQueryStage(leftPlan).get + val right = getShuffleQueryStage(rightPlan).get val leftStats = getStatistics(left) val rightStats = getStatistics(right) val numPartitions = leftStats.bytesByPartitionId.length val leftMedSize = medianSize(leftStats) val rightMedSize = medianSize(rightStats) - val leftSizeInfo = s"median size: $leftMedSize, max size: ${leftStats.bytesByPartitionId.max}" - val rightSizeInfo = s"median size: $rightMedSize," + - s" max size: ${rightStats.bytesByPartitionId.max}" logDebug( s""" |Try to optimize skewed join. - |Left side partition size: $leftSizeInfo - |Right side partition size: $rightSizeInfo + |Left side partition size: + |${getSizeInfo(leftMedSize, leftStats.bytesByPartitionId.max)} + |Right side partition size: + |${getSizeInfo(rightMedSize, rightStats.bytesByPartitionId.max)} """.stripMargin) val skewedPartitions = mutable.HashSet[Int]() @@ -171,16 +209,19 @@ case class OptimizeSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] { left, partitionId, leftMapIdStartIndices(i), leftEndMapId) val rightSkewedReader = SkewedPartitionReaderExec(right, partitionId, rightMapIdStartIndices(j), rightEndMapId) + val skewedLeft = reOptimizeChild(leftSkewedReader, leftPlan) + val skewedRight = reOptimizeChild(rightSkewedReader, rightPlan) subJoins += SortMergeJoinExec(leftKeys, rightKeys, joinType, condition, - s1.copy(child = leftSkewedReader), s2.copy(child = rightSkewedReader)) + skewedLeft, skewedRight) } } } logDebug(s"number of skewed partitions is ${skewedPartitions.size}") if (skewedPartitions.nonEmpty) { - val optimizedSmj = smj.transformDown { - case sort @ SortExec(_, _, shuffleStage: ShuffleQueryStageExec, _) => - sort.copy(child = PartialShuffleReaderExec(shuffleStage, skewedPartitions.toSet)) + val optimizedSmj = smj.transformUp { + case shuffleStage: ShuffleQueryStageExec if shuffleStage.id == left.id || + shuffleStage.id == right.id => + PartialShuffleReaderExec(shuffleStage, skewedPartitions.toSet) } subJoins += optimizedSmj UnionExec(subJoins) @@ -204,11 +245,27 @@ case class OptimizeSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] { val shuffleStages = collectShuffleStages(plan) if (shuffleStages.length == 2) { - // Currently we only support handling skewed join for 2 table join. - handleSkewJoin(plan) + // When multi table join, there will be too many complex combination to consider. + // Currently we only handle 2 table join like following two use cases. + // SMJ SMJ + // Sort Shuffle + // Shuffle or Shuffle + // Sort + // Shuffle + val optimizePlan = optimizeSkewJoin(plan) + val numShuffles = ensureRequirements.apply(optimizePlan).collect { + case e: ShuffleExchangeExec => e + }.length + + if (numShuffles > 0) { + logDebug("OptimizeSkewedJoin rule is not applied due" + + " to additional shuffles will be introduced.") + plan + } else { + optimizePlan + } } else { plan - } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala index f327e84563da9..985e1db2736fc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.Block._ import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{PartialShuffleReaderExec, SkewedPartitionReaderExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.util.collection.BitSet @@ -95,8 +96,20 @@ case class SortMergeJoinExec( s"${getClass.getSimpleName} should not take $x as the JoinType") } - override def requiredChildDistribution: Seq[Distribution] = - HashClusteredDistribution(leftKeys) :: HashClusteredDistribution(rightKeys) :: Nil + private def containSkewedReader(plan: SparkPlan): Boolean = plan match { + case s: SkewedPartitionReaderExec => true + case p: PartialShuffleReaderExec => true + case s: SortExec => containSkewedReader(s.child) + case _ => false + } + + override def requiredChildDistribution: Seq[Distribution] = { + if (containSkewedReader(left)) { + UnspecifiedDistribution :: UnspecifiedDistribution :: Nil + } else { + HashClusteredDistribution(leftKeys) :: HashClusteredDistribution(rightKeys) :: Nil + } + } override def outputOrdering: Seq[SortOrder] = joinType match { // For inner join, orders of both sides keys should be kept. 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 c2daae071afc5..78a1183664749 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 @@ -579,6 +579,32 @@ class AdaptiveQueryExecSuite } } + test("SPARK-30524: Do not optimize skew join if introduce additional shuffle") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_SKEWED_PARTITION_SIZE_THRESHOLD.key -> "100", + SQLConf.SHUFFLE_TARGET_POSTSHUFFLE_INPUT_SIZE.key -> "700") { + withTempView("skewData1", "skewData2") { + spark + .range(0, 1000, 1, 10) + .selectExpr("id % 2 as key1", "id as value1") + .createOrReplaceTempView("skewData1") + spark + .range(0, 1000, 1, 10) + .selectExpr("id % 1 as key2", "id as value2") + .createOrReplaceTempView("skewData2") + val (innerPlan, innerAdaptivePlan) = runAdaptiveAndVerifyResult( + "SELECT key1 FROM skewData1 join skewData2 ON key1 = key2 group by key1") + val innerSmj = findTopLevelSortMergeJoin(innerPlan) + assert(innerSmj.size == 1) + // Additional shuffle introduced, so disable the "OptimizeSkewedJoin" optimization + val innerSmjAfter = findTopLevelSortMergeJoin(innerAdaptivePlan) + assert(innerSmjAfter.size == 1) + } + } + } + test("SPARK-29544: adaptive skew join with different join types") { Seq("false", "true").foreach { reducePostShufflePartitionsEnabled => withSQLConf(