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

Commit

Permalink
Fixed nested FFI.
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao committed Jul 22, 2021
1 parent 8d7d08f commit a5035bd
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 16 deletions.
54 changes: 40 additions & 14 deletions arrow-pyarrow-integration-testing/tests/test_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,17 @@

class TestCase(unittest.TestCase):
def setUp(self):
self.old_allocated_rust = arrow_pyarrow_integration_testing.total_allocated_bytes()
self.old_allocated_rust = (
arrow_pyarrow_integration_testing.total_allocated_bytes()
)
self.old_allocated_cpp = pyarrow.total_allocated_bytes()

def tearDown(self):
# No leak of Rust
self.assertEqual(self.old_allocated_rust, arrow_pyarrow_integration_testing.total_allocated_bytes())
self.assertEqual(
self.old_allocated_rust,
arrow_pyarrow_integration_testing.total_allocated_bytes(),
)

# No leak of C++ memory
self.assertEqual(self.old_allocated_cpp, pyarrow.total_allocated_bytes())
Expand All @@ -46,6 +51,7 @@ def test_primitive_rust(self):
"""
Rust -> Python -> Rust
"""

def double(array):
array = array.to_pylist()
return pyarrow.array([x * 2 if x is not None else None for x in array])
Expand Down Expand Up @@ -102,17 +108,19 @@ def test_time32_python(self):
"""
Python -> Rust -> Python
"""
a = pyarrow.array([None, 1, 2], pyarrow.time32('s'))
a = pyarrow.array([None, 1, 2], pyarrow.time32("s"))
b = arrow_pyarrow_integration_testing.concatenate(a)
expected = pyarrow.array([None, 1, 2] + [None, 1, 2], pyarrow.time32('s'))
expected = pyarrow.array([None, 1, 2] + [None, 1, 2], pyarrow.time32("s"))
self.assertEqual(b, expected)

def test_list_array(self):
"""
Python -> Rust -> Python
"""
for _ in range(2):
a = pyarrow.array([[], None, [1, 2], [4, 5, 6]], pyarrow.list_(pyarrow.int64()))
a = pyarrow.array(
[[], None, [1, 2], [4, 5, 6]], pyarrow.list_(pyarrow.int64())
)
b = arrow_pyarrow_integration_testing.round_trip(a)

b.validate(full=True)
Expand All @@ -124,18 +132,36 @@ def test_struct_array(self):
Python -> Rust -> Python
"""
fields = [
('f1', pyarrow.int32()),
('f2', pyarrow.string()),
("f1", pyarrow.int32()),
("f2", pyarrow.string()),
]
a = pyarrow.array([
{"f1": 1, "f2": "a"},
None,
{"f1": 3, "f2": None},
{"f1": None, "f2": "d"},
{"f1": None, "f2": None},
], pyarrow.struct(fields))
a = pyarrow.array(
[
{"f1": 1, "f2": "a"},
None,
{"f1": 3, "f2": None},
{"f1": None, "f2": "d"},
{"f1": None, "f2": None},
],
pyarrow.struct(fields),
)
b = arrow_pyarrow_integration_testing.round_trip(a)

b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type

def test_list_list_array(self):
"""
Python -> Rust -> Python
"""
for _ in range(2):
a = pyarrow.array(
[[None], None, [[1], [2]], [[4, 5], [6]]],
pyarrow.list_(pyarrow.list_(pyarrow.int64())),
)
b = arrow_pyarrow_integration_testing.round_trip(a)

b.validate(full=True)
assert a.to_pylist() == b.to_pylist()
assert a.type == b.type
21 changes: 21 additions & 0 deletions src/ffi/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -168,4 +168,25 @@ mod tests {

test_round_trip(array)
}

#[test]
fn test_list_list() -> Result<()> {
let data = vec![
Some(vec![
Some(vec![None]),
Some(vec![Some(2)]),
Some(vec![Some(3)]),
]),
None,
Some(vec![Some(vec![Some(4), None, Some(6)])]),
];

let mut array =
MutableListArray::<i32, MutableListArray<i32, MutablePrimitiveArray<i32>>>::new();
array.try_extend(data)?;

let array: ListArray<i32> = array.into();

test_round_trip(array)
}
}
4 changes: 2 additions & 2 deletions src/ffi/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -644,8 +644,8 @@ impl<'a> ArrowArrayRef for ArrowArrayChild<'a> {
self.schema
}

fn child(&self, _: usize) -> Result<ArrowArrayChild> {
todo!()
fn child(&self, index: usize) -> Result<ArrowArrayChild> {
create_child(self.array, self.schema, self.parent.clone(), index)
}
}

Expand Down

0 comments on commit a5035bd

Please sign in to comment.