Skip to content

Commit c863a2c

Browse files
authored
Fix FFI and add support for Struct type (#287)
* fix: support nested types in FFI Ported from https://github.com/jorgecarleitao/arrow2 Fix #20 Fix #251 Signed-off-by: roee88 <roee88@gmail.com> * Removed Clone from FFI_ArrowArray Signed-off-by: roee88 <roee88@gmail.com> * Add nesting to FFI struct test Signed-off-by: roee88 <roee88@gmail.com>
1 parent de44c8c commit c863a2c

File tree

2 files changed

+387
-341
lines changed

2 files changed

+387
-341
lines changed

arrow/src/array/ffi.rs

Lines changed: 42 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -22,111 +22,40 @@ use std::convert::TryFrom;
2222
use crate::{
2323
error::{ArrowError, Result},
2424
ffi,
25+
ffi::ArrowArrayRef,
2526
};
2627

2728
use super::ArrayData;
28-
use crate::datatypes::DataType;
29-
use crate::ffi::ArrowArray;
3029

3130
impl TryFrom<ffi::ArrowArray> for ArrayData {
3231
type Error = ArrowError;
3332

3433
fn try_from(value: ffi::ArrowArray) -> Result<Self> {
35-
let child_data = value.children()?;
36-
37-
let child_type = if !child_data.is_empty() {
38-
Some(child_data[0].data_type().clone())
39-
} else {
40-
None
41-
};
42-
43-
let data_type = value.data_type(child_type)?;
44-
45-
let len = value.len();
46-
let offset = value.offset();
47-
let null_count = value.null_count();
48-
let buffers = value.buffers()?;
49-
let null_bit_buffer = value.null_bit_buffer();
50-
51-
Ok(ArrayData::new(
52-
data_type,
53-
len,
54-
Some(null_count),
55-
null_bit_buffer,
56-
offset,
57-
buffers,
58-
child_data,
59-
))
34+
value.to_data()
6035
}
6136
}
6237

6338
impl TryFrom<ArrayData> for ffi::ArrowArray {
6439
type Error = ArrowError;
6540

6641
fn try_from(value: ArrayData) -> Result<Self> {
67-
// If parent is nullable, then children also must be nullable
68-
// so we pass this nullable to the creation of hte child data
69-
let nullable = match value.data_type() {
70-
DataType::List(field) => field.is_nullable(),
71-
DataType::LargeList(field) => field.is_nullable(),
72-
_ => false,
73-
};
74-
75-
let len = value.len();
76-
let offset = value.offset() as usize;
77-
let null_count = value.null_count();
78-
let buffers = value.buffers().to_vec();
79-
let null_buffer = value.null_buffer().cloned();
80-
let child_data = value
81-
.child_data()
82-
.iter()
83-
.map(|arr| {
84-
let len = arr.len();
85-
let offset = arr.offset() as usize;
86-
let null_count = arr.null_count();
87-
let buffers = arr.buffers().to_vec();
88-
let null_buffer = arr.null_buffer().cloned();
89-
90-
// Note: the nullable comes from the parent data.
91-
unsafe {
92-
ArrowArray::try_new(
93-
arr.data_type(),
94-
len,
95-
null_count,
96-
null_buffer,
97-
offset,
98-
buffers,
99-
vec![],
100-
nullable,
101-
)
102-
.expect("infallible")
103-
}
104-
})
105-
.collect::<Vec<_>>();
106-
107-
unsafe {
108-
ffi::ArrowArray::try_new(
109-
value.data_type(),
110-
len,
111-
null_count,
112-
null_buffer,
113-
offset,
114-
buffers,
115-
child_data,
116-
nullable,
117-
)
118-
}
42+
unsafe { ffi::ArrowArray::try_new(value) }
11943
}
12044
}
12145

12246
#[cfg(test)]
12347
mod tests {
12448
use crate::error::Result;
12549
use crate::{
126-
array::{Array, ArrayData, Int64Array, UInt32Array, UInt64Array},
50+
array::{
51+
Array, ArrayData, BooleanArray, Int64Array, StructArray, UInt32Array,
52+
UInt64Array,
53+
},
54+
datatypes::{DataType, Field},
12755
ffi::ArrowArray,
12856
};
12957
use std::convert::TryFrom;
58+
use std::sync::Arc;
13059

13160
fn test_round_trip(expected: &ArrayData) -> Result<()> {
13261
// create a `ArrowArray` from the data.
@@ -165,4 +94,37 @@ mod tests {
16594
let data = array.data();
16695
test_round_trip(data)
16796
}
97+
98+
#[test]
99+
fn test_struct() -> Result<()> {
100+
let inner = StructArray::from(vec![
101+
(
102+
Field::new("a1", DataType::Boolean, false),
103+
Arc::new(BooleanArray::from(vec![true, true, false, false]))
104+
as Arc<dyn Array>,
105+
),
106+
(
107+
Field::new("a2", DataType::UInt32, false),
108+
Arc::new(UInt32Array::from(vec![1, 2, 3, 4])),
109+
),
110+
]);
111+
112+
let array = StructArray::from(vec![
113+
(
114+
Field::new("a", inner.data_type().clone(), false),
115+
Arc::new(inner) as Arc<dyn Array>,
116+
),
117+
(
118+
Field::new("b", DataType::Boolean, false),
119+
Arc::new(BooleanArray::from(vec![false, false, true, true]))
120+
as Arc<dyn Array>,
121+
),
122+
(
123+
Field::new("c", DataType::UInt32, false),
124+
Arc::new(UInt32Array::from(vec![42, 28, 19, 31])),
125+
),
126+
]);
127+
let data = array.data();
128+
test_round_trip(data)
129+
}
168130
}

0 commit comments

Comments
 (0)