Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Simplified API to import from FFI #854

Merged
merged 1 commit into from
Feb 19, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion arrow-pyarrow-integration-testing/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,8 @@ fn to_rust_array(ob: PyObject, py: Python) -> PyResult<Arc<dyn Array>> {
)?;

let field = unsafe { ffi::import_field_from_c(schema.as_ref()).map_err(PyO3ArrowError::from)? };
let array = unsafe { ffi::import_array_from_c(array, &field).map_err(PyO3ArrowError::from)? };
let array =
unsafe { ffi::import_array_from_c(array, field.data_type).map_err(PyO3ArrowError::from)? };

Ok(array.into())
}
Expand Down
2 changes: 1 addition & 1 deletion examples/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ unsafe fn import(
schema: &ffi::Ffi_ArrowSchema,
) -> Result<Box<dyn Array>> {
let field = ffi::import_field_from_c(schema)?;
ffi::import_array_from_c(array, &field)
ffi::import_array_from_c(array, field.data_type)
}

fn main() -> Result<()> {
Expand Down
2 changes: 1 addition & 1 deletion src/array/binary/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ unsafe impl<O: Offset> ToFfi for BinaryArray<O> {

impl<O: Offset, A: ffi::ArrowArrayRef> FromFfi<A> for BinaryArray<O> {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();

let validity = unsafe { array.validity() }?;
let offsets = unsafe { array.buffer::<O>(1) }?;
Expand Down
2 changes: 1 addition & 1 deletion src/array/boolean/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ unsafe impl ToFfi for BooleanArray {

impl<A: ffi::ArrowArrayRef> FromFfi<A> for BooleanArray {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();
let validity = unsafe { array.validity() }?;
let values = unsafe { array.bitmap(1) }?;
Ok(Self::from_data(data_type, values, validity))
Expand Down
2 changes: 1 addition & 1 deletion src/array/fixed_size_binary/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ unsafe impl ToFfi for FixedSizeBinaryArray {

impl<A: ffi::ArrowArrayRef> FromFfi<A> for FixedSizeBinaryArray {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();
let validity = unsafe { array.validity() }?;
let values = unsafe { array.buffer::<u8>(1) }?;

Expand Down
2 changes: 1 addition & 1 deletion src/array/fixed_size_list/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ unsafe impl ToFfi for FixedSizeListArray {

impl<A: ffi::ArrowArrayRef> FromFfi<A> for FixedSizeListArray {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();
let validity = unsafe { array.validity() }?;
let child = unsafe { array.child(0)? };
let values = ffi::try_from(child)?.into();
Expand Down
2 changes: 1 addition & 1 deletion src/array/list/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ unsafe impl<O: Offset> ToFfi for ListArray<O> {

impl<O: Offset, A: ffi::ArrowArrayRef> FromFfi<A> for ListArray<O> {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();
let validity = unsafe { array.validity() }?;
let offsets = unsafe { array.buffer::<O>(1) }?;
let child = unsafe { array.child(0)? };
Expand Down
2 changes: 1 addition & 1 deletion src/array/map/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ unsafe impl ToFfi for MapArray {

impl<A: ffi::ArrowArrayRef> FromFfi<A> for MapArray {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();
let validity = unsafe { array.validity() }?;
let offsets = unsafe { array.buffer::<i32>(1) }?;
let child = array.child(0)?;
Expand Down
2 changes: 1 addition & 1 deletion src/array/null.rs
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ unsafe impl ToFfi for NullArray {

impl<A: ffi::ArrowArrayRef> FromFfi<A> for NullArray {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();
Ok(Self::from_data(data_type, array.array().len()))
}
}
2 changes: 1 addition & 1 deletion src/array/primitive/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ unsafe impl<T: NativeType> ToFfi for PrimitiveArray<T> {

impl<T: NativeType, A: ffi::ArrowArrayRef> FromFfi<A> for PrimitiveArray<T> {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();
let validity = unsafe { array.validity() }?;
let values = unsafe { array.buffer::<T>(1) }?;

Expand Down
2 changes: 1 addition & 1 deletion src/array/struct_/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ unsafe impl ToFfi for StructArray {

impl<A: ffi::ArrowArrayRef> FromFfi<A> for StructArray {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();
let fields = Self::get_fields(&data_type);

let validity = unsafe { array.validity() }?;
Expand Down
5 changes: 2 additions & 3 deletions src/array/union/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,8 @@ unsafe impl ToFfi for UnionArray {

impl<A: ffi::ArrowArrayRef> FromFfi<A> for UnionArray {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let field = array.field();
let data_type = field.data_type().clone();
let fields = Self::get_fields(field.data_type());
let data_type = array.data_type().clone();
let fields = Self::get_fields(&data_type);

let mut types = unsafe { array.buffer::<i8>(0) }?;
let offsets = if Self::is_sparse(&data_type) {
Expand Down
2 changes: 1 addition & 1 deletion src/array/utf8/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ unsafe impl<O: Offset> ToFfi for Utf8Array<O> {

impl<O: Offset, A: ffi::ArrowArrayRef> FromFfi<A> for Utf8Array<O> {
unsafe fn try_from_ffi(array: A) -> Result<Self> {
let data_type = array.field().data_type().clone();
let data_type = array.data_type().clone();
let validity = unsafe { array.validity() }?;
let offsets = unsafe { array.buffer::<O>(1) }?;
let values = unsafe { array.buffer::<u8>(2)? };
Expand Down
2 changes: 1 addition & 1 deletion src/ffi/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use crate::{array::*, datatypes::PhysicalType};
/// * the interface is not valid (e.g. a null pointer)
pub unsafe fn try_from<A: ArrowArrayRef>(array: A) -> Result<Box<dyn Array>> {
use PhysicalType::*;
Ok(match array.field().data_type().to_physical_type() {
Ok(match array.data_type().to_physical_type() {
Null => Box::new(NullArray::try_from_ffi(array)?),
Boolean => Box::new(BooleanArray::try_from_ffi(array)?),
Primitive(primitive) => with_match_primitive_type!(primitive, |$T| {
Expand Down
51 changes: 23 additions & 28 deletions src/ffi/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,9 @@ use crate::{
bytes::{Bytes, Deallocation},
Buffer,
},
datatypes::{DataType, Field, PhysicalType},
datatypes::{DataType, PhysicalType},
error::{ArrowError, Result},
ffi::schema::get_field_child,
ffi::schema::get_child,
types::NativeType,
};

Expand Down Expand Up @@ -317,32 +317,32 @@ fn buffer_len(array: &Ffi_ArrowArray, data_type: &DataType, i: usize) -> Result<

fn create_child(
array: &Ffi_ArrowArray,
field: &Field,
field: &DataType,
parent: Arc<ArrowArray>,
index: usize,
) -> Result<ArrowArrayChild<'static>> {
let field = get_field_child(field, index)?;
let data_type = get_child(field, index)?;
assert!(index < array.n_children as usize);
assert!(!array.children.is_null());
unsafe {
let arr_ptr = *array.children.add(index);
assert!(!arr_ptr.is_null());
let arr_ptr = &*arr_ptr;

Ok(ArrowArrayChild::from_raw(arr_ptr, field, parent))
Ok(ArrowArrayChild::from_raw(arr_ptr, data_type, parent))
}
}

fn create_dictionary(
array: &Ffi_ArrowArray,
field: &Field,
data_type: &DataType,
parent: Arc<ArrowArray>,
) -> Result<Option<ArrowArrayChild<'static>>> {
if let DataType::Dictionary(_, values, _) = field.data_type() {
let field = Field::new("", values.as_ref().clone(), true);
if let DataType::Dictionary(_, values, _) = data_type {
let data_type = values.as_ref().clone();
assert!(!array.dictionary.is_null());
let array = unsafe { &*array.dictionary };
Ok(Some(ArrowArrayChild::from_raw(array, field, parent)))
Ok(Some(ArrowArrayChild::from_raw(array, data_type, parent)))
} else {
Ok(None)
}
Expand Down Expand Up @@ -371,12 +371,7 @@ pub trait ArrowArrayRef: std::fmt::Debug {
/// The caller must guarantee that the buffer `index` corresponds to a bitmap.
/// This function assumes that the bitmap created from FFI is valid; this is impossible to prove.
unsafe fn buffer<T: NativeType>(&self, index: usize) -> Result<Buffer<T>> {
create_buffer::<T>(
self.array(),
self.field().data_type(),
self.deallocation(),
index,
)
create_buffer::<T>(self.array(), self.data_type(), self.deallocation(), index)
}

/// # Safety
Expand All @@ -390,18 +385,18 @@ pub trait ArrowArrayRef: std::fmt::Debug {
/// # Safety
/// The caller must guarantee that the child `index` is valid per c data interface.
unsafe fn child(&self, index: usize) -> Result<ArrowArrayChild> {
create_child(self.array(), self.field(), self.parent().clone(), index)
create_child(self.array(), self.data_type(), self.parent().clone(), index)
}

fn dictionary(&self) -> Result<Option<ArrowArrayChild>> {
create_dictionary(self.array(), self.field(), self.parent().clone())
create_dictionary(self.array(), self.data_type(), self.parent().clone())
}

fn n_buffers(&self) -> usize;

fn parent(&self) -> &Arc<ArrowArray>;
fn array(&self) -> &Ffi_ArrowArray;
fn field(&self) -> &Field;
fn data_type(&self) -> &DataType;
}

/// Struct used to move an Array from and to the C Data Interface.
Expand All @@ -426,19 +421,19 @@ pub trait ArrowArrayRef: std::fmt::Debug {
#[derive(Debug)]
pub struct ArrowArray {
array: Box<Ffi_ArrowArray>,
field: Field,
data_type: DataType,
}

impl ArrowArray {
pub fn new(array: Box<Ffi_ArrowArray>, field: Field) -> Self {
Self { array, field }
pub fn new(array: Box<Ffi_ArrowArray>, data_type: DataType) -> Self {
Self { array, data_type }
}
}

impl ArrowArrayRef for Arc<ArrowArray> {
/// the data_type as declared in the schema
fn field(&self) -> &Field {
&self.field
fn data_type(&self) -> &DataType {
&self.data_type
}

fn parent(&self) -> &Arc<ArrowArray> {
Expand All @@ -457,14 +452,14 @@ impl ArrowArrayRef for Arc<ArrowArray> {
#[derive(Debug)]
pub struct ArrowArrayChild<'a> {
array: &'a Ffi_ArrowArray,
field: Field,
data_type: DataType,
parent: Arc<ArrowArray>,
}

impl<'a> ArrowArrayRef for ArrowArrayChild<'a> {
/// the data_type as declared in the schema
fn field(&self) -> &Field {
&self.field
fn data_type(&self) -> &DataType {
&self.data_type
}

fn parent(&self) -> &Arc<ArrowArray> {
Expand All @@ -481,10 +476,10 @@ impl<'a> ArrowArrayRef for ArrowArrayChild<'a> {
}

impl<'a> ArrowArrayChild<'a> {
fn from_raw(array: &'a Ffi_ArrowArray, field: Field, parent: Arc<ArrowArray>) -> Self {
fn from_raw(array: &'a Ffi_ArrowArray, data_type: DataType, parent: Arc<ArrowArray>) -> Self {
Self {
array,
field,
data_type,
parent,
}
}
Expand Down
6 changes: 3 additions & 3 deletions src/ffi/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ pub(crate) use ffi::{ArrowArray, ArrowArrayRef};
use std::sync::Arc;

use crate::array::Array;
use crate::datatypes::Field;
use crate::datatypes::{DataType, Field};
use crate::error::Result;

pub use ffi::Ffi_ArrowArray;
Expand Down Expand Up @@ -50,7 +50,7 @@ pub unsafe fn import_field_from_c(field: &Ffi_ArrowSchema) -> Result<Field> {
/// valid according to the [C data interface](https://arrow.apache.org/docs/format/CDataInterface.html) (FFI).
pub unsafe fn import_array_from_c(
array: Box<Ffi_ArrowArray>,
field: &Field,
data_type: DataType,
) -> Result<Box<dyn Array>> {
try_from(Arc::new(ArrowArray::new(array, field.clone())))
try_from(Arc::new(ArrowArray::new(array, data_type)))
}
16 changes: 8 additions & 8 deletions src/ffi/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -455,14 +455,14 @@ fn to_format(data_type: &DataType) -> String {
}
}

pub(super) fn get_field_child(field: &Field, index: usize) -> Result<Field> {
match (index, field.data_type()) {
(0, DataType::List(field)) => Ok(field.as_ref().clone()),
(0, DataType::FixedSizeList(field, _)) => Ok(field.as_ref().clone()),
(0, DataType::LargeList(field)) => Ok(field.as_ref().clone()),
(0, DataType::Map(field, _)) => Ok(field.as_ref().clone()),
(index, DataType::Struct(fields)) => Ok(fields[index].clone()),
(index, DataType::Union(fields, _, _)) => Ok(fields[index].clone()),
pub(super) fn get_child(data_type: &DataType, index: usize) -> Result<DataType> {
match (index, data_type) {
(0, DataType::List(field)) => Ok(field.data_type().clone()),
(0, DataType::FixedSizeList(field, _)) => Ok(field.data_type().clone()),
(0, DataType::LargeList(field)) => Ok(field.data_type().clone()),
(0, DataType::Map(field, _)) => Ok(field.data_type().clone()),
(index, DataType::Struct(fields)) => Ok(fields[index].data_type().clone()),
(index, DataType::Union(fields, _, _)) => Ok(fields[index].data_type().clone()),
(child, data_type) => Err(ArrowError::OutOfSpec(format!(
"Requested child {} to type {:?} that has no such child",
child, data_type
Expand Down
3 changes: 2 additions & 1 deletion tests/it/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@ fn _test_round_trip(array: Arc<dyn Array>, expected: Box<dyn Array>) -> Result<(

// import references
let result_field = unsafe { ffi::import_field_from_c(schema_ptr.as_ref())? };
let result_array = unsafe { ffi::import_array_from_c(array_ptr, &result_field)? };
let result_array =
unsafe { ffi::import_array_from_c(array_ptr, result_field.data_type.clone())? };

assert_eq!(&result_array, &expected);
assert_eq!(result_field, field);
Expand Down