[CARMEL-6352] Adjust scan partition size dynamically considering potential cost (#1110)

* [CARMEL-6352] Adjust scan partition size dynamically considering potential cost

* fix ut

* Add optimize tag

* minor

* Add limit for partition number

Author: xingchaozh

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala

@@ -142,6 +142,7 @@ case class Stack(children: Seq[Expression]) extends Generator {
   private lazy val numRows = children.head.eval().asInstanceOf[Int]
   private lazy val numFields = Math.ceil((children.length - 1.0) / numRows).toInt
 
+  def getNumRows: Int = numRows
   /**
    * Return true iff the first child exists and has a foldable IntegerType.
    */

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/hints.scala

@@ -17,8 +17,8 @@
 
 package org.apache.spark.sql.catalyst.plans.logical
 
+import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.expressions.Attribute
-import org.apache.spark.util.Utils
 
 /**
  * A general hint for the child that is not yet resolved. This node is generated by the parser and
@@ -63,13 +63,16 @@ case class ResolvedParallelHint(child: LogicalPlan,
 
 trait HintPlaceHolderInfo
 
-case class TableParallelInfo(tableIdentifier: String, partitionSize: Option[Long],
-                             partitionNumber: Option[Int]) extends HintPlaceHolderInfo {
+case class TableParallelInfo(tableIdentifier: Option[TableIdentifier],
+                             partitionSize: Option[Long],
+                             partitionNumber: Option[Int],
+                             partitionSizeReduceRadio: Int = 1) extends HintPlaceHolderInfo {
   override def toString: String = {
     Seq(
-      s"tableIdentifier=${tableIdentifier}",
-      s"partitionSize=${partitionSize}",
-      s"partitionNumber=${partitionNumber}")
+      s"table=${tableIdentifier.getOrElse("None")}",
+      s"size=${partitionSize}",
+      s"number=${partitionNumber}",
+      s"reduceRadio=${partitionSizeReduceRadio}")
       .mkString("[", ", ", "]")
   }
 }

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

@@ -1311,6 +1311,15 @@ object SQLConf {
           "must be greater than 1.0")
       .createWithDefault(10.0)
 
+  val AUTO_ADJUST_SCAN_SIZE_MAX_PARTITIONS =
+    buildConf("spark.sql.sources.autoAdjustScanSize.maxPartitions")
+    .doc("The maximum number of partitions allowed when scan size is adjusted.")
+    .version("3.0.0")
+    .intConf
+    .checkValue(_ > 0,
+      "the value of spark.sql.sources.autoAdjustScanSize.maxPartitions must be greater than 0")
+    .createWithDefault(100000)
+
   val AUTO_BUCKETED_SCAN_ENABLED =
     buildConf("spark.sql.sources.bucketing.autoBucketedScan.enabled")
       .doc("When true, decide whether to do bucketed scan on input tables based on query plan " +
@@ -1335,6 +1344,14 @@ object SQLConf {
           "must be great than or equal to 1")
       .createWithDefault(6000)
 
+  val AUTO_ADJUST_SCAN_PARTITION_SIZE_ENABLED =
+    buildConf("spark.sql.autoAdjustScanPartitionSize.enabled")
+      .internal()
+      .doc("When true, we will adjust scan partition size dynamically.")
+      .version("3.0.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val REMOVE_REDUNDANT_PARTIAL_AGGREGATES_ENABLED =
     buildConf("spark.sql.removeRedundantPartialAggregates.enabled")
       .internal()
@@ -4414,6 +4431,8 @@ class SQLConf extends Serializable with Logging {
 
   def initPartialListingCount: Int = getConf(SQLConf.INIT_PARTIAL_LISTING_COUNT)
 
+  def autoAdjustScanSizeMaxPartitions: Int = getConf(SQLConf.AUTO_ADJUST_SCAN_SIZE_MAX_PARTITIONS)
+
   def bucketingEnabled: Boolean = getConf(SQLConf.BUCKETING_ENABLED)
 
   def bucketingRatio: Double = getConf(SQLConf.BUCKETING_RATIO)

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

@@ -909,6 +909,32 @@ case class FileSourceScanExec(
         originSize * fsRelation.sparkSession.sessionState.conf.minParallelHintPartitionSizeRatio),
         originSize * fsRelation.sparkSession.sessionState.conf.maxParallelHintPartitionSizeRatio).
         longValue()
+    } else if (tableParallelInfo.nonEmpty
+      && tableParallelInfo.get.partitionSizeReduceRadio > 1) {
+      val specifiedMaxSplitBytes = Math.
+        max(originSize / tableParallelInfo.get.partitionSizeReduceRadio,
+        originSize * fsRelation.sparkSession.sessionState.conf.minParallelHintPartitionSizeRatio).
+        longValue()
+
+      val partitionNumber =
+        FilePartition.minPartitionNumberBySpecifiedSize(fsRelation.sparkSession,
+        selectedPartitions, specifiedMaxSplitBytes)
+
+      val originPartitionNumber =
+        FilePartition.minPartitionNumberBySpecifiedSize(fsRelation.sparkSession,
+        selectedPartitions, originSize)
+
+      if (partitionNumber > conf.autoAdjustScanSizeMaxPartitions) {
+        if (originPartitionNumber < conf.autoAdjustScanSizeMaxPartitions) {
+          FilePartition.maxSplitBytesBySpecifiedNumber(
+            fsRelation.sparkSession, selectedPartitions, conf.autoAdjustScanSizeMaxPartitions)
+        } else {
+          logInfo(s"Fallback to origin scan size for table ${tableParallelInfo}")
+          originSize
+        }
+      } else {
+        specifiedMaxSplitBytes
+      }
     } else if (relation.sparkSession.sessionState.conf.bucketingEnabled &&
       relation.bucketSpec.isDefined &&
       disableBucketedScan) { // Check if bucketing scan disabled by planner

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

@@ -35,7 +35,7 @@ import org.apache.spark.sql.catalyst.util.StringUtils.PlanStringConcat
 import org.apache.spark.sql.catalyst.util.truncatedString
 import org.apache.spark.sql.execution.QueryExecution.skipAuthTag
 import org.apache.spark.sql.execution.adaptive.{AdaptiveExecutionContext, EnsureRepartitionForWriting, InsertAdaptiveSparkPlan}
-import org.apache.spark.sql.execution.bucketing.DisableUnnecessaryBucketedScan
+import org.apache.spark.sql.execution.bucketing.{AdjustScanPartitionSizeDynamically, DisableUnnecessaryBucketedScan}
 import org.apache.spark.sql.execution.dynamicpruning.PlanDynamicPruningFilters
 import org.apache.spark.sql.execution.exchange.{EliminateShuffleExec, EnsureRequirements, ExchangePushDownThroughAggregate}
 import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery
@@ -373,6 +373,7 @@ object QueryExecution {
       EnsureRepartitionForWriting,
       EliminateShuffleExec,
       DisableUnnecessaryBucketedScan,
+      AdjustScanPartitionSizeDynamically,
       ApplyColumnarRulesAndInsertTransitions(sparkSession.sessionState.columnarRules),
       CollapseCodegenStages()) ++
       (if (subquery) {

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

@@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreeNodeTag
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec._
-import org.apache.spark.sql.execution.bucketing.DisableUnnecessaryBucketedScan
+import org.apache.spark.sql.execution.bucketing.{AdjustScanPartitionSizeDynamically, DisableUnnecessaryBucketedScan}
 import org.apache.spark.sql.execution.exchange._
 import org.apache.spark.sql.execution.ui.{SparkListenerSQLAdaptiveExecutionUpdate, SparkListenerSQLAdaptiveSQLMetricUpdates, SQLPlanMetric}
 import org.apache.spark.sql.internal.SQLConf
@@ -96,7 +96,8 @@ case class AdaptiveSparkPlanExec(
     RemoveRedundantPartialAggregates,
     EnsureRepartitionForWriting,
     EliminateShuffleExec,
-    DisableUnnecessaryBucketedScan
+    DisableUnnecessaryBucketedScan,
+    AdjustScanPartitionSizeDynamically
   ) ++ context.session.sessionState.queryStagePrepRules
 
   @transient private val initialPlan = context.session.withActive {

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

@@ -43,6 +43,7 @@ object OptimizeSkewedRangePartition extends Rule[SparkPlan]  {
   private def handleSkewedRangePartition(plan: SortExec, child: SparkPlan): SparkPlan = {
     val queryStageInfo = child match {
       case ShuffleStage(s: ShuffleStageInfo) => Option(s)
+      case _ => None
     }
 
     if (queryStageInfo.isEmpty ||

sql/core/src/main/scala/org/apache/spark/sql/execution/bucketing/AdjustScanPartitionSizeDynamically.scala

@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.bucketing
+
+import org.apache.spark.sql.catalyst.expressions.Stack
+import org.apache.spark.sql.catalyst.expressions.aggregate.{ApproximatePercentile, Percentile}
+import org.apache.spark.sql.catalyst.plans.logical.TableParallelInfo
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution._
+import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec
+import org.apache.spark.sql.execution.exchange.Exchange
+import org.apache.spark.sql.execution.joins.{BroadcastNestedLoopJoinExec, BroadcastRangeJoinExec}
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Adjust scan partition size dynamically considering potential cost.
+ * We check operators on path from root/exchange to scan operator to calculate the cost
+ */
+object AdjustScanPartitionSizeDynamically extends Rule[SparkPlan] {
+  def apply(plan: SparkPlan): SparkPlan = {
+    if (!conf.getConf(SQLConf.AUTO_ADJUST_SCAN_PARTITION_SIZE_ENABLED)) {
+      plan
+    } else {
+      adjustScanPartitionSize(plan)
+    }
+  }
+
+  // Visit operators from root/exchange to current scan
+  private def adjustScanPartitionSize(plan: SparkPlan, inputReduceRadio: Int = 1): SparkPlan = {
+    var reduceRadio = inputReduceRadio
+    plan match {
+      case e: Exchange =>
+        e.mapChildren(adjustScanPartitionSize(_, 1))
+      case scan: FileSourceScanExec =>
+        if (!scan.bucketedScan && reduceRadio > 1) {
+          if (scan.tableParallelInfo.isEmpty || (scan.tableParallelInfo.get.partitionNumber.isEmpty
+            && scan.tableParallelInfo.get.partitionSize.isEmpty)) {
+            val tableParallelInfo = Some(TableParallelInfo(scan.tableIdentifier, None,
+              None, reduceRadio))
+
+            val newScan = scan.copy(tableParallelInfo = tableParallelInfo)
+            scan.logicalLink.foreach(newScan.setLogicalLink)
+            newScan.addOptimizeTag(s"created by ${simpleRuleName}")
+            newScan
+          } else {
+            scan
+          }
+        } else {
+          scan
+        }
+      case o =>
+        reduceRadio = reduceRadio + visitCostInterestingOperator(o)
+        o.mapChildren(adjustScanPartitionSize(_, reduceRadio))
+    }
+  }
+
+  private def visitCostInterestingOperator(plan: SparkPlan): Int = {
+    var reduceRatio = 0
+    plan match {
+      case ObjectHashAggregateExec(_, _, aggregateExpr, _, _, _, _) =>
+        aggregateExpr.foreach(_.aggregateFunction match {
+          case _: Percentile => reduceRatio += 1
+          case _: ApproximatePercentile => reduceRatio += 1
+          case _ =>
+        })
+      case _: BroadcastNestedLoopJoinExec => reduceRatio += 3
+      case expand: ExpandExec => reduceRatio += expand.projections.size
+      case GenerateExec(stack: Stack, _, _, _, _) => reduceRatio += stack.getNumRows
+      case _ =>
+    }
+    reduceRatio
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/CacheParallelHint.scala

@@ -53,7 +53,7 @@ private[sql] object CacheParallelHint
         logicalRelation.catalogTable.foreach(catalogTable => {
           SparkContext.getActive.foreach(sc => {
             val tableParallelInfo =
-              TableParallelInfo(catalogTable.identifier.toString, r.partitionSize,
+              TableParallelInfo(Some(catalogTable.identifier), r.partitionSize,
                 r.partitionNumber)
 
             val queryPlanningContext = if (QueryPlanningTracker.getCurrent.isEmpty) {

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala

@@ -216,20 +216,20 @@ object FileSourceStrategy extends Strategy with SQLConfHelper with Logging {
 
       if (conf.parallelHintEnabled) {
         l.catalogTable.foreach(catalogTable => {
-          val tableIdentifier = catalogTable.identifier.toString
+          val tableIdentifier = catalogTable.identifier
 
           plan.collectFirst {
             case r: ResolvedParallelHint if ((r.partitionSize.nonEmpty ||
               r.partitionNumber.nonEmpty)) => r
           }.foreach(r => {
-            tableParallelInfo = Some(TableParallelInfo(tableIdentifier, r.partitionSize,
+            tableParallelInfo = Some(TableParallelInfo(Some(tableIdentifier), r.partitionSize,
               r.partitionNumber))
           })
 
           if (tableParallelInfo.isEmpty) {
             QueryPlanningTracker.getCurrent.foreach(f => {
               tableParallelInfo = Option(
-                f.queryPlanningContext.tableParallelHintMap.get(tableIdentifier)
+                f.queryPlanningContext.tableParallelHintMap.get(tableIdentifier.toString)
               )
             })
           }