Spark datafusion comet向量化插件CometPlugin

背景

Apache Datafusion Comet 是苹果公司开源的加速Spark运行的向量化项目。

本项目采用了 Spark插件化 + Brotobuf + Arrow + DataFusion 架构形式

其中

Spark插件是利用 SparkPlugin 插件，其中分为 DriverPlugin 和 ExecutorPlugin ,这两个插件在driver和 Executor启动的时候就会调用
Brotobuf 是用来序列化 spark对应的表达式以及计划，用来传递给 native 引擎去执行,利用了体积小，速度快的特性
Arrow 是用来 spark 和 native 引擎进行高效的数据交换(native执行的结果或者spark执行的数据结果)，主要在JNI中利用Arrow IPC 列式存储以及零拷贝等特点进行进程间数据交换
DataFusion 主要是利用Rust native以及Arrow内存格式实现的向量化执行引擎，Spark中主要offload对应的算子到该引擎中去执行

本文基于 datafusion comet 截止到2026年1月13号的main分支的最新代码(对应的commit为 eef5f28a0727d9aef043fa2b87d6747ff68b827a)

主要分析 CometPlugin spark的插件

分析

Spark插件CometPlugin

复制代码

class CometPlugin extends SparkPlugin with Logging {
  override def driverPlugin(): DriverPlugin = new CometDriverPlugin

  override def executorPlugin(): ExecutorPlugin = null
}

该CometPlugin在SparkContext初始化中使用到：

复制代码

_plugins = PluginContainer(this, _resources.asJava)

对于 CometDriverPlugin 这个DriverPlugin 主要来设置向量化转换的Rule(主要通过CometSparkSessionExtensions)。

复制代码

  override def init(sc: SparkContext, pluginContext: PluginContext): ju.Map[String, String] = {
    logInfo("CometDriverPlugin init")

    if (!CometSparkSessionExtensions.isOffHeapEnabled(sc.getConf) &&
      !sc.getConf.getBoolean(COMET_ONHEAP_ENABLED.key, false)) {
      logWarning("Comet plugin is disabled because Spark is not running in off-heap mode.")
      return Collections.emptyMap[String, String]
    }

    // register CometSparkSessionExtensions if it isn't already registered
    CometDriverPlugin.registerCometSessionExtension(sc.conf)

    if (CometSparkSessionExtensions.shouldOverrideMemoryConf(sc.getConf)) {
      val execMemOverhead = if (sc.getConf.contains(EXECUTOR_MEMORY_OVERHEAD.key)) {
        sc.getConf.getSizeAsMb(EXECUTOR_MEMORY_OVERHEAD.key)
      } else {
        // By default, executorMemory * spark.executor.memoryOverheadFactor, with minimum of 384MB
        val executorMemory =
          sc.getConf.getSizeAsMb(EXECUTOR_MEMORY.key, EXECUTOR_MEMORY_DEFAULT)
        val memoryOverheadFactor = sc.getConf.get(EXECUTOR_MEMORY_OVERHEAD_FACTOR)
        val memoryOverheadMinMib = getMemoryOverheadMinMib(sc.getConf)

        Math.max((executorMemory * memoryOverheadFactor).toLong, memoryOverheadMinMib)
      }

      val cometMemOverhead = CometSparkSessionExtensions.getCometMemoryOverheadInMiB(sc.getConf)
      sc.conf.set(EXECUTOR_MEMORY_OVERHEAD.key, s"${execMemOverhead + cometMemOverhead}M")
      val newExecMemOverhead = sc.getConf.getSizeAsMb(EXECUTOR_MEMORY_OVERHEAD.key)

      logInfo(s"""
         Overriding Spark memory configuration for Comet:
           - Spark executor memory overhead: ${execMemOverhead}MB
           - Comet memory overhead: ${cometMemOverhead}MB
           - Updated Spark executor memory overhead: ${newExecMemOverhead}MB
         """)
    } else {
      logInfo("Comet is running in unified memory mode and sharing off-heap memory with Spark")
    }

    Collections.emptyMap[String, String]
  }

如果说 spark.memory.offHeap.enabled(默认为false) 和 spark.comet.exec.onHeap.enabled(默认为false),都为false的话,则直接返回
CometDriverPlugin.registerCometSessionExtension(sc.conf)

这里会设置 CometSparkSessionExtensions 为sparkSession扩展，主要是设置CometScanRule 和 CometExecRule 规则
设置Executor端内存(前提是comet native 执行加上comet shuffle开启以及使用堆外内存)

如果 spark.comet.enabled (默认 true)，(spark.comet.exec.shuffle.enabled (默认true)，spark.comet.exec.enabled(默认是true)) ,spark.memory.offHeap.enabled(默认false) 都开启的前提下,才会继续设置内存，否则直接退出

jvm额外内存 = (Spark executor jvm运行的内存(spark.executor.memory 默认1g)*(spark.executor.memoryOverheadFactor 默认0.1) , spark.executor.minMemoryOverhead 默认384MB) 的最大值

comet额外内存 = 如果是堆外则0MB ，否则(spark.comet.memoryOverhead 默认 1GB)

最终设置 spark.executor.memoryOverhead 为 (jvm额外内存 + comet额外内存)