[SPARK-28356][SQL] Do not reduce the number of partitions for repartition in adaptive execution

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

@@ -312,6 +312,16 @@ object SQLConf {
     .booleanConf
     .createWithDefault(true)
 
+  val REDUCE_POST_SHUFFLE_PARTITIONS_FOR_REPARTITION =
+    buildConf("spark.sql.adaptive.reducePostShufflePartitionsForRepartition")
+    .doc("When true and reducePostShufflePartitions in adaptive execution is enabled, it also " +
+      "reduces the number of post-shuffle partitions for shuffles introduced by repartition. " +
+      "By default this is false, so users get the desired partition number by calling " +
+      "repartition even in adaptive execution.")
+    .internal()
+    .booleanConf
+    .createWithDefault(false)
+
   val SHUFFLE_MIN_NUM_POSTSHUFFLE_PARTITIONS =
     buildConf("spark.sql.adaptive.minNumPostShufflePartitions")
       .doc("The advisory minimum number of post-shuffle partitions used in adaptive execution.")
@@ -1951,6 +1961,9 @@ class SQLConf extends Serializable with Logging {
 
   def reducePostShufflePartitionsEnabled: Boolean = getConf(REDUCE_POST_SHUFFLE_PARTITIONS_ENABLED)
 
+  def reducePostShufflePartitionsForRepartition: Boolean =
+    getConf(REDUCE_POST_SHUFFLE_PARTITIONS_FOR_REPARTITION)
+
   def minNumPostShufflePartitions: Int =
     getConf(SHUFFLE_MIN_NUM_POSTSHUFFLE_PARTITIONS)
 

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

@@ -703,7 +703,9 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
 
       case logical.Repartition(numPartitions, shuffle, child) =>
         if (shuffle) {
-          ShuffleExchangeExec(RoundRobinPartitioning(numPartitions), planLater(child)) :: Nil
+          val supportAdaptive = conf.reducePostShufflePartitionsForRepartition
+          ShuffleExchangeExec(RoundRobinPartitioning(numPartitions),
+            planLater(child), supportAdaptive) :: Nil
         } else {
           execution.CoalesceExec(numPartitions, planLater(child)) :: Nil
         }
@@ -736,7 +738,8 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
       case r: logical.Range =>
         execution.RangeExec(r) :: Nil
       case r: logical.RepartitionByExpression =>
-        exchange.ShuffleExchangeExec(r.partitioning, planLater(r.child)) :: Nil
+        val supportAdaptive = conf.reducePostShufflePartitionsForRepartition
+        exchange.ShuffleExchangeExec(r.partitioning, planLater(r.child), supportAdaptive) :: Nil
       case ExternalRDD(outputObjAttr, rdd) => ExternalRDDScanExec(outputObjAttr, rdd) :: Nil
       case r: LogicalRDD =>
         RDDScanExec(r.output, r.rdd, "ExistingRDD", r.outputPartitioning, r.outputOrdering) :: Nil

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

@@ -140,6 +140,8 @@ case class ShuffleQueryStageExec(
       case _ =>
     }
   }
+
+  def supportAdaptive: Boolean = plan.supportAdaptive
 }
 
 /**

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

@@ -61,12 +61,16 @@ case class ReduceNumShufflePartitions(conf: SQLConf) extends Rule[SparkPlan] {
       // If not all leaf nodes are query stages, it's not safe to reduce the number of
       // shuffle partitions, because we may break the assumption that all children of a spark plan
       // have same number of output partitions.
+      return plan
+    }
+
+    val shuffleStages = plan.collect {
+      case stage: ShuffleQueryStageExec => stage
+      case ReusedQueryStageExec(_, stage: ShuffleQueryStageExec, _) => stage
+    }
+    if (!shuffleStages.forall(_.supportAdaptive)) {
       plan
     } else {
-      val shuffleStages = plan.collect {
-        case stage: ShuffleQueryStageExec => stage
-        case ReusedQueryStageExec(_, stage: ShuffleQueryStageExec, _) => stage
-      }
       val shuffleMetrics = shuffleStages.map { stage =>
         val metricsFuture = stage.mapOutputStatisticsFuture
         assert(metricsFuture.isCompleted, "ShuffleQueryStageExec should already be ready")

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala

@@ -94,7 +94,7 @@ case class EnsureRequirements(conf: SQLConf) extends Rule[SparkPlan] {
             val defaultPartitioning = distribution.createPartitioning(targetNumPartitions)
             child match {
               // If child is an exchange, we replace it with a new one having defaultPartitioning.
-              case ShuffleExchangeExec(_, c) => ShuffleExchangeExec(defaultPartitioning, c)
+              case ShuffleExchangeExec(_, c, _) => ShuffleExchangeExec(defaultPartitioning, c)
               case _ => ShuffleExchangeExec(defaultPartitioning, child)
             }
           }
@@ -191,7 +191,7 @@ case class EnsureRequirements(conf: SQLConf) extends Rule[SparkPlan] {
 
   def apply(plan: SparkPlan): SparkPlan = plan.transformUp {
     // TODO: remove this after we create a physical operator for `RepartitionByExpression`.
-    case operator @ ShuffleExchangeExec(upper: HashPartitioning, child) =>
+    case operator @ ShuffleExchangeExec(upper: HashPartitioning, child, _) =>
       child.outputPartitioning match {
         case lower: HashPartitioning if upper.semanticEquals(lower) => child
         case _ => operator

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala

@@ -43,7 +43,8 @@ import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, RecordCo
  */
 case class ShuffleExchangeExec(
     override val outputPartitioning: Partitioning,
-    child: SparkPlan) extends Exchange {
+    child: SparkPlan,
+    supportAdaptive: Boolean = true) extends Exchange {
 
   // NOTE: coordinator can be null after serialization/deserialization,
   //       e.g. it can be null on the Executor side

sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala

@@ -1365,7 +1365,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
         val agg = cp.groupBy('id % 2).agg(count('id))
 
         agg.queryExecution.executedPlan.collectFirst {
-          case ShuffleExchangeExec(_, _: RDDScanExec) =>
+          case ShuffleExchangeExec(_, _: RDDScanExec, _) =>
           case BroadcastExchangeExec(_, _: RDDScanExec) =>
         }.foreach { _ =>
           fail(

sql/core/src/test/scala/org/apache/spark/sql/execution/ReduceNumShufflePartitionsSuite.scala

@@ -576,6 +576,7 @@ class ReduceNumShufflePartitionsSuite extends SparkFunSuite with BeforeAndAfterA
 
   test("Union two datasets with different pre-shuffle partition number") {
     val test: SparkSession => Unit = { spark: SparkSession =>
+      spark.conf.set(SQLConf.REDUCE_POST_SHUFFLE_PARTITIONS_FOR_REPARTITION.key, "true")
       val dataset1 = spark.range(3)
       val dataset2 = spark.range(3)
 
@@ -592,4 +593,40 @@ class ReduceNumShufflePartitionsSuite extends SparkFunSuite with BeforeAndAfterA
     }
     withSparkSession(test, 100, None)
   }
+
+  test("Do not reduce the number of shuffle partition for repartition") {
+    val test: SparkSession => Unit = { spark: SparkSession =>
+      spark.conf.set(SQLConf.REDUCE_POST_SHUFFLE_PARTITIONS_FOR_REPARTITION.key, "false")
+      val ds = spark.range(3)
+      val resultDf = ds.repartition(2, ds.col("id")).toDF()
+
+      checkAnswer(resultDf,
+        Seq((0), (1), (2)).map(i => Row(i)))
+      val finalPlan = resultDf.queryExecution.executedPlan
+        .asInstanceOf[AdaptiveSparkPlanExec].executedPlan
+      assert(finalPlan.collect { case p: CoalescedShuffleReaderExec => p }.length == 0)
+    }
+    withSparkSession(test, 200, None)
+  }
+
+  test("Reduce the number of shuffle partition for repartition") {
+    val test: SparkSession => Unit = { spark: SparkSession =>
+      spark.conf.set(SQLConf.REDUCE_POST_SHUFFLE_PARTITIONS_FOR_REPARTITION.key, "true")
+      val ds = spark.range(3)
+      val resultDf = ds.repartition(2, ds.col("id")).toDF()
+
+      checkAnswer(resultDf,
+        Seq((0), (1), (2)).map(i => Row(i)))
+      val finalPlan = resultDf.queryExecution.executedPlan
+        .asInstanceOf[AdaptiveSparkPlanExec].executedPlan
+      val shuffleReaders = finalPlan.collect {
+        case reader: CoalescedShuffleReaderExec => reader
+      }
+      assert(shuffleReaders.length === 1)
+      shuffleReaders.foreach { reader =>
+        assert(reader.outputPartitioning.numPartitions === 1)
+      }
+    }
+    withSparkSession(test, 200, None)
+  }
 }