深入剖析 Apache DataFusion Comet：Rust Native 物理执行计划的构建与优化

在大数据计算领域，性能优化是永恒的主题。Apache Spark作为主流计算引擎，其性能瓶颈常出现在JVM的序列化与GC开销上。苹果公司开源的 Apache DataFusion Comet 项目，正是为了解决这一痛点而生。它通过向量化执行和Rust Native引擎，将Spark算子的执行效率推向新的高度。本文将聚焦于Comet的核心环节——Rust Native物理执行计划的创建过程，深入解析其跨语言协作与高性能设计的精妙之处。

一、Comet架构概览：跨语言协作的向量化加速器

在深入代码之前，理解Comet的整体架构至关重要。它并非一个独立的执行引擎，而是一个精巧的Spark插件化加速方案，其核心思想是“将合适的计算卸载到更快的引擎中执行”。整个架构融合了多种技术栈，各司其职：

• Spark Plugin：作为入口，分为DriverPlugin和ExecutorPlugin，分别在Driver和Executor启动时加载，负责生命周期的管理。
• Protocol Buffers (Protobuf)：承担跨语言通信的“信使”角色。Spark的表达式和逻辑/物理计划被序列化为Protobuf格式，传递给Rust Native引擎。其体积小、序列化/反序列化速度快的特性，是保证低延迟通信的关键。
• Apache Arrow：作为进程间高效数据交换的“高速公路”。它定义了统一的内存列式格式，使得Spark（JVM）与DataFusion（Rust）之间可以通过Arrow IPC进行零拷贝或极低开销的数据传递，彻底避免了传统序列化带来的性能损耗。
• DataFusion (Rust Native)：真正的执行核心。这是一个用Rust编写的、基于Arrow内存格式的向量化查询引擎。Spark会将过滤、投影、聚合等计算密集型算子“下推”给DataFusion执行，充分利用Rust的性能优势和向量化处理能力。

本文的分析基于DataFusion Comet项目在 eef5f28a0727d9aef043fa2b87d6747ff68b827a 提交时的最新代码，我们将重点关注物理执行计划在Rust侧的构建逻辑。

二、计划创建的入口：Java侧的桥梁与参数准备

物理计划的创建始于Spark的Executor端。当Spark决定将某个算子（对应CometNativeExec节点）交由Comet执行时，会调用其doExecuteColumnar方法，进而触发Native物理计划的构建。这一过程主要封装在 CometExecIterator 类中。

该类提供了两个核心的创建入口：

1. CometNativeShuffleWriter：在启用Native Shuffle时，用于构造Native的Shuffle Writer计划，相比JVM写入Shuffle文件，能获得更高的I/O效率。
2. CometNativeExec：最常用的入口，用于为单个Native算子创建执行计划，后续从Native引擎获取计算结果。

我们重点分析后者。其Java侧的调用逻辑如下：

 private val plan = {
    val conf = SparkEnv.get.conf
    val localDiskDirs = SparkEnv.get.blockManager.getLocalDiskDirs
    // serialize Comet related Spark configs in protobuf format
    val protobufSparkConfigs = CometExecIterator.serializeCometSQLConfs()
    // Create keyUnwrapper if encryption is enabled
    val keyUnwrapper = if (encryptedFilePaths.nonEmpty) {
      val unwrapper = new CometFileKeyUnwrapper()
      val hadoopConf: Configuration = broadcastedHadoopConfForEncryption.get.value.value
      encryptedFilePaths.foreach(filePath =>
        unwrapper.storeDecryptionKeyRetriever(filePath, hadoopConf))
      unwrapper
    } else {
      null
    }
    val memoryConfig = CometExecIterator.getMemoryConfig(conf)
    nativeLib.createPlan(
      id,
      cometBatchIterators,
      protobufQueryPlan,
      protobufSparkConfigs,
      numParts,
      nativeMetrics,
      metricsUpdateInterval = COMET_METRICS_UPDATE_INTERVAL.get(),
      cometTaskMemoryManager,
      localDiskDirs,
      batchSize = COMET_BATCH_SIZE.get(),
      memoryConfig.offHeapMode,
      memoryConfig.memoryPoolType,
      memoryConfig.memoryLimit,
      memoryConfig.memoryLimitPerTask,
      taskAttemptId,
      taskCPUs,
      keyUnwrapper)
  }

这段代码在 Executor端 执行。它主要完成以下几项准备工作：

• 配置收集：通过protobufSparkConfigs方法，收集所有以 spark.comet 开头的Spark配置项，并将其序列化为Protobuf格式的字节流。
• 内存配置获取：调用CometExecIterator.getMemoryConfig获取任务的内存配置参数。
• 参数传递：将序列化后的Spark计划、配置、内存信息、度量节点引用等，通过JNI传递给Rust函数。

这里，offHeapMode 为 spark.comet.exec.memoryPool，其值为 fair_unified。方法最终返回的是一个指向Rust侧创建的 ExecutionContext 对象的原始指针（内存地址）。

三、Rust侧的核心逻辑：反序列化、配置与会话构建

当调用通过JNI进入Rust世界后，真正的计划创建在 jni_api.rs 模块的 pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan 函数中展开。这个过程可以分解为几个清晰的步骤：

1. 配置反序列化与环境准备

