Spark3中Catalog组件设计和自定义扩展Catalog实现

Spark 3 中Catalog组件设计

catalog管理类继承关系

在Spark 3中，我们可以轻松的对内置的访问Hive metastore 的catalog进行自定义功能扩展。我们先梳理一下整个系统实现的类设计，以及catalog的初始化流程

Delta 是 Databrick 提供的一种扩展的文件存储格式，同时也提供了相关的SQL开发接口。我们看一下Delta项目中 DeltaCatalog 类的继承和实现功能关系。

CatalogPlugin: 每个独立的 namespaces 需要实现该接口, 默认的实现 V2SessionCatalog 的 namespace = Array("default")
    TableCatalog: 定义 Table相关操作的接口，包括 loadTable(), createTable(), alterTable(), dropTable() 等
        V2SessionCatalog: Spark 默认catalog实现，通过代理 SessionCatalog (也就是下面的 externalCatalog ) 对象，实现各种表的实际操作。
        CatalogExtension: 这个类同样继承了TableCatalog, 通过setDelegateCatalog()方法把上面的V2SessionCatalog 实例进行代理。
            DelegatingCatalogExtension: 该类实现了上面抽象类的所有默认方法，可以通过`super.func(ident)`调用默认实现，然后扩展我们自己要实现的逻辑。
                DeltaCatalog: 详细实现参考 Delta 提供的样例代码

总结一下，如果我们要实现自定义的Catalog，继承 TableCatalog 类就可以了；如果我们要实现对现有的catalog进行功能扩展，要继承 DelegatingCatalogExtension ，该类会提供很多方法的默认实现，我们要做扩展的话，直接扩展对应逻辑即可。

Catalog 初始化过程

Spark 通过CatalogManager 可以同时管理内部连接多个catalog，通过 spark.sql.catalog.${name} 可以注册多个catalog，Spark默认的catalog由 spark.sql.catalog.spark_catalog 参数指定，通常的做法是，自定义catalog类继承 DelegatingCatalogExtension 实现，然后通过 spark.sql.catalog.spark_catalog 参数来指定自定义catalog类。

详细看一下 HiveExternalCatalog, v2SessionCatalog, spark_catalog 等对象的实例化和管理流程

1. SparkSession在创建的选项中启用HiveSupport SparkSession.enableHiveSupport() ， 在该方法内会设置参数 CATALOG_IMPLEMENTATION = hive， Spark SQL 的catalog类型(spark.sql.catalogImplementation 参数) 默认支持 hive 和 in-memory 两种，如果没有指定，默认为 in-memory.
2. session.sharedState.externalCatalog 是Spark Session 实际负责和外部系统交互的catalog, 根据上面设置的参数，分别会实例化出 HiveExternalCatalog 和 InMemoryCatalog 两个实例。
3. 在BaseSessionStateBuilder / HiveSessionStateBuilder 中会使用上面的 externalCatalog 创建 catalog 对象，再根据 catalog 创建 v2SessionCatalog 对象
4. 根据 catalog 和 v2SessionCatalog 创建 CatalogManager 实例。 CatalogManager 通过 catalogs Map[catalog name, catalog instance] 对象来管理多个catalog。CatalogManager 的 defaultSessionCatalog 属性就是上面的 v2SessionCatalog 对象。
5. CatalogManager.catalog(name: String) 通过catalog的name返回catalog实例，如果没有该实例，则通过 Catalogs.load(name: String, ...) 方法进行实例化。
6. Catalogs.load(name: String, ...) 方法加载conf中配置的 spark.sql.catalog.${name} 类，并实例化/初始化(initialize) CatalogPlugin 对象
7. 有一个特殊的catalog，name = spark_catalog。如果 spark.sql.catalog.spark_catalog 参数为空(默认为空)时，返回CatalogManager 中的 defaultSessionCatalog 属性.
8. 如果 spark.sql.catalog.spark_catalog 参数已经配置， 对上面Catalogs.load() 出来的实例进行判断，如果发现上面加载的是CatalogExtension子类，自动调用其 setDelegateCatalog() 方法，将 catalogManager 中 defaultSessionCatalog 设置为其内部代理对象。

补充相关实现代码:

BaseSessionStateBuilder 实例的初始化

  protected lazy val catalog: SessionCatalog = {
    val catalog = new SessionCatalog(
      () => session.sharedState.externalCatalog,
      () => session.sharedState.globalTempViewManager,
      functionRegistry,
      conf,
      SessionState.newHadoopConf(session.sparkContext.hadoopConfiguration, conf),
      sqlParser,
      resourceLoader)
    parentState.foreach(_.catalog.copyStateTo(catalog))
    catalog
  }

  protected lazy val v2SessionCatalog = new V2SessionCatalog(catalog, conf)

  protected lazy val catalogManager = new CatalogManager(conf, v2SessionCatalog, catalog)

HiveSessionStateBuilder extends BaseSessionStateBuilder

