Spark DataFusion Comet 向量化：Rust Native ScanExec 与 Selection Vectors

allocate_and_fetch_batch 方法及后续说明

调用 allocate_and_fetch_batch 方法从 Java 端获取一批数据并赋值到传入的 FFI_ArrowArray 和 FFI_ArrowSchema 中。对于 Java 端的处理和之前的 exportSingleVector 方法处理一样，主要是使用 Data.exportVector 方法进行数据回传。

如果存在 Selection Vector，则使用 Arrow take 方法获取到真正的值，否则就是原值。其中对 &Arc 使用 &* 操作是将类型由引用 &Arc 依次转换为 Arc、T，最后再取引用 &T，用于不转移所有权地读取数据。

如果存在 Selection Vector，则返回真正的行数。最后组装成 InputBatch::new(inputs, Some(actual_num_rows)) 返回。

JNI 调用 JVM 的 exportSelectionIndices 方法

let _exported_count: i32 = unsafe { jni_call!(env, comet_batch_iterator(iter).export_selection_indices( JValueGen::Object(JObject::from(indices_array_obj).as_ref()), JValueGen::Object(JObject::from(indices_schema_obj).as_ref()) ) -> i32)? };

这里用到的 as_ref 方法通常用于将高级包装类型（如 JObject, JString 等）转换为对底层 JNI 指针（jobject）的共享引用。它使得在不改变对象所有权的情况下，可以安全地将对象传递给 JNIEnv 函数进行后续操作。

在 Java 侧的话，主要是 NativeUtil.exportSingleVector 的调用：

def exportSingleVector(vector: CometVector, arrayAddr: Long, schemaAddr: Long): Unit = {
  val valueVector = vector.getValueVector
  val provider = if (valueVector.getField.getDictionary != null) {
    vector.getDictionaryProvider
  } else {
    null
  }
  val arrowSchema = ArrowSchema.wrap(schemaAddr)
  val arrowArray = ArrowArray.wrap(arrayAddr)
  Data.exportVector(allocator, getFieldVector(valueVector, "export"), provider, arrowArray, arrowSchema)
}

Data.exportVector 使用这个方法使 Selection Vector 回传到 Rust 端。

判断是否存在 Selection Vectors，通过 JNI 调用 java 侧方法 hasSelectionVectors：

jni_call!(env, comet_batch_iterator(iter).has_selection_vectors() -> jni::sys::jboolean)

get_selection_indices 方法说明

fn get_selection_indices(
  env: &mut jni::JNIEnv,
  iter: &JObject,
  num_cols: usize,
) -> Result<Option<Vec<ArrayRef>>, CometError> {
  // Check if all columns have selection vectors
  let has_selection_vectors_result: jni::sys::jboolean = unsafe {
    jni_call!(env, comet_batch_iterator(iter).has_selection_vectors() -> jni::sys::jboolean)?
  };
  let has_selection_vectors = has_selection_vectors_result != 0;
  let selection_indices_arrays = if has_selection_vectors {
    // Allocate arrays for selection indices export (one per column)
    let mut indices_array_addrs = Vec::with_capacity(num_cols);
    let mut indices_schema_addrs = Vec::with_capacity(num_cols);
    for _ in 0..num_cols {
      let arrow_array = Rc::new(FFI_ArrowArray::empty());
      let arrow_schema = Rc::new(FFI_ArrowSchema::empty());
      indices_array_addrs.push(Rc::into_raw(arrow_array) as i64);
      indices_schema_addrs.push(Rc::into_raw(arrow_schema) as i64);
    }
    // Prepare JNI arrays for the export call
    let indices_array_obj = env.new_long_array(num_cols as jsize)?;
    let indices_schema_obj = env.new_long_array(num_cols as jsize)?;
    env.set_long_array_region(&indices_array_obj, 0, &indices_array_addrs)?;
    env.set_long_array_region(&indices_schema_obj, 0, &indices_schema_addrs)?;
    // Export selection indices from JVM
    let _exported_count: i32 = unsafe {
      jni_call!(env, comet_batch_iterator(iter).export_selection_indices(
        JValueGen::Object(JObject::from(indices_array_obj).as_ref()),
        JValueGen::Object(JObject::from(indices_schema_obj).as_ref())
      ) -> i32)?
    };
    // Convert to ArrayRef for easier handling
    let mut selection_arrays = Vec::with_capacity(num_cols);
    for i in 0..num_cols {
      let array_data = ArrayData::from_spark((indices_array_addrs[i], indices_schema_addrs[i]))?;
      selection_arrays.push(make_array(array_data));
      // Drop the references to the FFI arrays
      unsafe {
        Rc::from_raw(indices_array_addrs[i] as *const FFI_ArrowArray);
        Rc::from_raw(indices_schema_addrs[i] as *const FFI_ArrowSchema);
      }
    }
    Some(selection_arrays)
  } else {
    None
  };
  Ok(selection_indices_arrays)
}

Rust 侧处理逻辑

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)))

indices 为 Java 数组 selectionIndices 以 Arrow vector 的表示，便于其他语言能够以零拷贝的方式访问这些数据，后续会被传递给 Native(Rust) 层。

this.indices = CometVector.getVector(indicesVector, values.useDecimal128, values.getDictionaryProvider());

Java 侧实现

public boolean hasSelectionVectors() {
  if (currentBatch == null) {
    return false;
  }
  // Check if all columns are CometSelectionVector instances
  for (int i = 0; i < currentBatch.numCols(); i++) {
    if (!(currentBatch.column(i) instanceof CometSelectionVector)) {
      return false;
    }
  }
  return true;
}

这其中的 CometSelectionVector 就是对应上文中说到的 selection vector，具体的代码如下：

public class CometSelectionVector extends CometVector {
  /** The original vector containing all values */
  private final CometVector values;
  /**
   * The valid indices in the values vector. This array is converted into an Arrow vector so we can
   * transfer the data to native in one JNI call. This is used to represent the rowid mapping used
   * by Iceberg
   */
  private final int[] selectionIndices;
  /**
   * The indices vector containing selection indices. This is currently allocated by the JVM side
   * unlike the values vector which is allocated on the native side
   */
  private final CometVector indices;
  // ...
}