BUG: Slice Arrow buffer before passing it to numpy (pandas-dev#40896)

Add Arrow buffer slicing before handing it over to numpy which is
needed in case the Arrow buffer contains padding or offset.

Author: ThomasBlauthQC

pandas/core/arrays/_arrow_utils.py

@@ -14,6 +14,8 @@ def pyarrow_array_to_numpy_and_mask(arr, dtype):
     Convert a primitive pyarrow.Array to a numpy array and boolean mask based
     on the buffers of the Array.
 
+    At the moment pyarrow.BooleanArray is not supported.
+
     Parameters
     ----------
     arr : pyarrow.Array
@@ -25,8 +27,15 @@ def pyarrow_array_to_numpy_and_mask(arr, dtype):
         Tuple of two numpy arrays with the raw data (with specified dtype) and
         a boolean mask (validity mask, so False means missing)
     """
+
     buflist = arr.buffers()
-    data = np.frombuffer(buflist[1], dtype=dtype)[arr.offset : arr.offset + len(arr)]
+    # Since Arrow buffers might contain padding and the data might be offset,
+    # the buffer gets sliced here before handing it to numpy.
+    # See also https://github.com/pandas-dev/pandas/issues/40896
+    offset = arr.offset * arr.type.bit_width // 8
+    length = len(arr) * arr.type.bit_width // 8
+    data_buf = buflist[1][offset : offset + length]
+    data = np.frombuffer(data_buf, dtype=dtype)
     bitmask = buflist[0]
     if bitmask is not None:
         mask = pyarrow.BooleanArray.from_buffers(

pandas/tests/arrays/test_arrow_utils.py

@@ -0,0 +1,58 @@
+import numpy as np
+
+import pandas.util._test_decorators as td
+
+import pandas._testing as tm
+
+
+@td.skip_if_no("pyarrow")
+def test_pyarrow_array_to_numpy_and_mask(dtype, pa_array, np_expected, mask_expected):
+    """
+    Test conversion from pyarrow array to numpy array.
+
+    Also modifies the pyarrow buffer to contain padding and offset, which are
+    considered valid buffers by pyarrow.
+    See also https://github.com/pandas-dev/pandas/issues/40896
+    """
+    import pyarrow as pa
+
+    from pandas.core.arrays._arrow_utils import pyarrow_array_to_numpy_and_mask
+
+    dtype = np.int32
+    pa_array = pa.array([0, 1, 2], type=pa.int32())
+    np_expected = np.array([0, 1, 2], dtype=np.int32)
+    mask_expected = np.array([True, True, True])
+
+    data, mask = pyarrow_array_to_numpy_and_mask(pa_array, dtype)
+    tm.assert_numpy_array_equal(data, np_expected)
+    assert (mask == mask_expected).all()
+
+    mask_buffer = pa_array.buffers()[0]
+    data_buffer = pa_array.buffers()[1].to_pybytes()
+
+    # Add trailing padding to the buffer.
+    data_buffer_trail = pa.py_buffer(data_buffer + b"\x00")
+    pa_array_trail = pa.Array.from_buffers(
+        type=pa_array.type,
+        length=len(pa_array),
+        buffers=[mask_buffer, data_buffer_trail],
+        offset=pa_array.offset,
+    )
+    pa_array_trail.validate()
+    data, mask = pyarrow_array_to_numpy_and_mask(pa_array_trail, dtype)
+    tm.assert_numpy_array_equal(data, np_expected)
+    assert (mask == mask_expected).all()
+
+    # Add offset to the buffer.
+    offset = b"\x00" * (pa_array.type.bit_width // 8)
+    data_buffer_offset = pa.py_buffer(offset + data_buffer)
+    pa_array_offset = pa.Array.from_buffers(
+        type=pa_array.type,
+        length=len(pa_array),
+        buffers=[mask_buffer, data_buffer_offset],
+        offset=pa_array.offset + 1,
+    )
+    pa_array_offset.validate()
+    data, mask = pyarrow_array_to_numpy_and_mask(pa_array_offset, dtype)
+    tm.assert_numpy_array_equal(data, np_expected)
+    assert (mask == mask_expected).all()