Implement custom spark optimization rules

Catalyst optimizer

Spark SQL Use a call catalyst The optimizer for all uses spark sql and dataframe dsl Query optimization . Optimized queries are better than using RDD Programs written directly run faster .catalyst yes rule based Optimizer , Many optimization rules are provided internally , These internal optimization rules will be introduced in detail later , Today, we will mainly discuss how to do this without modifying the source code , Write and apply our customized optimization rules in the way of plug-ins .

Write an optimization rule in practice

Let's make a simple optimization rule here , Realization function ： If we select A numeric field of type to multiply 1.0 This string , Let's optimize the multiplication calculation .

Write optimization rules

import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.expressions._

object MultiplyOptimizationRule extends Rule[LogicalPlan] with Logging {
    
  def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
    
    case Multiply(left,right) if right.isInstanceOf[Literal] &&
      right.asInstanceOf[Literal].value.asInstanceOf[Double] == 1.0 =>
      logInfo("MyRule  Optimization rules take effect ")
      left
  }
}

Register optimization rules

adopt spark Provide an interface to register the optimization rules we have written

spark.experimental.extraOptimizations = Seq(MultiplyOptimizationRule)

test result

Let's test it on the command line , We can see Project In the selected field ,(cast(id#7L as double) * 1.0) AS id2#12 Has been optimized to cast(id#7L as double) AS id2#14

scala> val df = spark.range(10).selectExpr("id", "concat('wankun-',id) as name")
df: org.apache.spark.sql.DataFrame = [id: bigint, name: string]

scala> val multipliedDF = df.selectExpr("id * cast(1.0 as double) as id2")
multipliedDF: org.apache.spark.sql.DataFrame = [id2: double]

scala> println(multipliedDF.queryExecution.optimizedPlan.numberedTreeString)
00 Project [(cast(id#7L as double) * 1.0) AS id2#12]
01 +- Range (0, 10, step=1, splits=Some(1))

import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.expressions._

object MultiplyOptimizationRule extends Rule[LogicalPlan] with Logging {
    
  def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
    
    case Multiply(left,right) if right.isInstanceOf[Literal] &&
      right.asInstanceOf[Literal].value.asInstanceOf[Double] == 1.0 =>
      logInfo("MyRule  Optimization rules take effect ")
      left
  }
}

scala> spark.experimental.extraOptimizations = Seq(MultiplyOptimizationRule)
spark.experimental.extraOptimizations: Seq[org.apache.spark.sql.catalyst.rules.Rule[org.apache.spark.sql.catalyst.plans.logical.LogicalPlan]] = List(MultiplyOptimizationRule$@675d209c)

scala>

scala> val multipliedDFWithOptimization = df.selectExpr("id * cast(1.0 as double) as id2")
multipliedDFWithOptimization: org.apache.spark.sql.DataFrame = [id2: double]

scala> println(multipliedDFWithOptimization.queryExecution.optimizedPlan.numberedTreeString)
00 Project [cast(id#7L as double) AS id2#14]
01 +- Range (0, 10, step=1, splits=Some(1))

Add hook and extension point functions

Through the above example , We go through spark Interface programming provided , It can be implemented to add our customized optimization rules .
But our spark-sql Tools don't allow us to add rules directly , in addition ,catalyst There's also... Inside Analysis, Logical Optimization, Physical Planning Multiple stages , If we want to make a function expansion in these places , It's not convenient . So in Spark 2.2 Version introduces a more powerful feature , Add hooks and extension points .

Write custom optimization rules and extension points

Or take the function of realizing function as an example :

package com.wankun.sql.optimizer

import org.apache.spark.internal.Logging
import org.apache.spark.sql.catalyst.expressions.{
    Literal, Multiply}
import org.apache.spark.sql.{
    SparkSession, SparkSessionExtensions}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.types.Decimal

/** * @author kun.wan * @date 2020-03-03. */
case class MyRule(spark: SparkSession) extends Rule[LogicalPlan] with Logging {
    
  override def apply(plan: LogicalPlan): LogicalPlan = {
    
    logInfo(" Start using  MyRule  Optimize the rules ")
    plan transformAllExpressions {
    
      case Multiply(left, right) if right.isInstanceOf[Literal] &&
        right.asInstanceOf[Literal].value.isInstanceOf[Decimal] &&
        right.asInstanceOf[Literal].value.asInstanceOf[Decimal].toDouble == 1.0 =>
        logInfo("MyRule  Optimization rules take effect ")
        left
    }
  }
}

class MyExtensions extends (SparkSessionExtensions => Unit) with Logging {
    
  def apply(e: SparkSessionExtensions): Unit = {
    
    logInfo(" Get into MyExtensions The extension point ")
    e.injectResolutionRule(MyRule)
  }
}

Package the above code as spark-extensions-1.0.jar

Configure and enable customization Spark Expand

spark-sql --master local --conf spark.sql.extensions=com.wankun.sql.optimizer.MyExtensions --jars /Users/wankun/ws/wankun/spark-extensions/target/spark-extensions-1.0.jar 

Test optimization rules

You can see plan By Analyzed after , Multiplication disappears , Customized optimization rules have taken effect .

spark-sql> explain extended
         > with stu as (
         >   select 1 as id, 'wankun-1' as name
         >   union
         >   select 2 as id, 'wankun-2' as name
         >   union
         >   select 3 as id, 'wankun-3' as name
         > )
         > select id * 1.0
         > from stu;
20/03/04 01:56:16 INFO MyRule: org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54)  Start using  MyRule  Optimize the rules 
20/03/04 01:56:17 INFO MyRule: org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54)  Start using  MyRule  Optimize the rules 
20/03/04 01:56:17 INFO MyRule: org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54)  Start using  MyRule  Optimize the rules 
20/03/04 01:56:17 INFO MyRule: org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54) MyRule  Optimization rules take effect 
20/03/04 01:56:17 INFO MyRule: org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54)  Start using  MyRule  Optimize the rules 
20/03/04 01:56:17 INFO MyRule: org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54)  Start using  MyRule  Optimize the rules 
20/03/04 01:56:17 INFO MyRule: org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54)  Start using  MyRule  Optimize the rules 
20/03/04 01:56:18 INFO CodeGenerator: org.apache.spark.internal.Logging$class.logInfo(Logging.scala:54) Code generated in 156.862003 ms
== Parsed Logical Plan ==
CTE [stu]
:  +- SubqueryAlias `stu`
:     +- Distinct
:        +- Union
:           :- Distinct
:           :  +- Union
:           :     :- Project [1 AS id#0, wankun-1 AS name#1]
:           :     :  +- OneRowRelation
:           :     +- Project [2 AS id#2, wankun-2 AS name#3]
:           :        +- OneRowRelation
:           +- Project [3 AS id#4, wankun-3 AS name#5]
:              +- OneRowRelation
+- 'Project [unresolvedalias(('id * 1.0), None)]
   +- 'UnresolvedRelation `stu`