子类的catalog返回的是 HiveSessionCatalog 实例

  override protected lazy val catalog: HiveSessionCatalog = {
    val catalog = new HiveSessionCatalog(
      () => externalCatalog,
      () => session.sharedState.globalTempViewManager,
      new HiveMetastoreCatalog(session),
      functionRegistry,
      conf,
      SessionState.newHadoopConf(session.sparkContext.hadoopConfiguration, conf),
      sqlParser,
      resourceLoader)
    parentState.foreach(_.catalog.copyStateTo(catalog))
    catalog
  }

测试自定义Catalog

每个catalog都有一个name，内部有一组namespaces数组，系统默认的 catalog name = spark_catalog, namespace = Array(“default”)。我们来测试使用一个自定义的catalog。

编译和打包测试类

备注: 测试类不能在spark shell窗口中直接写，REPL底层使用的是 scala 内置的 scala.tools.nsc.interpreter 包下面的工具进行的代码即时编译和加载，但是对应的class文件被修改了， 对应的类无参构造方法没有，所以会出问题。稳妥一点的方法还是把测试类编译为Jar再进行测试

package org.apache.spark.wankun.catalog

import org.apache.spark.sql.connector.catalog.CatalogPlugin
import org.apache.spark.sql.util.CaseInsensitiveStringMap

// CatalogPlugin 子类必须要提供无参的构造方法，可以在 initialize 方法中进行初始化
class DummyCatalog extends CatalogPlugin {
  override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {
    _name = name
  }
  private var _name: String = null
  override def name(): String = _name
  override def defaultNamespace(): Array[String] = Array("a", "b")
}

切换catalog测试

// spark-shell --master local --jars /Users/wankun/ws/wankun/sbtstudy/target/scala-2.12/sbtstudy_2.12-1.0.jar

scala> spark.sessionState.catalogManager.currentCatalog
res0: org.apache.spark.sql.connector.catalog.CatalogPlugin = V2SessionCatalog(spark_catalog)

scala> spark.sessionState.catalogManager.currentCatalog.name
res1: String = spark_catalog

scala> spark.sessionState.catalogManager.currentNamespace
res2: Array[String] = Array(default)

scala>

scala> import org.apache.spark.wankun.catalog.DummyCatalog
import org.apache.spark.wankun.catalog.DummyCatalog

scala> spark.sessionState.conf.setConfString("spark.sql.catalog.dummy", classOf[DummyCatalog].getName)

scala> spark.sessionState.catalogManager.currentCatalog
res4: org.apache.spark.sql.connector.catalog.CatalogPlugin = V2SessionCatalog(spark_catalog)

scala> spark.sessionState.catalogManager.setCurrentCatalog("dummy")

scala> spark.sessionState.catalogManager.currentCatalog
res6: org.apache.spark.sql.connector.catalog.CatalogPlugin = org.apache.spark.wankun.catalog.DummyCatalog@735666c4

scala> spark.sessionState.catalogManager.currentCatalog.name
res7: String = dummy

scala> spark.sessionState.catalogManager.currentNamespace
res8: Array[String] = Array(a, b)

自定义JDBC和Kafka数据源的Catalog实战

上面基本分析清楚了spark sql中自定义扩展catalog的实现原理，在spark sql的解析过程中，通过ResolveTables Rule 可以帮我们自动根据catalog(第一个namespace)，找到并加载对应的 NonSessionCatalogAndIdentifier，该catalog根据我们传入的剩余的namespace 和name， 返回我们自定义的Table；最后通过DataSourceV2Relation.create(table, Some(catalog), Some(ident)) 返回Table包装的Relaiton对象。

规则代码如下:

  object ResolveTables extends Rule[LogicalPlan] {
    def apply(plan: LogicalPlan): LogicalPlan = ResolveTempViews(plan).resolveOperatorsUp {
      case u: UnresolvedRelation =>
        lookupV2Relation(u.multipartIdentifier)
          .map { rel =>
            val ident = rel.identifier.get
            SubqueryAlias(rel.catalog.get.name +: ident.namespace :+ ident.name, rel)
          }.getOrElse(u)
      ...
    }

    /**
     * Performs the lookup of DataSourceV2 Tables from v2 catalog.
     */
    private def lookupV2Relation(identifier: Seq[String]): Option[DataSourceV2Relation] =
      expandRelationName(identifier) match {
        case NonSessionCatalogAndIdentifier(catalog, ident) =>
          CatalogV2Util.loadTable(catalog, ident) match {
            case Some(table) =>
              Some(DataSourceV2Relation.create(table, Some(catalog), Some(ident)))
            case None => None
          }
        case _ => None
      }
  }

最后我们来实现自定义JDBC和Kafka数据源的Catalog， 详细代码参考: JDBC catalog实现 和 Kafka Catalog实现

参考

• Spark Catalog Plugin 机制介绍
• Easy Guide to Create a Custom Read Data Source in Apache Spark 3
• SparkSQL DatasourceV2 之 Multiple Catalog