auto calculate the initial partition number with ae (apache#61)

Author: luzhonghao

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

@@ -91,7 +91,7 @@ case class HandleSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] {
       partitionId: Int,
       medianSize: Long,
       medianRowCount: Long): Array[Int] = {
-    val stats = queryStageInput.childStage.stats
+    val stats = queryStageInput.childStage.statsPlan
     val size = stats.bytesByPartitionId.get(partitionId)
     val rowCount = stats.recordStatistics.get.recordsByPartitionId(partitionId)
     val factor = Math.max(size / medianSize, rowCount / medianRowCount)
@@ -110,8 +110,8 @@ case class HandleSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] {
       left: QueryStageInput,
       right: QueryStageInput): Boolean = {
       supportedJoinTypes.contains(joinType) &&
-      left.childStage.stats.getPartitionStatistics.isDefined &&
-      right.childStage.stats.getPartitionStatistics.isDefined
+      left.childStage.statsPlan.getPartitionStatistics.isDefined &&
+      right.childStage.statsPlan.getPartitionStatistics.isDefined
   }
 
   private def supportSplitOnLeftPartition(joinType: JoinType) = joinType != RightOuter
@@ -128,8 +128,8 @@ case class HandleSkewedJoin(conf: SQLConf) extends Rule[SparkPlan] {
       SortExec(_, _, right: ShuffleQueryStageInput, _))
       if supportOptimization(joinType, left, right) =>
 
-      val leftStats = left.childStage.stats.getPartitionStatistics.get
-      val rightStats = right.childStage.stats.getPartitionStatistics.get
+      val leftStats = left.childStage.statsPlan.getPartitionStatistics.get
+      val rightStats = right.childStage.statsPlan.getPartitionStatistics.get
       val numPartitions = leftStats.bytesByPartitionId.length
       val (leftMedSize, leftMedRowCount) = medianSizeAndRowCount(leftStats)
       val (rightMedSize, rightMedRowCount) = medianSizeAndRowCount(rightStats)

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

@@ -40,7 +40,8 @@ case class OptimizeJoin(conf: SQLConf) extends Rule[SparkPlan] {
   }
 
   private def canBroadcast(plan: SparkPlan): Boolean = {
-    plan.stats.sizeInBytes >= 0 && plan.stats.sizeInBytes <= conf.adaptiveBroadcastJoinThreshold
+    val ret = plan.statsPlan.sizeInBytes >= 0
+    ret && plan.statsPlan.sizeInBytes <= conf.adaptiveBroadcastJoinThreshold
   }
 
   private def removeSort(plan: SparkPlan): SparkPlan = {
@@ -85,9 +86,9 @@ case class OptimizeJoin(conf: SQLConf) extends Rule[SparkPlan] {
     // reading them in local shuffle read.
     broadcastSidePlan match {
       case broadcast: ShuffleQueryStageInput
-        if broadcast.childStage.stats.bytesByPartitionId.isDefined =>
+        if broadcast.childStage.statsPlan.bytesByPartitionId.isDefined =>
           val (startIndicies, endIndicies) = calculatePartitionStartEndIndices(broadcast.childStage
-            .stats.bytesByPartitionId.get)
+            .statsPlan.bytesByPartitionId.get)
           childrenPlans.foreach {
             case input: ShuffleQueryStageInput =>
               input.partitionStartIndices = Some(startIndicies)

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

@@ -351,7 +351,7 @@ class PlannerSuite extends SharedSQLContext {
     )
     withSQLConf(SQLConf.SHUFFLE_TARGET_POSTSHUFFLE_INPUT_SIZE.key -> "1") {
 
-      val totalInputFileSize = inputPlan.collectLeaves().map(_.stats.sizeInBytes).sum
+      val totalInputFileSize = inputPlan.collectLeaves().map(_.statsPlan.sizeInBytes).sum
       val expectedNum = Math.ceil(
         totalInputFileSize.toLong * 1.0 / conf.targetPostShuffleInputSize).toInt