== Analyzed Logical Plan ==
id: decimal(10,0)
Project [cast(id#0 as decimal(10,0)) AS id#8]
+- SubqueryAlias `stu`
   +- Distinct
      +- Union
         :- Distinct
         :  +- Union
         :     :- Project [1 AS id#0, wankun-1 AS name#1]
         :     :  +- OneRowRelation
         :     +- Project [2 AS id#2, wankun-2 AS name#3]
         :        +- OneRowRelation
         +- Project [3 AS id#4, wankun-3 AS name#5]
            +- OneRowRelation

== Optimized Logical Plan ==
Aggregate [id#0, name#1], [cast(id#0 as decimal(10,0)) AS id#8]
+- Union
   :- Project [1 AS id#0, wankun-1 AS name#1]
   :  +- OneRowRelation
   :- Project [2 AS id#2, wankun-2 AS name#3]
   :  +- OneRowRelation
   +- Project [3 AS id#4, wankun-3 AS name#5]
      +- OneRowRelation

== Physical Plan ==
*(5) HashAggregate(keys=[id#0, name#1], functions=[], output=[id#8])
+- Exchange hashpartitioning(id#0, name#1, 200)
   +- *(4) HashAggregate(keys=[id#0, name#1], functions=[], output=[id#0, name#1])
      +- Union
         :- *(1) Project [1 AS id#0, wankun-1 AS name#1]
         :  +- Scan OneRowRelation[]
         :- *(2) Project [2 AS id#2, wankun-2 AS name#3]
         :  +- Scan OneRowRelation[]
         +- *(3) Project [3 AS id#4, wankun-3 AS name#5]
            +- Scan OneRowRelation[]
Time taken: 1.945 seconds, Fetched 1 row(s)

sparkSession There are extension points left for users in ,Spark catalyst The extension point of is SPARK-18127 Introduced in ,Spark Users can go to SQL Extend custom implementations at various stages of processing , Very powerful and efficient

• injectOptimizerRule – add to optimizer Custom rule ,optimizer Be responsible for the optimization of logical execution plan , In our example, we extend the logic optimization rules .
• injectParser – add to parser Custom rule ,parser be responsible for SQL analysis .
• injectPlannerStrategy – add to planner strategy Custom rule ,planner Responsible for the generation of physical execution plan .
• injectResolutionRule – add to Analyzer Custom rule to Resolution Stage ,analyzer Responsible for generating logical execution plans .
• injectPostHocResolutionRule – add to Analyzer Custom rule to Post Resolution Stage .
• injectCheckRule – add to Analyzer Customize Check The rules .

Reference documents

• Catalyst: Allow adding custom optimizers
• Add hooks and extension points to Spark
• How to Extend Apache Spark with Customized Optimizations
• Introduction to Spark 2.0 - Part 6 : Custom Optimizers in Spark SQL
• understand spark sql The best way to optimize the strategy is to implement a