fix(ffi): handle null data buffers from empty arrays (#3276)

* test(python): validate roundtripping of empty arrays

* test: check more types

* test: parameterize rust side

* fix: also check for zero length buffers

* fix: handle null data buffers

* fix: only allow zero-length buffers to be null

Author: wjones127

arrow-pyarrow-integration-testing/src/lib.rs

@@ -20,6 +20,7 @@
 
 use std::sync::Arc;
 
+use arrow::array::new_empty_array;
 use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;
 
@@ -70,6 +71,13 @@ fn double_py(lambda: &PyAny, py: Python) -> PyResult<bool> {
     Ok(array == expected)
 }
 
+#[pyfunction]
+fn make_empty_array(datatype: PyArrowType<DataType>, py: Python) -> PyResult<PyObject> {
+    let array = new_empty_array(&datatype.0);
+
+    array.data().to_pyarrow(py)
+}
+
 /// Returns the substring
 #[pyfunction]
 fn substring(
@@ -134,6 +142,7 @@ fn round_trip_record_batch_reader(
 fn arrow_pyarrow_integration_testing(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(double))?;
     m.add_wrapped(wrap_pyfunction!(double_py))?;
+    m.add_wrapped(wrap_pyfunction!(make_empty_array))?;
     m.add_wrapped(wrap_pyfunction!(substring))?;
     m.add_wrapped(wrap_pyfunction!(concatenate))?;
     m.add_wrapped(wrap_pyfunction!(round_trip_type))?;

arrow-pyarrow-integration-testing/tests/test_sql.py

@@ -193,6 +193,38 @@ def test_time32_python():
     del b
     del expected
 
+
+@pytest.mark.parametrize("datatype", _supported_pyarrow_types, ids=str)
+def test_empty_array_python(datatype):
+    """
+    Python -> Rust -> Python
+    """
+    if datatype == pa.float16():
+        pytest.skip("Float 16 is not implemented in Rust")
+
+    a = pa.array([], datatype)
+    b = rust.round_trip_array(a)
+    b.validate(full=True)
+    assert a.to_pylist() == b.to_pylist()
+    assert a.type == b.type
+    del a
+    del b
+
+
+@pytest.mark.parametrize("datatype", _supported_pyarrow_types, ids=str)
+def test_empty_array_rust(datatype):
+    """
+    Rust -> Python
+    """
+    a = pa.array([], type=datatype)
+    b = rust.make_empty_array(datatype)
+    b.validate(full=True)
+    assert a.to_pylist() == b.to_pylist()
+    assert a.type == b.type
+    del a
+    del b
+
+
 def test_binary_array():
     """
     Python -> Rust -> Python

arrow/src/ffi.rs

@@ -123,7 +123,7 @@ use std::{
 use bitflags::bitflags;
 
 use crate::array::{layout, ArrayData};
-use crate::buffer::Buffer;
+use crate::buffer::{Buffer, MutableBuffer};
 use crate::datatypes::DataType;
 use crate::error::{ArrowError, Result};
 use crate::util::bit_util;
@@ -578,7 +578,7 @@ unsafe fn create_buffer(
     index: usize,
     len: usize,
 ) -> Option<Buffer> {
-    if array.buffers.is_null() {
+    if array.buffers.is_null() || array.n_buffers == 0 {
         return None;
     }
     let buffers = array.buffers as *mut *const u8;
@@ -657,13 +657,20 @@ pub trait ArrowArrayRef {
 
                 let len = self.buffer_len(index)?;
 
-                unsafe { create_buffer(self.owner().clone(), self.array(), index, len) }
-                    .ok_or_else(|| {
-                        ArrowError::CDataInterface(format!(
-                            "The external buffer at position {} is null.",
-                            index - 1
-                        ))
-                    })
+                match unsafe {
+                    create_buffer(self.owner().clone(), self.array(), index, len)
+                } {
+                    Some(buf) => Ok(buf),
+                    None if len == 0 => {
+                        // Null data buffer, which Rust doesn't allow. So create
+                        // an empty buffer.
+                        Ok(MutableBuffer::new(0).into())
+                    }
+                    None => Err(ArrowError::CDataInterface(format!(
+                        "The external buffer at position {} is null.",
+                        index - 1
+                    ))),
+                }
             })
             .collect()
     }