Spark Datafusion Comet 向量化Rust Native--Native算子(CometNativeExec)从JVM中获取RecordBatch数据

背景

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的Spark Datafusion Comet 向量化Rust Native–执行Datafusion计划中涉及到的获取RecordBatch数据细节实现

执行计划planner.create_plan，拉取数据

这里的方法会返回scans, root_op,如下:

let (scans, root_op) = planner.create_plan(
&exec_context.spark_plan,
&mut exec_context.input_sources.clone(),
exec_context.partition_count,
)?;

从JVM传过来的参数可知，该input_sources 是所有该物理计划的输入，该方法会把 protobuf Operator一一映射为DataFusion physical plan,这里以Scan为例子(其他算子依次类推)：

OpStruct::Scan(scan) => {
let data_types = scan.fields.iter().map(to_arrow_datatype).collect_vec();

// If it is not test execution context for unit test, we should have at least one
// input source
if self.exec_context_id != TEST_EXEC_CONTEXT_ID && inputs.is_empty() {
return Err(GeneralError("No input for scan".to_string()));
}

// Consumes the first input source for the scan
let input_source =
if self.exec_context_id == TEST_EXEC_CONTEXT_ID && inputs.is_empty() {
// For unit test, we will set input batch to scan directly by `set_input_batch`.
None
} else {
Some(inputs.remove(0))
};

// The `ScanExec` operator will take actual arrays from Spark during execution
let scan = ScanExec::new(
self.exec_context_id,
input_source,
&scan.source,
data_types,
scan.arrow_ffi_safe,
)?;

Ok((
vec![scan.clone()],
Arc::new(SparkPlan::new(spark_plan.plan_id, Arc::new(scan), vec![])),
)) }

• to_arrow_datatype依次映射数据类型到dataFusion对应的类型
• inputs.remove(0)选择第一个输入的RDD为ScanExec和Scans 返回一个(scans,ScanExec)元组类型

注意这里的 ScanExec实现了get_next_batch以及ExecutionPlan（来自DataFusion的物理计划）的execute方法：

pub fn get_next_batch(&mut self) -> Result<(), CometError> {
if self.input_source.is_none() {
// This is a unit test. We don't need to call JNI.
return Ok(());
}
let mut timer = self.baseline_metrics.elapsed_compute().timer();

let mut current_batch = self.batch.try_lock().unwrap();
if current_batch.is_none() {
let next_batch = ScanExec::get_next(
self.exec_context_id,
self.input_source.as_ref().unwrap().as_obj(),
self.data_types.len(),
self.arrow_ffi_safe,
)?;
*current_batch = Some(next_batch);
}

timer.stop();

Ok(())
}
…
fn execute(
&self,
partition: usize,
_: Arc<TaskContext>,
) -> datafusion::common::Result<SendableRecordBatchStream> {
Ok(Box::pin(ScanStream::new(
self.clone(),
self.schema(),
partition,
self.baseline_metrics.clone(),
)))
}

注意：这两个方法也是相互配合使用，execute获取对应的RecordBatch Stream数据（loop循环判断调用），get_next_batch会填充RecordBatch Stream数据 并且DataFusion物理计划execute方法用到的RecordBatch数据是通过后续JNI调用JVM的方法获取到的 这 execute返回的ScanStream类型的对象中的build_record_batch方法，会使用Arrow中cast_with_options进行字典解码，组装成native_plan.execute返回的RecordBatch结果,具体的可以看Function cast_with_options Copy item path

而这里的get_next_batch方法将会被pull_input_batches方法调用，这个方法调用JNI从JVM端获取数据：

fn pull_input_batches(exec_context: &mut ExecutionContext) -> Result<(), CometError> {
exec_context.scans.iter_mut().try_for_each(|scan| {
scan.get_next_batch()?;
Ok::<(), CometError>(())
})
}
…
/// Pull next input batch from JVM.
pub fn get_next_batch(&mut self) -> Result<(), CometError> {
if self.input_source.is_none() {
// This is a unit test. We don't need to call JNI.
return Ok(());
}
let mut timer = self.baseline_metrics.elapsed_compute().timer();

let mut current_batch = self.batch.try_lock().unwrap();
if current_batch.is_none() {
let next_batch = ScanExec::get_next(
self.exec_context_id,
self.input_source.as_ref().unwrap().as_obj(),
self.data_types.len(),
self.arrow_ffi_safe,
)?;
*current_batch = Some(next_batch);
}

timer.stop();

Ok(())
}

