Spark Datafusion Comet 向量化Rule--CometScanRule分析

背景

Apache Datafusion Comet 是苹果公司开源的加速Spark运行的向量化项目。
本项目采用了 Spark插件化 + Protobuf + Arrow + DataFusion 架构形式
其中

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

本文基于 datafusion comet 截止到2026年1月13号的main分支的最新代码(对应的commit为 eef5f28a0727d9aef043fa2b87d6747ff68b827a)
主要分析 CometSparkSessionExtensions 中的向量化规则转换

CometScanRule分析

上代码：

class CometSparkSessionExtensions
    extends (SparkSessionExtensions => Unit)
    with Logging
    with ShimCometSparkSessionExtensions {
  override def apply(extensions: SparkSessionExtensions): Unit = {
    extensions.injectColumnar { session => CometScanColumnar(session) }
    extensions.injectColumnar { session => CometExecColumnar(session) }
    extensions.injectQueryStagePrepRule { session => CometScanRule(session) }
    extensions.injectQueryStagePrepRule { session => CometExecRule(session) }
  }

  case class CometScanColumnar(session: SparkSession) extends ColumnarRule {
    override def preColumnarTransitions: Rule[SparkPlan] = CometScanRule(session)
  }

  case class CometExecColumnar(session: SparkSession) extends ColumnarRule {
    override def preColumnarTransitions: Rule[SparkPlan] = CometExecRule(session)

    override def postColumnarTransitions: Rule[SparkPlan] =
      EliminateRedundantTransitions(session)
  }
}

这里的 CometScanColumnar 最终也是调用 CometScanRule 方法,所以这里只分析 CometScanRule 这个规则：

case class CometScanRule(session: SparkSession) extends Rule[SparkPlan] with CometTypeShim {

  import CometScanRule._

  private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get()

  override def apply(plan: SparkPlan): SparkPlan = {
    val newPlan = _apply(plan)
    if (showTransformations && !newPlan.fastEquals(plan)) {
      logInfo(s"""
           |=== Applying Rule $ruleName ===
           |${sideBySide(plan.treeString, newPlan.treeString).mkString("\n")}
           |""".stripMargin)
    }
    newPlan
  }

这里有个spark.comet.explain.rules配置（默认是false）,用来表明经过这个规则转换后，计划的前后变化。
对于_apply方法：

  private def _apply(plan: SparkPlan): SparkPlan = {
    if (!isCometLoaded(conf)) return plan

    def isSupportedScanNode(plan: SparkPlan): Boolean = plan match {
      case _: FileSourceScanExec => true
      case _: BatchScanExec => true
      case _ => false
    }

    def hasMetadataCol(plan: SparkPlan): Boolean = {
      plan.expressions.exists(_.exists {
        case a: Attribute =>
          a.isMetadataCol
        case _ => false
      })
    }

    def transformScan(plan: SparkPlan): SparkPlan = plan match {
      case scan if !isCometScanEnabled(conf) =>
        withInfo(scan, "Comet Scan is not enabled")

      case scan if hasMetadataCol(scan) =>
        withInfo(scan, "Metadata column is not supported")

      // data source V1
      case scanExec: FileSourceScanExec =>
        transformV1Scan(scanExec)

      // data source V2
      case scanExec: BatchScanExec =>
        if (isIcebergMetadataTable(scanExec)) {
          withInfo(scanExec, "Iceberg Metadata tables are not supported")
        } else {
          transformV2Scan(scanExec)
        }
    }

    plan.transform {
      case scan if isSupportedScanNode(scan) => transformScan(scan)
    }
  }

• 首先 isCometLoaded方法，这里会有个lazy val isBigEndian: Boolean = ByteOrder.nativeOrder().equals(ByteOrder.BIG_ENDIAN)
  判断平台是不是大端的判断，如果是大端的话，就直接退出。
  这里为什么是大端就直接退出了呢，是因为一般来说 如果涉及到网路传输的时候，会涉及到多字节整数的存储问题，如果网络传输方和接收方在大小端不一致的时候
  会导致传输和接收的数据不一致的问题。
  而在Datafusion Comet这个项目中使用小端进行数据的交互，主要是在Shuffle数据文件的时候，具体的可以见 NativeBatchDecoderIterator :

    private val longBuf = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN)
  

  这里会以小端的形式进行存储，所以平台测如果是大端的话就直接退出了

• 保证 spark.comet.enabled(默认为true)为True

• 还有保证rust 动态库被成功加载
  NativeBase.isLoaded 就是来判断对应的动态库是不是已经被成功加载，如果没成功加载，则直接报错
  动态库的加载 在 NativeBase的 静态代码块中

• transformScan主要用来对scan的计划进行转换，分Datasource V1和V2

  • datasource V1 (FileSourceScanExec)
    这里的代码比较多，直接说重点
    • 只有relation为HadoopFsRelation类型的才继续
    • 如果存在动态分区裁剪的话且(spark.comet.dppFallback.enabled为true（默认是true）)，就直接返回Spark 自带的FileSourceScanExec，
    • 对于 spark.comet.scan.impl 为auto的（默认为auto）, 如果不是不是本地文件以及S3文件系统，则设置为 SCAN_NATIVE_COMET
      最终组成CometScanExec(HadoopFsRelation(fileFormat = new CometParquetFileFormat))
  • datasource V2(BatchScanExec)
    • 只有Scan为ParquetScan才继续
    • 并且满足一定的数据类型才可以
    • 如果scan中有聚合下推，则返回Spark BatchScanExec额scan,否则最终会返回CometBatchScanExec