Spark Catalyst 添加自己的规则

abstract class Rule[TreeType <: TreeNode[_]] extends Logging {

  /** Name for this rule, automatically inferred based on class name. */
  val ruleName: String = {
    val className = getClass.getName
    if (className endsWith "$") className.dropRight(1) else className
  }

  def apply(plan: TreeType): TreeType
}

abstract class GenericStrategy[PhysicalPlan <: TreeNode[PhysicalPlan]] extends Logging {

  /**
   * Returns a placeholder for a physical plan that executes `plan`. This placeholder will be
   * filled in automatically by the QueryPlanner using the other execution strategies that are
   * available.
   */
  protected def planLater(plan: LogicalPlan): PhysicalPlan

  def apply(plan: LogicalPlan): Seq[PhysicalPlan]
}

  // 注入一个Rule
	type RuleBuilder = SparkSession => Rule[LogicalPlan]
	// 用于check而已，只是check LogicalPlan，如果不通过，会抛异常，通过则不做任何操作，所以返回类型为Unit
  type CheckRuleBuilder = SparkSession => LogicalPlan => Unit
	// 注入一个Strategy
  type StrategyBuilder = SparkSession => Strategy
	// 注入一个Parser, 用于语法解析
  type ParserBuilder = (SparkSession, ParserInterface) => ParserInterface

  def injectResolutionRule(builder: RuleBuilder): Unit 

  def injectPostHocResolutionRule(builder: RuleBuilder): Unit 

def injectCheckRule(builder: CheckRuleBuilder): Unit

def injectOptimizerRule(builder: RuleBuilder): Unit 

def injectPlannerStrategy(builder: StrategyBuilder): Unit

def injectParser(builder: ParserBuilder): Unit

case class MyResolutionRule(spark: SparkSession) extends Rule[LogicalPlan] {
  override def apply(plan: LogicalPlan): LogicalPlan = plan
}
case class MyPostHocResolutionRule(spark: SparkSession) extends Rule[LogicalPlan] {
  override def apply(plan: LogicalPlan): LogicalPlan = plan
}

case class MyOptimizerRule(spark: SparkSession) extends Rule[LogicalPlan] {
  override def apply(plan: LogicalPlan): LogicalPlan = plan
}

case class MyCheckRule(spark: SparkSession) extends (LogicalPlan => Unit) {
  override def apply(plan: LogicalPlan): Unit = { }
}

case class MySparkStrategy(spark: SparkSession) extends SparkStrategy {
  override def apply(plan: LogicalPlan): Seq[SparkPlan] = Seq.empty
}

case class MyParser(spark: SparkSession, delegate: ParserInterface) extends ParserInterface {
  override def parsePlan(sqlText: String): LogicalPlan =
    delegate.parsePlan(sqlText)

  override def parseExpression(sqlText: String): Expression =
    delegate.parseExpression(sqlText)

  override def parseTableIdentifier(sqlText: String): TableIdentifier =
    delegate.parseTableIdentifier(sqlText)

  override def parseFunctionIdentifier(sqlText: String): FunctionIdentifier =
    delegate.parseFunctionIdentifier(sqlText)

  override def parseTableSchema(sqlText: String): StructType =
    delegate.parseTableSchema(sqlText)

  override def parseDataType(sqlText: String): DataType =
    delegate.parseDataType(sqlText)
}

class MyExtensions extends (SparkSessionExtensions => Unit) {
  def apply(e: SparkSessionExtensions): Unit = {
    e.injectResolutionRule(MyResolutionRule)
    e.injectPostHocResolutionRule(MyPostHocResolutionRule)
    e.injectCheckRule(MyCheckRule)
    e.injectOptimizerRule(MyOptimizerRule)
    e.injectPlannerStrategy(MySparkStrategy)
    e.injectParser(MyParser)
  }
}

import org.apache.spark.internal.Logging
import org.apache.spark.sql.{AnalysisException, SparkSession, SparkSessionExtensions}
import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project}
import org.apache.spark.sql.execution.datasources.{InsertIntoDataSourceCommand, LogicalRelation}
import org.apache.spark.sql.execution.datasources.greenplum.GreenplumRelation

case class GreenPlumColumnChecker(spark: SparkSession) extends (LogicalPlan => Unit) with Logging {
  override def apply(plan: LogicalPlan): Unit = plan match {
    case InsertIntoDataSourceCommand(LogicalRelation(_: GreenplumRelation, output, _, _),
    Project(_, c), _) =>
      // The real output of sub query, which is not be casted.
      val realOutput = c.output
      if (realOutput.size != output.size || realOutput.zip(output).exists(
        ats => ats._1.name != ats._2.name)) {
        throw new AnalysisException(
          s"""
             | The column names of GreenPlum table are not consistent with the
             | projects output names of subQuery.
           """.stripMargin)
      }
    case InsertIntoDataSourceCommand(LogicalRelation(_: GreenplumRelation, _, _, _), query, _) =>
      query.
      logWarning(s"GreenPlumColumnChecker: The query of this GreenPlumRelation " +
        s"is a ${query.getClass.getName}.")
    case _ =>
  }
}

class GreenPlumColumnCheckerExtension extends (SparkSessionExtensions => Unit) {
  override def apply(e: SparkSessionExtensions): Unit = {
    e.injectCheckRule(GreenPlumColumnChecker)
  }
}