背景
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调用对应的 JavaCometBatchIterator 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，便于进行数据拉取
-
3万+
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