• timer为计算此次从JVM中拉取数据的耗时
• 调用ScanExec::get_next方法获取对应的批次数据/// Invokes JNI call to get next batch.
  fn get_next(
  exec_context_id: i64,
  iter: &JObject,
  num_cols: usize,
  arrow_ffi_safe: bool,
  ) -> Result<InputBatch, CometError> {
  …
  let mut env = JVMClasses::get_env()?;

  let num_rows: i32 = unsafe {
  jni_call!(&mut env,
  comet_batch_iterator(iter).has_next() -> i32)?
  };

  if num_rows == -1 {
  return Ok(InputBatch::EOF);
  }

  // Check for selection vectors and get selection indices if needed from
  // JVM via FFI
  // Selection vectors can be provided by, for instance, Iceberg to
  // remove rows that have been deleted.
  let selection_indices_arrays = Self::get_selection_indices(&mut env, iter, num_cols)?;

  // fetch batch data from JVM via FFI
  let (num_rows, array_addrs, schema_addrs) =
  Self::allocate_and_fetch_batch(&mut env, iter, num_cols)?;

  let mut inputs: Vec<ArrayRef> = Vec::with_capacity(num_cols);

  // Process each column
  for i in 0..num_cols {
  let array_ptr = array_addrs[i];
  let schema_ptr = schema_addrs[i];
  let array_data = ArrayData::from_spark((array_ptr, schema_ptr))?;

  // TODO: validate array input data
  // array_data.validate_full()?;

  let array = make_array(array_data);

  // Apply selection if selection vectors exist (applies to all columns)
  let array = if let Some(ref selection_arrays) = selection_indices_arrays {
  let indices = &selection_arrays[i];
  // Apply the selection using Arrow's take kernel
  match take(&*array, &**indices, None) {
  Ok(selected_array) => selected_array,
  Err(e) => {
  return Err(CometError::from(ExecutionError::ArrowError(format!(
  "Failed to apply selection for column {i}: {e}",
  ))));
  }
  }
  } else {
  array
  };

  let array = if arrow_ffi_safe {
  // ownership of this array has been transferred to native
  // but we still need to unpack dictionary arrays
  copy_or_unpack_array(&array, &CopyMode::UnpackOrClone)?
  } else {
  // it is necessary to copy the array because the contents may be
  // overwritten on the JVM side in the future
  copy_array(&array)
  };

  inputs.push(array);

  // Drop the Arcs to avoid memory leak
  unsafe {
  Rc::from_raw(array_ptr as *const FFI_ArrowArray);
  Rc::from_raw(schema_ptr as *const FFI_ArrowSchema);
  }
  }

  // If selection was applied, determine the actual row count from the selected arrays
  let actual_num_rows = if let Some(ref selection_arrays) = selection_indices_arrays {
  if !selection_arrays.is_empty() {
  // Use the length of the first selection array as the actual row count
  selection_arrays[0].len()
  } else {
  num_rows as usize
  }
  } else {
  num_rows as usize
  };
  Ok(InputBatch::new(inputs, Some(actual_num_rows)))
  }

  • JVMClasses::get_env() 得到JavaVM对象，并调用attach_current_thread方法将当前Rust线程关联到JVM
  • 通过JNI调用对应的 Java CometBatchIterator hasNext方法，获取下一批返回的数据行数 其中，有宏jni_call的定义，如下: macro_rules! jni_call {
    ($env:expr, $clsname:ident($obj:expr).$method:ident($($args:expr),* $(,)?) -> $ret:ty) => {{
    let method_id = paste::paste! {
    $crate::jvm_bridge::JVMClasses::get().[<$clsname>].[<method_ $method>]
    };
    let ret_type = paste::paste! {
    $crate::jvm_bridge::JVMClasses::get().[<$clsname>].[<method_ $method _ret>]
    }.clone();
    let args = $crate::jvm_bridge::jvalues!($($args,)*);

    // Call the JVM method and obtain the returned value
    let ret = $env.call_method_unchecked($obj, method_id, ret_type, args);

    // Check if JVM has thrown any exception, and handle it if so.
    let result = if let Some(exception) = $crate::jvm_bridge::check_exception($env)? {
    Err(exception.into())
    } else {
    $crate::jvm_bridge::jni_map_error!($env, ret)
    };

    result.and_then(|result| $crate::jvm_bridge::jni_map_error!($env, <$ret>::try_from(result)))
    }}
    }
    这里还会用到paste::paste!宏，该宏使用说明见这里，主要是它允许在 [<…>] 内的标识符被拼接成单个标识符 该此调用的Java方法都被提前封装在JVMClasses中

  • 如果返回的行数为-1,则说明已经没有数据可以读取了，直接返回InputBatch::EOF
  • get_selection_indices 通过JNI获取到 selection indices，这里暂且跳过，后续会说明
  • allocate_and_fetch_batch通过 Arrow FFI 结构以及JNI获取到批次数据
  • 对于每一列应用selection_indices_arrays来获取真正的数据，这个后续会说明
• 赋值当前的批次数据给ScanExec.batch作为给exec_context.stream中的最终数据返回给JVM端 这里的以ScanExec计划为例子(该execute方法返回的是ScanStream)： impl Stream for ScanStream<'_> {
  type Item = DataFusionResult<RecordBatch>;

  fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
  let mut timer = self.baseline_metrics.elapsed_compute().timer();
  let mut scan_batch: std::sync::MutexGuard<'_, Option<InputBatch>> = self.scan.batch.try_lock().unwrap();
  let input_batch = &*scan_batch;
  let input_batch = if let Some(batch) = input_batch {
  batch
  } else {
  timer.stop();
  return Poll::Pending;
  };
  let result = match input_batch {
  InputBatch::EOF => Poll::Ready(None),
  InputBatch::Batch(columns, num_rows) => {
  self.baseline_metrics.record_output(*num_rows);
  let maybe_batch = self.build_record_batch(columns, *num_rows);
  Poll::Ready(Some(maybe_batch))
  }
  };
  *scan_batch = None;
  timer.stop();
  result
  }
  }

  • 记录本次拉取所有批次数据的耗时

    let mut timer = self.baseline_metrics.elapsed_compute().timer();

  • 初始的 batch值为None,所以根据

    let input_batch = &*scan_batch;
    let input_batch = if let Some(batch) = input_batch {
    batch
    } else {
    timer.stop();
    return Poll::Pending;
    };

    这里判断返回Poll::Pending,所以在主流程中会继续调用pull_input_batches拉取数据

  • 如果返回的是InputBatch::EOF这表明数据已经处理完了,则返回Poll::Ready(None)，则在主流程中返回-1

  • 如果存在数据，则调用build_record_batch方法返回Poll::Ready(Some(batch))批数据,在主流程中获取拉取的数据，并返回给JVM中

  • 并设置batch为None，便于进行数据拉取