[arrow-rs] branch master updated: fix(ffi): handle null data buffers from empty arrays (#3276)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 2e806b0c6 fix(ffi): handle null data buffers from empty arrays (#3276)
2e806b0c6 is described below

commit 2e806b0c634e7c5487811c28329b8ef4c5f1ba77
Author: Will Jones <willjones...@gmail.com>
AuthorDate: Thu Dec 8 02:07:01 2022 -0800

    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
---
 arrow-pyarrow-integration-testing/src/lib.rs       |  9 ++++++
 .../tests/test_sql.py                              | 32 ++++++++++++++++++++++
 arrow/src/ffi.rs                                   | 25 +++++++++++------
 3 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/arrow-pyarrow-integration-testing/src/lib.rs 
b/arrow-pyarrow-integration-testing/src/lib.rs
index 2e74f0cf6..cf94b0dd4 100644
--- a/arrow-pyarrow-integration-testing/src/lib.rs
+++ b/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))?;
diff --git a/arrow-pyarrow-integration-testing/tests/test_sql.py 
b/arrow-pyarrow-integration-testing/tests/test_sql.py
index a19edf0cc..5a8bec792 100644
--- a/arrow-pyarrow-integration-testing/tests/test_sql.py
+++ b/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
diff --git a/arrow/src/ffi.rs b/arrow/src/ffi.rs
index fc8dc654a..0c1c1fa54 100644
--- a/arrow/src/ffi.rs
+++ b/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()
     }