扩展Spark Catalyst，打造自定义的Spark SQL引擎

2022-01-30 05:00:09

Apache Spark是大数据处理领域最常用的计算引擎之一，被应用在各种各样的场景中，除了易用的API，稳定高效的处理引擎，可扩展性也是Spark能够得到广泛应用的一个重要原因。Spark中最常见的扩展就是基于DataSource API添加对新数据源的支持，除了Spark内置的HDFS，S3，Kafka等数据源，Parquet，Orc，Avro等数据类型，还有很多第三方的DataSource Plugin使得Spark成为大数据领域可以处理数据源类型最丰富的计算引擎。当然，除了DataSource，Spark还有支持很多其他的扩展点，今天我们要介绍的是Spark SQL Catalyst的扩展点，以及如何通过这些扩展点实现一些有意思的功能，打造自定义的Spark SQL引擎。

在Spark2.2版本中，引入了新的扩展点，使得用户可以在Spark session中自定义自己的parser，analyzer，optimizer以及physical planning stragegy rule。

【注意】本篇文章需要读者对Scala语言，SQL引擎基本概念以及Spark Catalyst引擎有基本的了解。

Spark Catalyst扩展点

Spark catalyst的扩展点在SPARK-18127中被引入，Spark用户可以在SQL处理的各个阶段扩展自定义实现，非常强大高效，下面我们具体看看其提供的接口和在Spark中的实现。

SparkSessionExtensions

SparkSessionExtensions保存了所有用户自定义的扩展规则，自定义规则保存在成员变量中，对于不同阶段的自定义规则，SparkSessionExtensions提供了不同的接口。

新增自定义规则

用户可以通过SparkSessionExtensions提供的inject开头的方法添加新的自定义规则，具体的inject接口如下：

injectOptimizerRule – 添加optimizer自定义规则，optimizer负责逻辑执行计划的优化。
injectParser – 添加parser自定义规则，parser负责SQL解析。
injectPlannerStrategy – 添加planner strategy自定义规则，planner负责物理执行计划的生成。
injectResolutionRule – 添加Analyzer自定义规则到Resolution阶段，analyzer负责逻辑执行计划生成。
injectPostHocResolutionRule – 添加Analyzer自定义规则到Post Resolution阶段。
injectCheckRule – 添加Analyzer自定义Check规则。

【注】Spark Catalyst的SQL处理分成parser，analyzer，optimizer以及planner等多个步骤，其中analyzer，optimizer等步骤内部也分为多个阶段，以Analyzer为例，analyse规则切分到不同的batch中，每个batch的执行策略可能不尽相同，有的只会执行一遍，有的会迭代执行直到满足一定条件。具体每个步骤的每个阶段的具体实现请参考Spark源码，本文篇幅有限，不再详述。

获取自定义规则

SparkSessionExtensions对应每一种自定义规则也都有一个build开头的方法用于获取对应类型的自定义规则，Spark session在初始化的时候，通过这些方法获取自定义规则并传递给parser，analyzer，optimizer以及planner等对象。

buildOptimizerRules
buildParser
buildPlannerStrategies
buildResolutionRules
buildPostHocResolutionRules
buildCheckRules

配置自定义规则

在Spark中，用户自定义的规则可以通过两种方式配置生效：

使用SparkSession.Builder中的withExtenstion方法，withExtension方法是一个高阶函数，接收一个自定义函数作为参数，这个自定义函数以SparkSessionExtensions作为参数，用户可以实现这个函数，通过SparkSessionExtensions的inject开头的方法添加用户自定义规则。
通过Spark配置参数，具体参数名为spark.sql.extensions。用户可以将1中的自定义函数实现定义为一个类，将完整类名作为参数值。

具体的用法用户可以参考org.apache.spark.sql.SparkSessionExtensionSuite测试用例中的Spark代码。

扩展Spark Catalyst实现SQL检查

