Add a new parititioning CoalescedHashPartitioning

Author: LantaoJin

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

@@ -340,3 +340,28 @@ case class BroadcastPartitioning(mode: BroadcastMode) extends Partitioning {
     case _ => false
   }
 }
+
+/**
+ * With AE, multiple partitions in hash partitioned output could be coalesced
+ * to a single partition. CoalescedHashPartitioning is designed for such case.
+ */
+case class CoalescedHashPartitioning(
+    expressions: Seq[Expression],
+    numPartitions: Int)
+  extends Expression with Partitioning with Unevaluable {
+
+  override def children: Seq[Expression] = expressions
+  override def nullable: Boolean = false
+  override def dataType: DataType = IntegerType
+
+  override def satisfies0(required: Distribution): Boolean = {
+    super.satisfies0(required) || {
+      required match {
+        case ClusteredDistribution(requiredClustering, requiredNumPartitions) =>
+          expressions.forall(x => requiredClustering.exists(_.semanticEquals(x))) &&
+            (requiredNumPartitions.isEmpty || requiredNumPartitions.get == numPartitions)
+        case _ => false
+      }
+    }
+  }
+}

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

@@ -23,7 +23,7 @@ import scala.collection.mutable.ArrayBuffer
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
-import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning}
+import org.apache.spark.sql.catalyst.plans.physical.{CoalescedHashPartitioning, HashPartitioning, Partitioning, UnknownPartitioning}
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
@@ -65,6 +65,14 @@ case class CustomShuffleReaderExec private(
         case _ =>
           throw new IllegalStateException("operating on canonicalization plan")
       }
+    } else if (partitionSpecs.nonEmpty &&
+      partitionSpecs.forall(_.isInstanceOf[CoalescedPartitionSpec])) {
+      child match {
+        case ShuffleQueryStageExec(_, ShuffleExchangeExec(p: HashPartitioning, _, _)) =>
+          CoalescedHashPartitioning(p.expressions, partitionSpecs.size)
+        case _ =>
+          throw new IllegalStateException("operating on canonicalization plan")
+      }
     } else {
       UnknownPartitioning(partitionSpecs.length)
     }

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

@@ -738,6 +738,24 @@ class AdaptiveQueryExecSuite
       SQLConf.SHUFFLE_PARTITIONS.key -> "100",
       SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "800",
       SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "800") {
+
+      // SMJ
+      //   Sort
+      //     CustomShuffleReader(coalesced)
+      //       Shuffle
+      //   Sort
+      //     HashAggregate
+      //       CustomShuffleReader(coalesced)
+      //         Shuffle
+      // -->
+      // SMJ
+      //   Sort
+      //     CustomShuffleReader(coalesced and skew)
+      //       Shuffle
+      //   Sort
+      //     HashAggregate
+      //       CustomShuffleReader(coalesced)
+      //         Shuffle
       withTempView("skewData1", "skewData2") {
         spark
           .range(0, 1000, 1, 10)
@@ -747,14 +765,14 @@ class AdaptiveQueryExecSuite
               .otherwise('id).as("key1"),
             'id as "value1")
           .createOrReplaceTempView("skewData1")
-
         spark
           .range(0, 1000, 1, 10)
           .select(
             when('id < 250, 249)
               .otherwise('id).as("key2"),
             'id as "value2")
           .createOrReplaceTempView("skewData2")
+
         val sqlText =
           """
             |SELECT * FROM
@@ -764,7 +782,7 @@ class AdaptiveQueryExecSuite
             |    SELECT skewData2.key2, sum(skewData2.value2) AS sum2
             |    FROM skewData2 GROUP BY skewData2.key2
             |  ) AS data2
-            |ON data1.key1 = data2.key2
+            |ON data1.key1 = data2.key2 LIMIT 10
             |""".stripMargin
 
         val (_, adaptivePlan) = runAdaptiveAndVerifyResult(sqlText)