[SPARK-40703][SQL] Introduce shuffle on SinglePartition to improve pa…

…rallelism

### What changes were proposed in this pull request?

This PR fixes a performance regression issue when one side of a join uses `HashPartitioning` with `ShuffleExchange` while the other side uses `SinglePartition`. In this case, Spark will re-shuffle the side with `HashPartitioning` and both sides will end up with only a single partition. This could hurt query performance a lot if the side with `HashPartitioning` contains a lot of input data.

### Why are the changes needed?

After SPARK-35703, when Spark sees that one side of the join has `ShuffleExchange` (meaning it needs to be shuffled anyways), and the other side doesn't, it'll try to avoid shuffling the side without `ShuffleExchange`. For instance:

```
ShuffleExchange(HashPartition(200)) <-> HashPartition(150)
```

will be converted into
```
ShuffleExchange(HashPartition(150)) <-> HashPartition(150)
```

However, when the side without `ShuffleExchange` is `SinglePartition`, like the following:
```
ShuffleExchange(HashPartition(150)) <-> SinglePartition
```

Spark will also do the same which causes the left-hand side to only use one partition. This can hurt job parallelism dramatically, especially when using DataSource V2, since `SinglePartition` is used by the V2 scan. On the other hand, it seems DataSource V1 won't be impacted much as it always report `UnknownPartitioning` in this situation.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Added new unit tests in `EnsureRequirementsSuite`.

Closes #38196 from sunchao/SPARK-40703.

Authored-by: Chao Sun <sunchao@apple.com>
Signed-off-by: Yuming Wang <yumwang@ebay.com>
(cherry picked from commit bde6423)
Signed-off-by: Yuming Wang <yumwang@ebay.com>

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala

@@ -529,7 +529,6 @@ trait ShuffleSpec {
    * clustering expressions.
    *
    * This will only be called when:
-   *  - [[canCreatePartitioning]] returns true.
    *  - [[isCompatibleWith]] returns false on the side where the `clustering` is from.
    */
   def createPartitioning(clustering: Seq[Expression]): Partitioning =
@@ -542,7 +541,7 @@ case object SinglePartitionShuffleSpec extends ShuffleSpec {
     other.numPartitions == 1
   }
 
-  override def canCreatePartitioning: Boolean = true
+  override def canCreatePartitioning: Boolean = false
 
   override def createPartitioning(clustering: Seq[Expression]): Partitioning =
     SinglePartition

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala

@@ -367,7 +367,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
       assert(HashShuffleSpec(HashPartitioning(Seq($"a", $"b"), 10), distribution)
         .canCreatePartitioning)
     }
-    assert(SinglePartitionShuffleSpec.canCreatePartitioning)
+    assert(!SinglePartitionShuffleSpec.canCreatePartitioning)
     withSQLConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
       assert(ShuffleSpecCollection(Seq(
         HashShuffleSpec(HashPartitioning(Seq($"a"), 10), distribution),

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

@@ -73,9 +73,13 @@ case class EnsureRequirements(
       case _ => false
     }.map(_._2)
 
+    // Special case: if all sides of the join are single partition
+    val allSinglePartition =
+      childrenIndexes.forall(children(_).outputPartitioning == SinglePartition)
+
     // If there are more than one children, we'll need to check partitioning & distribution of them
     // and see if extra shuffles are necessary.
-    if (childrenIndexes.length > 1) {
+    if (childrenIndexes.length > 1 && !allSinglePartition) {
       val specs = childrenIndexes.map(i => {
         val requiredDist = requiredChildDistributions(i)
         assert(requiredDist.isInstanceOf[ClusteredDistribution],

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

@@ -262,7 +262,7 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
           val numExchanges = collect(plan) {
             case exchange: ShuffleExchangeExec => exchange
           }.length
-          assert(numExchanges === 3)
+          assert(numExchanges === 5)
         }
 
         {
@@ -278,7 +278,7 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
           val numExchanges = collect(plan) {
             case exchange: ShuffleExchangeExec => exchange
           }.length
-          assert(numExchanges === 3)
+          assert(numExchanges === 5)
         }
 
       }

sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala

@@ -433,8 +433,10 @@ class EnsureRequirementsSuite extends SharedSparkSession {
         exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2)
       EnsureRequirements.apply(smjExec) match {
         case SortMergeJoinExec(_, _, _, _,
-        SortExec(_, _, DummySparkPlan(_, _, SinglePartition, _, _), _),
-        SortExec(_, _, ShuffleExchangeExec(SinglePartition, _, _), _), _) =>
+        SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _), _),
+        SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _), _), _) =>
+          assert(left.numPartitions == 5)
+          assert(right.numPartitions == 5)
         case other => fail(other.toString)
       }
 
@@ -690,6 +692,45 @@ class EnsureRequirementsSuite extends SharedSparkSession {
     }
   }
 
+  test("SPARK-40703: shuffle for SinglePartitionShuffleSpec") {
+    withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> 20.toString) {
+      // We should re-shuffle the side with single partition when the other side is
+      // `HashPartitioning` with shuffle node, and respect the minimum parallelism.
+      var plan1: SparkPlan = ShuffleExchangeExec(
+        outputPartitioning = HashPartitioning(exprA :: Nil, 10),
+        DummySparkPlan())
+      var plan2 = DummySparkPlan(outputPartitioning = SinglePartition)
+      var smjExec = SortMergeJoinExec(exprA :: Nil, exprC :: Nil, Inner, None, plan1, plan2)
+      EnsureRequirements.apply(smjExec) match {
+        case SortMergeJoinExec(leftKeys, rightKeys, _, _,
+        SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _), _),
+        SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _), _), _) =>
+          assert(leftKeys === Seq(exprA))
+          assert(rightKeys === Seq(exprC))
+          assert(left.numPartitions == 20)
+          assert(right.numPartitions == 20)
+        case other => fail(other.toString)
+      }
+
+      // We should also re-shuffle the side with only a single partition even the other side does
+      // not have `ShuffleExchange`, but just `HashPartitioning`. However in this case the minimum
+      // shuffle parallelism will be ignored since we don't want to introduce extra shuffle.
+      plan1 = DummySparkPlan(
+        outputPartitioning = HashPartitioning(exprA :: Nil, 10))
+      plan2 = DummySparkPlan(outputPartitioning = SinglePartition)
+      smjExec = SortMergeJoinExec(exprA :: Nil, exprC :: Nil, Inner, None, plan1, plan2)
+      EnsureRequirements.apply(smjExec) match {
+        case SortMergeJoinExec(leftKeys, rightKeys, _, _,
+        SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _),
+        SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _), _), _) =>
+          assert(leftKeys === Seq(exprA))
+          assert(rightKeys === Seq(exprC))
+          assert(right.numPartitions == 10)
+        case other => fail(other.toString)
+      }
+    }
+  }
+
   test("Check with KeyGroupedPartitioning") {
     // simplest case: identity transforms
     var plan1 = DummySparkPlan(