back to use propagate template method in QueryPlanner.plan for extension strategy compability

Author: mingbo.pb

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

@@ -60,7 +60,9 @@ abstract class QueryPlanner[PhysicalPlan <: TreeNode[PhysicalPlan]] {
     // Obviously a lot to do here still...
 
     // Collect physical plan candidates.
-    val candidates = strategies.iterator.flatMap(_(plan))
+    val candidates = strategies.iterator.flatMap({
+      _(plan).map(propagate(_, plan))
+    })
 
     // The candidates may contain placeholders marked as [[planLater]],
     // so try to replace them by their child plans.
@@ -102,4 +104,7 @@ abstract class QueryPlanner[PhysicalPlan <: TreeNode[PhysicalPlan]] {
 
   /** Prunes bad plans to prevent combinatorial explosion. */
   protected def prunePlans(plans: Iterator[PhysicalPlan]): Iterator[PhysicalPlan]
+
+  /** Propagate logicalPlan properties to PhysicalPlan */
+  protected def propagate(plan: PhysicalPlan, logicalPlan: LogicalPlan): PhysicalPlan
 }

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

@@ -65,6 +65,10 @@ class SparkPlanner(
     plans
   }
 
+  override protected def propagate(plan: SparkPlan, logicalPlan: LogicalPlan): SparkPlan = {
+    plan.withStats(logicalPlan.stats)
+  }
+
   /**
    * Used to build table scan operators where complex projection and filtering are done using
    * separate physical operators.  This function returns the given scan operator with Project and

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

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