Spark Datafusion Comet 向量化Rust Native--CometShuffleExchangeExec怎么控制读写

背景

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)
主要分析Rust Native的Shuffle算子 CometShuffleExchangeExec怎么控制 Native Shuffle Writer还是Columnar Shuffle Writer

CometShuffleExchangeExec

在Spark中 ShuffleExchangeExec 这个算子起到了承上启下的作用，他控制了 Shuffle writer的写姿势，也控制了Shuffle read的读姿势，在Spark Comet中CometShuffleExchangeExec也是一样，其中最关键的在于其生成的CometShuffledBatchRDD:

  protected override def doExecuteColumnar(): RDD[ColumnarBatch] = {
    // Returns the same CometShuffledBatchRDD if this plan is used by multiple plans.
    if (cachedShuffleRDD == null) {
      cachedShuffleRDD = new CometShuffledBatchRDD(shuffleDependency, readMetrics)
    }
    cachedShuffleRDD
  }

其中 shuffleDependency是Shuffle Writer的关键，CometShuffledBatchRDD 是Shuffle Read的关键

• ShuffleDependency

     lazy val shuffleDependency: ShuffleDependency[Int, _, _] =
    if (shuffleType == CometNativeShuffle) {
      val dep = CometShuffleExchangeExec.prepareShuffleDependency(
        inputRDD.asInstanceOf[RDD[ColumnarBatch]],
        child.output,
        outputPartitioning,
        serializer,
        metrics)
      metrics("numPartitions").set(dep.partitioner.numPartitions)
      val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
      SQLMetrics.postDriverMetricUpdates(
        sparkContext,
        executionId,
        metrics("numPartitions") :: Nil)
      dep
    } else if (shuffleType == CometColumnarShuffle) {
      val dep = CometShuffleExchangeExec.prepareJVMShuffleDependency(
        inputRDD.asInstanceOf[RDD[InternalRow]],
        child.output,
        outputPartitioning,
        serializer,
        metrics)
      metrics("numPartitions").set(dep.partitioner.numPartitions)
      val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
      SQLMetrics.postDriverMetricUpdates(
        sparkContext,
        executionId,
        metrics("numPartitions") :: Nil)
      dep
    } else {
      throw new UnsupportedOperationException(
        s"Unsupported shuffle type: ${shuffleType.getClass.getName}")
    }
  
  

  • 该ShuffleDependency 会根据之前生成的Comet计划的时候的类型(是CometNativeShuffle,还是CometColumnarShuffle）进行不同的判断
    具体可以参考Spark Datafusion Comet 向量化Rule–CometExecRule分析 规则转换分析，简单的说就是如果shuffle的子节点都是CometNative的话，就是CometNativeShuffle,
    此时生成的CometShuffleDependency就是CometNativeShuffle的，否则就是CometColumnarShuffle，且生成的CometShuffleDependency是CometColumnarShuffle,这样在注册到CometShuffleManager 的时候才会根据不同的类型返回不同的ShuffleHandle,这样在MapTask获取不同的Writer的时候才会选择CometNativeShuffleWriter还是CometUnsafeShuffleWriter
  • 规范化Shuffle writer的写的分区布局
    在此前阶段CometShuffleExchangeExec此时的outputPartitioning只是Partitioning的逻辑描述，此后，将会转换为Partitioner类型的具体的可计算的物理算子
    • 对于CometColumnarShuffle而言，会根据不同的Partitioning生成不同的Partitioner ,且传给CometShuffleDependency的RDD是带有partitionId的Product2 类型的RDD,后续会根据这个partitionId进行排序。
    • 对于CometNativeShuffle来说,和CometColumnarShuffle 一样，只不过传给CometShuffleDependency的RDD是partitionId为0的Product2 类型的RDD,后续的写数据是由Rust Native 的dataFusion组建控制，它这里会构建对应的分区算法,所以这里的分区数随便设置为0

• CometShuffledBatchRDD
  这个是ReduceTask计算的时候用到,会从Map端拉取数据shuffleHandle选择不同的reader,当然这里的reader都是同一个CometBlockStoreShuffleReader