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