[CARMEL-6383] Allow introduce new shuffle in skew handling (#1161)

* [CARMEL-6383] Allow introduce new shuffle in skew handling

* fix code style

* Support multi skew joins in EliminateSkewOptimzeIntroducedShuffle

* fix ut

* fix ut

* Add tag to detect newly introduced shuffle

* Add tags back when updateShuffleReads

Author: xingchaozh

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -696,6 +696,12 @@ object SQLConf {
       .booleanConf
       .createWithDefault(false)
 
+  val ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN =
+    buildConf("spark.sql.adaptive.forceOptimizeSkewedJoin")
+      .doc("When true, force enable OptimizeSkewedJoin even if it introduces extra shuffle.")
+      .version("3.3.0")
+      .fallbackConf(ALLOW_ADDITIONAL_SHUFFLE)
+
   val SKEW_JOIN_SKEWED_PARTITION_FACTOR =
     buildConf("spark.sql.adaptive.skewJoin.skewedPartitionFactor")
       .doc("A partition is considered as skewed if its size is larger than this factor " +

sql/core/src/main/scala/org/apache/spark/sql/execution/RemoveRedundantSorts.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution
 
 import org.apache.spark.sql.catalyst.expressions.SortOrder
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.window.WindowExec
 import org.apache.spark.sql.internal.SQLConf
 
 /**
@@ -42,5 +43,9 @@ object RemoveRedundantSorts extends Rule[SparkPlan] {
         if SortOrder.orderingSatisfies(child.outputOrdering, orders) &&
           child.outputPartitioning.satisfies(s.requiredChildDistribution.head) =>
       child
+    case w @ WindowExec(_, _, _, _, s1 @ WindowSortLimitExec(
+    _, _, _, _, _, _, s2: WindowSortLimitExec, _))
+      if conf.enableWindowLimit && s2.redundantWith(s1) =>
+      w.copy(child = s2)
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala

@@ -184,6 +184,28 @@ case class WindowSortLimitExec(
   override protected def doProduce(ctx: CodegenContext): String = {
     doProduce(ctx, classOf[AbstractUnsafeExternalRowSorter].getName)
   }
+
+  def redundantWith(o: WindowSortLimitExec): Boolean = {
+    val partitionSpecSame = partitionSpec.length == o.partitionSpec.length &&
+      partitionSpec.zip(o.partitionSpec).forall {
+      case (l, r) => l.semanticEquals(r)
+    }
+
+    val sortOrderInWindowSame = sortOrderInWindow.length == o.sortOrderInWindow.length &&
+      sortOrderInWindow.zip(o.sortOrderInWindow).forall {
+      case (l, r) => l.semanticEquals(r)
+    }
+
+    val sortOrderAcrossWindowsSame =
+      sortOrderAcrossWindows.length == o.sortOrderAcrossWindows.length &&
+      sortOrderAcrossWindows.zip(o.sortOrderAcrossWindows).forall {
+        case (l, r) => l.semanticEquals(r)
+      }
+
+    partitionSpecSame && sortOrderInWindowSame && sortOrderAcrossWindowsSame &&
+      global == o.global && partitionLimit == o.partitionLimit &&
+      maxBufferSize == o.maxBufferSize && testSpillFrequency == o.testSpillFrequency
+  }
 }
 
 abstract class SortExecBase(

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala

@@ -97,7 +97,10 @@ case class AdaptiveSparkPlanExec(
     EnsureRepartitionForWriting,
     EliminateShuffleExec,
     DisableUnnecessaryBucketedScan,
-    AdjustScanPartitionSizeDynamically
+    AdjustScanPartitionSizeDynamically,
+    OptimizeSkewedJoin, // ensureRequirements
+    EliminateSkewOptimzeIntroducedShuffle,
+    removeRedundantSorts
   ) ++ context.session.sessionState.queryStagePrepRules
 
   @transient private val initialPlan = context.session.withActive {
@@ -109,16 +112,15 @@ case class AdaptiveSparkPlanExec(
   @transient private val queryStageOptimizerRules: Seq[Rule[SparkPlan]] = Seq(
     PlanAdaptiveDynamicPruningFilters(initialPlan),
     ReuseAdaptiveSubquery(context.subqueryCache),
-
-    OptimizeSkewedJoin,
     CoalesceShufflePartitions(context.session),
     // The following two rules need to make use of 'CustomShuffleReaderExec.partitionSpecs'
     // added by `CoalesceShufflePartitions`. So they must be executed after it.
     OptimizeSkewedRangePartition,
     OptimizeSkewedInsert,
     OptimizeLocalShuffleReader,
     ensureRequirements,
-    EliminateSkewOptimzeIntroducedShuffle
+    EliminateSkewOptimzeIntroducedShuffle,
+    removeRedundantSorts
   )
 
   // A list of physical optimizer rules to be applied right after a new stage is created. The input
@@ -128,7 +130,8 @@ case class AdaptiveSparkPlanExec(
     CollapseCodegenStages()
   )
 
-  @transient private val costEvaluator = SimpleCostEvaluator
+  @transient private val costEvaluator =
+    SimpleCostEvaluator(conf.getConf(SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN))
 
   @volatile private[sql] var currentPhysicalPlan = initialPlan
 

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala

@@ -231,7 +231,7 @@ case class CoalesceShufflePartitions(session: SparkSession)
   }
 
   private def collectShuffleStageInfos(plan: SparkPlan): Seq[ShuffleStageInformation] = plan match {
-    case ShuffleStageInformation(stage, specs) => Seq(new ShuffleStageInformation(stage, specs))
+    case ShuffleStageInformation(stage, specs, _) => Seq(new ShuffleStageInformation(stage, specs))
     case _ => plan.children.flatMap(collectShuffleStageInfos)
   }
 
@@ -242,7 +242,7 @@ case class CoalesceShufflePartitions(session: SparkSession)
     // Even for shuffle exchange whose input RDD has 0 partition, we should still update its
     // `partitionStartIndices`, so that all the leaf shuffles in a stage have the same
     // number of output partitions.
-    case ShuffleStageInformation(stage, _) =>
+    case ShuffleStageInformation(stage, _, optimizeTags) =>
       specsMap.get(stage.id).map { specs =>
         // Since we may not submit tasks which read empty partition, and the MapOutputStatistics
         // above could be all '0' of each partition size. And the coalesced partitionSpecs could
@@ -251,7 +251,9 @@ case class CoalesceShufflePartitions(session: SparkSession)
         val normalizedSpecs = if (specs.isEmpty) {
           specs :+ CoalescedPartitionSpec(0, numPartitions)
         } else specs
-        CustomShuffleReaderExec(stage, normalizedSpecs)
+        val newCustomShuffleReaderExec = CustomShuffleReaderExec(stage, normalizedSpecs)
+        optimizeTags.foreach(newCustomShuffleReaderExec.addOptimizeTag(_))
+        newCustomShuffleReaderExec
       }.getOrElse(plan)
     case other => other.mapChildren(updateShuffleReads(_, specsMap, numPartitions))
   }
@@ -262,11 +264,12 @@ private class ShuffleStageInformation(val shuffleStage: ShuffleQueryStageExec,
 
 private object ShuffleStageInformation {
   def unapply(plan: SparkPlan)
-  : Option[(ShuffleQueryStageExec, Option[Seq[ShufflePartitionSpec]])] = plan match {
+  : Option[(ShuffleQueryStageExec,
+    Option[Seq[ShufflePartitionSpec]], Seq[String])] = plan match {
     case stage: ShuffleQueryStageExec =>
-      Some((stage, None))
-    case CustomShuffleReaderExec(s: ShuffleQueryStageExec, partitionSpecs) =>
-      Some((s, Some(partitionSpecs)))
+      Some((stage, None, Seq.empty[String]))
+    case c@CustomShuffleReaderExec(s: ShuffleQueryStageExec, partitionSpecs) =>
+      Some((s, Some(partitionSpecs), c.getOptimizeTags()))
     case _ => None
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/EliminateSkewOptimzeIntroducedShuffle.scala

@@ -18,7 +18,7 @@ package org.apache.spark.sql.execution.adaptive
 
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.execution.{CoalescedPartitionSpec, SparkPlan}
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.execution.joins.SortMergeJoinExec
 
 /**
@@ -43,8 +43,8 @@ private[adaptive] object EliminateSkewOptimzeIntroducedShuffle extends Rule[Spar
     val skewedJoins = plan.collect {
       case s: SortMergeJoinExec if s.isSkewJoin => s
     }
-    // Sanity check. Suppose this rule should take effect iff there is one skewed join operator.
-    if (skewedJoins.size != 1) {
+    // Sanity check. Suppose this rule should take effect if at least one skewed join operator.
+    if (skewedJoins.isEmpty) {
       plan
     } else {
       plan.transformUp {
@@ -60,8 +60,8 @@ private[adaptive] object EliminateSkewOptimzeIntroducedShuffle extends Rule[Spar
     }
   }
 
-  private def existSkewedReader(plan: SparkPlan): Boolean =
-    plan.find {
+  private def existSkewedReader(child: SparkPlan): Boolean =
+    child.find {
       // All CoalescedPartitionSpec and there are some duplicate specs.
       case CustomShuffleReaderExec(s: ShuffleQueryStageExec, partitions) =>
         partitions.forall(_.isInstanceOf[CoalescedPartitionSpec]) &&

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeLocalShuffleReader.scala

@@ -114,12 +114,12 @@ object OptimizeLocalShuffleReader extends Rule[SparkPlan] {
       s match {
         case proj1 @ ProjectExec(_, sort@SortExec(_, true, proj2@ProjectExec(_,
         c@CustomShuffleReaderExec(_: ShuffleQueryStageExec, _)), _)) =>
-          proj1.withNewChildren(Seq(sort.withNewChildren(
-            Seq(proj2.withNewChildren(Seq(createLocalReader(c)))))))
+          proj1.withNewChildren(sort.withNewChildren(
+            proj2.withNewChildren(createLocalReader(c) :: Nil) :: Nil) :: Nil)
         case limit @ LocalLimitExec(_, proj1@ProjectExec(_, sort@SortExec(_, true,
         proj2@ProjectExec(_, c@CustomShuffleReaderExec(_: ShuffleQueryStageExec, _)), _))) =>
-          limit.withNewChildren(Seq(proj1.withNewChildren(Seq(sort.withNewChildren(
-            Seq(proj2.withNewChildren(Seq(createLocalReader(c)))))))))
+          limit.withNewChildren(proj1.withNewChildren(sort.withNewChildren(
+            proj2.withNewChildren(createLocalReader(c) :: Nil) :: Nil) :: Nil) :: Nil)
         case _ => createLocalReader(s)
       }
     case s: SparkPlan =>
@@ -180,13 +180,15 @@ object OptimizeLocalShuffleReader extends Rule[SparkPlan] {
         case _ => false
       }
     case ProjectExec(_, sort @ SortExec(_, true, ProjectExec(_, CustomShuffleReaderExec(
-    _: ShuffleQueryStageExec, _)), _))
+    s: ShuffleQueryStageExec, _)), _))
       if conf.getConf(SQLConf.OPTIMIZE_RANGE_PARTITION_SKEW_ENABLED) =>
+      supportLocalReader(s.shuffle) &&
         sort.getTagValue(OptimizeSkewedRangePartition.SKEWED_RANGE_PARTITION_TAG).nonEmpty
     case LocalLimitExec(_, ProjectExec(_, sort@SortExec(_, true, ProjectExec(_,
-    CustomShuffleReaderExec(_: ShuffleQueryStageExec, _)), _)))
+    CustomShuffleReaderExec(s: ShuffleQueryStageExec, _)), _)))
       if conf.getConf(SQLConf.OPTIMIZE_RANGE_PARTITION_SKEW_ENABLED) =>
-      sort.getTagValue(OptimizeSkewedRangePartition.SKEWED_RANGE_PARTITION_TAG).nonEmpty
+      supportLocalReader(s.shuffle) &&
+        sort.getTagValue(OptimizeSkewedRangePartition.SKEWED_RANGE_PARTITION_TAG).nonEmpty
     case _ => false
   }