propagate the statistics from logical plan to physical plan in Strategy.apply method

Author: pengbo

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/QueryPlanner.scala

@@ -65,8 +65,6 @@ abstract class QueryPlanner[PhysicalPlan <: TreeNode[PhysicalPlan]] {
     // The candidates may contain placeholders marked as [[planLater]],
     // so try to replace them by their child plans.
     val plans = candidates.flatMap { candidate =>
-      propagateProperty(candidate, plan)
-
       val placeholders = collectPlaceholders(candidate)
 
       if (placeholders.isEmpty) {
@@ -96,9 +94,6 @@ abstract class QueryPlanner[PhysicalPlan <: TreeNode[PhysicalPlan]] {
     pruned
   }
 
-  protected def propagateProperty(candidate: PhysicalPlan, plan: LogicalPlan): Unit = {
-  }
-
   /**
    * Collects placeholders marked using [[GenericStrategy#planLater planLater]]
    * by [[strategies]].

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

@@ -53,10 +53,6 @@ class SparkPlanner(
    */
   def extraPlanningStrategies: Seq[Strategy] = Nil
 
-  override protected def propagateProperty(candidate: SparkPlan, plan: LogicalPlan): Unit = {
-    candidate.withStats(plan.stats)
-  }
-
   override protected def collectPlaceholders(plan: SparkPlan): Seq[(SparkPlan, LogicalPlan)] = {
     plan.collect {
       case placeholder @ PlanLater(logicalPlan) => placeholder -> logicalPlan

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

@@ -49,6 +49,12 @@ import org.apache.spark.sql.types.StructType
 abstract class SparkStrategy extends GenericStrategy[SparkPlan] {
 
   override protected def planLater(plan: LogicalPlan): SparkPlan = PlanLater(plan)
+
+  override def apply(plan: LogicalPlan): Seq[SparkPlan] = {
+    doApply(plan).map(sparkPlan => sparkPlan.withStats(plan.stats))
+  }
+
+  protected def doApply(plan: LogicalPlan): Seq[SparkPlan]
 }
 
 case class PlanLater(plan: LogicalPlan) extends LeafExecNode {
@@ -67,7 +73,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * Plans special cases of limit operators.
    */
   object SpecialLimits extends Strategy {
-    override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case ReturnAnswer(rootPlan) => rootPlan match {
         case Limit(IntegerLiteral(limit), Sort(order, true, child))
             if limit < conf.topKSortFallbackThreshold =>
@@ -209,7 +215,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         hint.rightHint.exists(_.strategy.contains(SHUFFLE_REPLICATE_NL))
     }
 
-    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
 
       // If it is an equi-join, we first look at the join hints w.r.t. the following order:
       //   1. broadcast hint: pick broadcast hash join if the join type is supported. If both sides
@@ -383,7 +389,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * on-demand, only when planning in a [[org.apache.spark.sql.execution.streaming.StreamExecution]]
    */
   object StatefulAggregationStrategy extends Strategy {
-    override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case _ if !plan.isStreaming => Nil
 
       case EventTimeWatermark(columnName, delay, child) =>
@@ -423,7 +429,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * Used to plan the streaming deduplicate operator.
    */
   object StreamingDeduplicationStrategy extends Strategy {
-    override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case Deduplicate(keys, child) if child.isStreaming =>
         StreamingDeduplicateExec(keys, planLater(child)) :: Nil
 
@@ -440,7 +446,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * Limit is unsupported for streams in Update mode.
    */
   case class StreamingGlobalLimitStrategy(outputMode: OutputMode) extends Strategy {
-    override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case ReturnAnswer(rootPlan) => rootPlan match {
         case Limit(IntegerLiteral(limit), child)
             if plan.isStreaming && outputMode == InternalOutputModes.Append =>
@@ -455,7 +461,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
   }
 
   object StreamingJoinStrategy extends Strategy {
-    override def apply(plan: LogicalPlan): Seq[SparkPlan] = {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = {
       plan match {
         case ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, condition, left, right, _)
           if left.isStreaming && right.isStreaming =>
@@ -476,7 +482,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * Used to plan the aggregate operator for expressions based on the AggregateFunction2 interface.
    */
   object Aggregation extends Strategy {
-    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case PhysicalAggregation(groupingExpressions, aggExpressions, resultExpressions, child)
         if aggExpressions.forall(expr => expr.isInstanceOf[AggregateExpression]) =>
         val aggregateExpressions = aggExpressions.map(expr =>
@@ -538,7 +544,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
   }
 
   object Window extends Strategy {
-    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case PhysicalWindow(
         WindowFunctionType.SQL, windowExprs, partitionSpec, orderSpec, child) =>
         execution.window.WindowExec(
@@ -556,7 +562,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
   protected lazy val singleRowRdd = sparkContext.parallelize(Seq(InternalRow()), 1)
 
   object InMemoryScans extends Strategy {
-    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case PhysicalOperation(projectList, filters, mem: InMemoryRelation) =>
         pruneFilterProject(
           projectList,
@@ -574,7 +580,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * be replaced with the real relation using the `Source` in `StreamExecution`.
    */
   object StreamingRelationStrategy extends Strategy {
-    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case s: StreamingRelation =>
         StreamingRelationExec(s.sourceName, s.output) :: Nil
       case s: StreamingExecutionRelation =>
@@ -590,7 +596,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * in streaming plans. Conversion for batch plans is handled by [[BasicOperators]].
    */
   object FlatMapGroupsWithStateStrategy extends Strategy {
-    override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case FlatMapGroupsWithState(
         func, keyDeser, valueDeser, groupAttr, dataAttr, outputAttr, stateEnc, outputMode, _,
         timeout, child) =>
@@ -608,7 +614,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    * Strategy to convert EvalPython logical operator to physical operator.
    */
   object PythonEvals extends Strategy {
-    override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case ArrowEvalPython(udfs, output, child) =>
         ArrowEvalPythonExec(udfs, output, planLater(child)) :: Nil
       case BatchEvalPython(udfs, output, child) =>
@@ -619,7 +625,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
   }
 
   object BasicOperators extends Strategy {
-    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case d: DataWritingCommand => DataWritingCommandExec(d, planLater(d.query)) :: Nil
       case r: RunnableCommand => ExecutedCommandExec(r) :: Nil
 

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

@@ -261,7 +261,7 @@ class FindDataSourceTable(sparkSession: SparkSession) extends Rule[LogicalPlan]
 case class DataSourceStrategy(conf: SQLConf) extends Strategy with Logging with CastSupport {
   import DataSourceStrategy._
 
-  def apply(plan: LogicalPlan): Seq[execution.SparkPlan] = plan match {
+  override protected def doApply(plan: LogicalPlan): Seq[execution.SparkPlan] = plan match {
     case PhysicalOperation(projects, filters, l @ LogicalRelation(t: CatalystScan, _, _, _)) =>
       pruneFilterProjectRaw(
         l,

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

@@ -136,7 +136,7 @@ object FileSourceStrategy extends Strategy with Logging {
     }
   }
 
-  def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+  override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
     case PhysicalOperation(projects, filters,
       l @ LogicalRelation(fsRelation: HadoopFsRelation, _, table, _)) =>
       // Filters on this relation fall into four categories based on where we can use them to avoid

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

@@ -102,7 +102,7 @@ object DataSourceV2Strategy extends Strategy with PredicateHelper {
 
   import DataSourceV2Implicits._
 
-  override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+  override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
     case PhysicalOperation(project, filters, relation: DataSourceV2Relation) =>
       val scanBuilder = relation.newScanBuilder()
 

sql/core/src/test/scala/org/apache/spark/sql/ExtraStrategiesSuite.scala

@@ -39,7 +39,7 @@ case class FastOperator(output: Seq[Attribute]) extends SparkPlan {
 }
 
 object TestStrategy extends Strategy {
-  def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+  override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
     case Project(Seq(attr), _) if attr.name == "a" =>
       FastOperator(attr.toAttribute :: Nil) :: Nil
     case _ => Nil

sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala

@@ -218,7 +218,7 @@ case class MyCheckRule(spark: SparkSession) extends (LogicalPlan => Unit) {
 }
 
 case class MySparkStrategy(spark: SparkSession) extends SparkStrategy {
-  override def apply(plan: LogicalPlan): Seq[SparkPlan] = Seq.empty
+  override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = Seq.empty
 }
 
 case class MyParser(spark: SparkSession, delegate: ParserInterface) extends ParserInterface {
@@ -272,7 +272,7 @@ case class MyCheckRule2(spark: SparkSession) extends (LogicalPlan => Unit) {
 }
 
 case class MySparkStrategy2(spark: SparkSession) extends SparkStrategy {
-  override def apply(plan: LogicalPlan): Seq[SparkPlan] = Seq.empty
+  override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = Seq.empty
 }
 
 object MyExtensions2 {

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

@@ -33,7 +33,7 @@ class SparkPlannerSuite extends SharedSQLContext {
 
     var planned = 0
     object TestStrategy extends Strategy {
-      def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+      override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
         case ReturnAnswer(child) =>
           planned += 1
           planLater(child) :: planLater(NeverPlanned) :: Nil

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveStrategies.scala

@@ -223,7 +223,7 @@ private[hive] trait HiveStrategies {
   val sparkSession: SparkSession
 
   object Scripts extends Strategy {
-    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case ScriptTransformation(input, script, output, child, ioschema) =>
         val hiveIoSchema = HiveScriptIOSchema(ioschema)
         ScriptTransformationExec(input, script, output, planLater(child), hiveIoSchema) :: Nil
@@ -236,7 +236,7 @@ private[hive] trait HiveStrategies {
    * applied.
    */
   object HiveTableScans extends Strategy {
-    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+    override protected def doApply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case PhysicalOperation(projectList, predicates, relation: HiveTableRelation) =>
         // Filter out all predicates that only deal with partition keys, these are given to the
         // hive table scan operator to be used for partition pruning.