[SPARK-30524] [SQL] Disable OptimizeSkewedJoin rule when introducing additional shuffle

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'

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
-
     }
   }
 }

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.

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(