有许多用户使用Spark SQL构建数据分析查询平台，会有许多的业务用户使用这个平台进行数据分析处理，由于业务用户SQL开发能力参差不齐，作为平台方很难约束用户，一个不合理的SQL查询不仅可能导致容易出错，很难维护，可能还会直接搞垮整个平台。通过Spark Catalyst扩展点，平台方可以解析所有用户的SQL，制定具体的SQL使用规范，对于不合理的SQL请求直接拒绝。本文我们以一个非常简单的规范为例，展示如何通过自定义规则检查SQL。SELECT 是一个非常常见的SQL查询方式，用于获取表的所有列数据，但是这种SQL的可维护性相对来说会比较差，表可能增加新列或者删除已有列，甚至列的展示顺序也可能发生变化，这些都会影响SQL执行的结果以及依赖此查询的后续查询，一个规范严格的数据平台SQL规范可能不允许这种情况发生。作为本文的示例，我们可以添加一个简单的检查，发现SELECT 请求就报错，不允许执行。

创建一个自定义Parser

通过集成ParserInterface，实现自定义Parser。

class StrictParser(parser: ParserInterface) extends ParserInterface {
  /**
   * Parse a string to a [[LogicalPlan]].
   */
  override def parsePlan(sqlText: String): LogicalPlan = {
    val logicalPlan = parser.parsePlan(sqlText)
    logicalPlan transform {
      case project @ Project(projectList, _) =>
        projectList.foreach {
          name =>
            if (name.isInstanceOf[UnresolvedStar]) {
              throw new RuntimeException("You must specify your project column set," +
                " * is not allowed.")
            }
        }
        project
    }
    logicalPlan
  }

  /**
   * Parse a string to an [[Expression]].
   */
  override def parseExpression(sqlText: String): Expression = parser.parseExpression(sqlText)

  /**
   * Parse a string to a [[TableIdentifier]].
   */
  override def parseTableIdentifier(sqlText: String): TableIdentifier =
    parser.parseTableIdentifier(sqlText)

  /**
   * Parse a string to a [[FunctionIdentifier]].
   */
  override def parseFunctionIdentifier(sqlText: String): FunctionIdentifier =
    parser.parseFunctionIdentifier(sqlText)

  /**
   * Parse a string to a [[StructType]]. The passed SQL string should be a comma separated
   * list of field definitions which will preserve the correct Hive metadata.
   */
  override def parseTableSchema(sqlText: String): StructType =
    parser.parseTableSchema(sqlText)

  /**
   * Parse a string to a [[DataType]].
   */
  override def parseDataType(sqlText: String): DataType = parser.parseDataType(sqlText)
}

创建扩展点函数

type ParserBuilder = (SparkSession, ParserInterface) => ParserInterface
type ExtensionsBuilder = SparkSessionExtensions => Unit
val parserBuilder: ParserBuilder = (_, parser) => new StrictParser(parser)
val extBuilder: ExtensionsBuilder = { e => e.injectParser(parserBuilder)}

这里面有两个函数，extBuilder函数用于SparkSession构建，SparkSessionExtensions.injectParser函数本身也是一个高阶函数，接收parserBuilder作为参数，将原生parser作为参数传递给自定义的StrictParser，并将StrictParser作为自定义parser插入SparkSessionExtensions中。

在SparkSession中启用自定义Parser

val spark = SparkSession
  .builder()
  .appName("Spark SQL basic example")
  .config("spark.master", "local[2]")
  .withExtensions(extBuilder)
  .getOrCreate()

测试代码

测试代码很简单，我们首先构建一个简单的测试集，验证包含SELECT *的SQL语句是否按照预想抛出异常即可，测试代码如下：

val spark = SparkSession
  .builder()
  .appName("Spark SQL basic example")
  .config("spark.master", "local[2]")
  .withExtensions(extBuilder)
  .getOrCreate()
val df = spark.read.json("examples/src/main/resources/people.json")
    df.toDF().write.saveAsTable("person")
    spark.sql("select * from person limit 3").show
spark.stop()

执行的结果如下：

扩展Spark Catalyst优化SQL执行

本次测试用例优化规则参考http://blog.madhukaraphatak.com/introduction-to-spark-two-part-6/，由于Spark版本依赖不同，具体实现代码略有不同，具体的测试代码如下：

