From 92be80ff9890ebe0e5b6f09e724c32eac6592d45 Mon Sep 17 00:00:00 2001 From: Dan King Date: Thu, 5 Sep 2024 12:24:45 -0400 Subject: [PATCH] teach vortex.encode to also accept nullable arrays (#697) 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. --- pyvortex/src/encode.rs | 3 ++- pyvortex/test/test_array.py | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/pyvortex/src/encode.rs b/pyvortex/src/encode.rs index 838ca044aa..850fc45563 100644 --- a/pyvortex/src/encode.rs +++ b/pyvortex/src/encode.rs @@ -24,7 +24,8 @@ pub fn encode<'py>(obj: &Bound<'py, PyAny>) -> PyResult> { 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> = obj.getattr("chunks")?.extract()?; diff --git a/pyvortex/test/test_array.py b/pyvortex/test/test_array.py index db93368baf..8acf7ab2fd 100644 --- a/pyvortex/test/test_array.py +++ b/pyvortex/test/test_array.py @@ -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)