首先，Rust需要解析从Java传来的配置信息。

 let bytes = env.convert_byte_array(serialized_spark_configs)?;
 let spark_configs: datafusion_comet_proto::spark_config::ConfigMap = serde::deserialize_config(bytes.as_slice())?;
 let spark_config: HashMap = spark_configs.entries.into_iter().collect();

代码使用 into_iter().collect() 将配置Map转换并消耗所有权。为了便捷地获取配置值，项目对HashMap进行了扩展，提供了 get_bool（用于布尔值）和 get_u64（用于整数值）等方法。这些方法会处理键不存在或值格式错误的情况，增强了健壮性。

接着，初始化 JVMClasses 和 JVMClasses::init(&mut env)，并使用 let env = unsafe { std::mem::transmute::<&mut JNIEnv, &'static mut JNIEnv>(env) }; 和 std::mem::transmute 将JNI环境转换为静态类型，以便在后续闭包中使用。

2. Spark计划反序列化与对象持有

接下来是反序列化Spark物理计划的核心步骤：

 let bytes = env.convert_byte_array(serialized_query)?;
 let spark_plan = serde::deserialize_op(bytes.as_slice())?;

这里，Protobuf格式的字节流被反序列化为Rust侧的 Protobuf 结构体。同时，为了防止Java垃圾回收器（GC）在Rust使用期间意外回收关键的Java对象（如度量节点CometMetricNode），代码会通过 let metrics = Arc::new(jni_new_global_ref!(env, metrics_node)?); 将其转换为GlobalRef并持久化持有。对 Array[CometBatchIterator] 和 CometTaskMemoryManager 的处理也是基于同样的原理，这是跨语言编程中资源管理的常见模式。

[AFFILIATE_SLOT_1]

3. 构建DataFusion执行上下文（SessionContext）

DataFusion的所有执行都发生在一个 SessionContext 中。创建上下文是计划执行前的关键准备：

 let exec_context: Box = Box::new(ExecutionContext {
              id,
              task_attempt_id,
              spark_plan,
              partition_count: partition_count as usize,
              root_op: None,
              scans: vec![],
              input_sources,
              stream: None,
              metrics,
              metrics_update_interval,
              metrics_last_update_time: Instant::now(),
              poll_count_since_metrics_check: 0,
              plan_creation_time,
              session_ctx: Arc::new(session),
              debug_native,
              explain_native,
              memory_pool_config,
              tracing_enabled,
          });
Ok(Box::into_raw(exec_context) as i64)

prepare_datafusion_session_context 函数负责创建和配置 DataFusion SessionContext。它主要通过 Spark Context 来设置各种配置项。一个值得注意的细节是：它将DataFusion任务的并行度设置为Spark单个Task分配的CPU核心数，这实现了资源层面的对齐与优化。

此外，该函数还负责注册函数库：

• datafusion::functions_nested::register_all: 注册DataFusion内置函数。
• register_datafusion_spark_function: 注册与Spark语义兼容的函数。
• datafusion_comet_spark_expr::register_all_comet_functions: 注册Comet项目自定义的函数扩展。

最终，函数构造出 ExecutionContext 并返回。创建过程的耗时会被记录在 plan_creation_time 中。

四、所有权转移与指针返回：跨语言内存管理

计划创建的最后一步，涉及Rust与Java之间内存所有权的转移，这是保证内存安全且不泄漏的核心。

创建的 CometPlan 对象通常被包装在 Box 中（智能指针）。为了将对象“交给”Java层管理，Rust侧需要放弃所有权。这是通过 Box::into_raw(exec_context) 方法实现的：

• Box::into_raw 会消耗掉 Box，返回一个指向堆上数据的裸指针（*mut CometPlan）。
• 此后，Rust的所有权系统和析构器将不再管理这块内存，内存管理的责任转移给了接收该指针的Java代码。
• Java侧在后续（如Task结束时）必须通过JNI回调Rust的释放函数，将指针传回并使用 Box::from_raw 重新构造Box，由Rust的析构器安全释放内存。这是一个标准的“谁分配，谁释放”的跨语言模式。

至此，一个指向 ExecutionContext 的指针被返回给Java层。Rust Native物理执行计划构建完成，随时可以接受数据并开始高效的向量化执行。

[AFFILIATE_SLOT_2]

五、总结与最佳实践启示

通过对Apache DataFusion Comet中Rust Native物理计划创建过程的深度剖析，我们可以得到以下几点关键启示：

1. 清晰的职责边界：Comet架构完美诠释了“让专业的工具做专业的事”。Spark负责调度、资源管理与生态系统集成，而高性能计算则卸载给专为性能而生的Rust DataFusion引擎。
2. 高效的数据交换：Arrow IPC与列式内存格式是打破JVM与Native引擎性能壁垒的基石，实现了近乎零拷贝的数据传递。
3. 严谨的跨语言资源管理：从Protobuf序列化到GlobalRef持有对象，再到通过裸指针转移内存所有权，每一步都体现了对跨语言编程中内存安全与生命周期的周密考虑。
4. 配置与执行环境隔离：为每个Task独立创建DataFusion SessionContext，确保了执行环境的隔离性与配置的灵活性。

这种将现代高性能语言（如Rust）与成熟大数据框架（Spark）结合的模式，为大数据计算引擎的演进提供了重要思路。理解其底层机制，不仅能帮助开发者更好地使用和调试Comet，也为构建类似的高性能跨语言系统提供了宝贵的范本。