val spark = SparkSession
  .builder()
  .appName("Spark SQL basic example")
  .config("spark.master", "local[2]")
  .getOrCreate()
  val df = spark.read
      .format("com.databricks.spark.csv")
      .option("header", "true")
      .option("delimiter", ";")
      .load("examples/src/main/resources/people.csv")
    df.toDF.write.saveAsTable("person")

    // scalastyle:off
    println(spark.sql("select age * 1 from person").queryExecution.optimizedPlan
      .numberedTreeString)
    // scalastyle:on
    spark.stop

在示例代码中，我们加载了一个csv文件，然后用age字段乘1，我们可以通过queryExecution.optimizedPlan获取优化后的执行计划。LogicalPlan以树状结构组织，通过numberedTreeString我们可以得到优化后执行计划的字符串信息。

00 Project [(cast(age#26 as double) * 1.0) AS (CAST(age AS DOUBLE) * CAST(1 AS DOUBLE))#28]
01 +- Relation[name#25,age#26,job#27] parquet

执行计划包括两个部分：

01 Relation - 标示我们通过csv文件创建的表。

00 Project - 标示Project投影操作。

由于age字段默认是string类型，可以看到spark默认会cast成double类型，然后乘以1.0.

创建自定义Optimizer规则

从之前的执行计划可以看到，每个age字段都会乘以1.0，但是实际上我们知道只并不必要，任何值乘以1都等于本身，我们可以据此添加自定义Optimizer规则省略掉乘以1.0的计算。

case class MultiplyOptimizationRule(spark: SparkSession) extends Rule[LogicalPlan] {
  def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
    case mul @ Multiply(left, right) if right.isInstanceOf[Literal] &&
        right.asInstanceOf[Literal].value.asInstanceOf[Double] == 1.0 =>
      left
  }
}

在SparkSession中启用自定义Optimizer规则

type ExtensionsBuilder = SparkSessionExtensions => Unit
val extBuilder: ExtensionsBuilder = { e => e.injectOptimizerRule(MultiplyOptimizationRule)}

val spark = SparkSession
  .builder()
  .appName("Spark SQL basic example")
  .config("spark.master", "local[2]")
  .withExtensions(extBuilder)
  .getOrCreate()

测试用例

使用同样的测试用例，启用自定义Optimizer规则，具体代码如下：

type ExtensionsBuilder = SparkSessionExtensions => Unit
val extBuilder: ExtensionsBuilder = { e => e.injectOptimizerRule(MultiplyOptimizationRule)}


// $example on:init_session$
val spark = SparkSession
  .builder()
  .appName("Spark SQL basic example")
  .config("spark.master", "local[2]")
  .withExtensions(extBuilder)
  .getOrCreate()

val df = spark.read
      .format("com.databricks.spark.csv")
      .option("header", "true")
      .option("delimiter", ";")
      .load("examples/src/main/resources/people.csv")
    df.toDF.write.saveAsTable("person")

    // scalastyle:off
    println(spark.sql("select age * 1 from person").queryExecution.optimizedPlan
      .numberedTreeString)
    // scalastyle:on
    spark.stop

添加自定义优化规则后，新的执行计划如下：

00 Project [cast(age#26 as double) AS (CAST(age AS DOUBLE) * CAST(1 AS DOUBLE))#28]
01 +- Relation[name#25,age#26,job#27] parquet

通过对比之前的执行计划，我们可以看到新的执行计划没有了乘以1.0的步骤。

总结

在Spark2.2版本中，引入了新的扩展点，使得用户可以在Spark session中自定义自己的parser，analyzer，optimizer以及physical planning stragegy rule。通过两个简单的示例，我们展示了如何通过Spark提供的扩展点实现SQL检查以及定制化的执行计划优化。Spark Catalyst高度的可扩展性使得我们可以非常方便的定制适合自己实际使用场景的SQL引擎，拓展了更多的可能性。我们可以实现特定的SQL方言，针对特殊的数据源做更深入的优化，进行SQL规范检查，针对特定执行环境制定特定的优化策略等等，本文抛砖引玉，希望能对读者有所裨益。

欢迎加群指正交流

码农公寓