fix a bug that no support a union node with differing number of partitions if we explicitly repartition them apache#98

Author: weixiuli

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

@@ -85,6 +85,33 @@ abstract class QueryStage extends UnaryExecNode {
       Future.sequence(shuffleStageFutures)(implicitly, QueryStage.executionContext), Duration.Inf)
   }
 
+  def getSupportAdaptiveFlag(queryStageInputs: Seq[ShuffleQueryStageInput]): Boolean = {
+    val queryStageInputsNumPartitions = queryStageInputs.map {
+      _.outputPartitioning match {
+        case hash: HashPartitioning => hash.numPartitions
+        case collection: PartitioningCollection =>
+          val PartitioningCollectionNumPartitions = collection.partitionings.map {
+            partitioning => {
+              if (partitioning.isInstanceOf[HashPartitioning]) {
+                partitioning.numPartitions
+              } else {
+                -1
+              }
+            }
+          }.distinct
+          if (PartitioningCollectionNumPartitions.length > 1) {
+            -1
+          } else {
+            PartitioningCollectionNumPartitions.head
+          }
+        case _ => -1
+      }
+    }.distinct
+    val supportAdaptiveFlag = (queryStageInputsNumPartitions.length == 1
+        && queryStageInputsNumPartitions.head != -1)
+    supportAdaptiveFlag
+  }
+
   private var prepared = false
 
   /**
@@ -127,14 +154,7 @@ abstract class QueryStage extends UnaryExecNode {
       val childMapOutputStatistics = queryStageInputs.map(_.childStage.mapOutputStatistics)
         .filter(_ != null).toArray
       // Right now, Adaptive execution only support HashPartitionings.
-      val supportAdaptive = queryStageInputs.forall {
-        _.outputPartitioning match {
-          case hash: HashPartitioning => true
-          case collection: PartitioningCollection =>
-            collection.partitionings.forall(_.isInstanceOf[HashPartitioning])
-          case _ => false
-        }
-      }
+      val supportAdaptive = getSupportAdaptiveFlag(queryStageInputs)
 
       if (childMapOutputStatistics.length > 0 && supportAdaptive) {
         val exchangeCoordinator = new ExchangeCoordinator(

sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/QueryStageSuite.scala

@@ -1001,4 +1001,29 @@ class QueryStageSuite extends SparkFunSuite with BeforeAndAfterAll {
       " union all select count(test.age) from test"),
       Row(1) :: Row(1) :: Row(2) :: Nil)
   }
+
+  test("different pre-shuffle partition number of datasets to union with adaptive") {
+    val sparkSession = defaultSparkSession
+    val dataset1 = sparkSession.range(1000)
+    val dataset2 = sparkSession.range(1001)
+
+    val compute = dataset1.repartition(505, dataset1.col("id"))
+      .union(dataset2.repartition(105, dataset2.col("id")))
+
+    assert(compute.orderBy("id").toDF("id").takeAsList(10).toArray
+      === Seq((0), (0), (1), (1), (2), (2), (3), (3), (4), (4)).map(i => Row(i)).toArray)
+    compute.explain()
+  }
+
+  test("different pre-shuffle partition number of datasets to join with adaptive") {
+    val sparkSession = defaultSparkSession
+    val dataset1 = sparkSession.range(1000)
+    val dataset2 = sparkSession.range(1001)
+    val compute = dataset1.repartition(105).toDF("key1")
+      .join(dataset1.repartition(505).toDF("key2"), col("key1") === col("key2"), "left")
+    assert(compute.orderBy("key1").toDF("key1","key2").select("key1").takeAsList(10).toArray
+      === Seq((0), (1), (2), (3), (4), (5), (6), (7), (8), (9)).map(i => Row(i)).toArray)
+    compute.explain()
+  }
+
 }