Skip to content

Commit

Permalink
teach vortex.encode to also accept nullable arrays (#697)
Browse files Browse the repository at this point in the history
This new test fails on develop:

        def test_array_with_nulls():
            a = pa.array([b'123', None])
    >       arr = vortex.encode(a)
E pyo3_runtime.PanicException: assertion failed: nulls.is_none()

It seems odd that `from_array` has a nullability parameter, but, in
principle, we could support an array which does not know its
nullability.
  • Loading branch information
danking authored Sep 5, 2024
1 parent 660b518 commit 92be80f
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 1 deletion.
3 changes: 2 additions & 1 deletion pyvortex/src/encode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@ pub fn encode<'py>(obj: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyArray>> {

if obj.is_instance(&pa_array)? {
let arrow_array = ArrowArrayData::from_pyarrow_bound(obj).map(make_array)?;
let enc_array = Array::from_arrow(arrow_array, false);
let is_nullable = arrow_array.is_nullable();
let enc_array = Array::from_arrow(arrow_array, is_nullable);
Bound::new(obj.py(), PyArray::new(enc_array))
} else if obj.is_instance(&chunked_array)? {
let chunks: Vec<Bound<PyAny>> = obj.getattr("chunks")?.extract()?;
Expand Down
6 changes: 6 additions & 0 deletions pyvortex/test/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,12 @@ def test_primitive_array_round_trip():
assert arr.to_arrow().combine_chunks() == a


def test_array_with_nulls():
a = pa.array([b"123", None])
arr = vortex.encode(a)
assert arr.to_arrow().combine_chunks() == a


def test_varbin_array_round_trip():
a = pa.array(["a", "b", "c"])
arr = vortex.encode(a)
Expand Down

0 comments on commit 92be80f

Please sign in to comment.