fix pyarrow interchange

Author: MarcoGorelli

doc/source/whatsnew/v2.2.2.rst

@@ -14,14 +14,16 @@ including other versions of pandas.
 Fixed regressions
 ~~~~~~~~~~~~~~~~~
 - :meth:`DataFrame.__dataframe__` was producing incorrect data buffers when the a column's type was a pandas nullable on with missing values (:issue:`56702`)
+- :meth:`DataFrame.__dataframe__` was producing incorrect data buffers when the a column's type was a pyarrow nullable on with missing values (:issue:`57664`)
 -
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_222.bug_fixes:
 
 Bug fixes
 ~~~~~~~~~
--
+- :meth:`DataFrame.__dataframe__` was showing bytemask instead of bitmask for ``'string[pyarrow]'`` validity buffer (:issue:`57762`)
+- :meth:`DataFrame.__dataframe__` was showing non-null validity buffer (instead of ``None``) ``'string[pyarrow]'`` without missing values (:issue:`57761`)
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_222.other:

pandas/core/interchange/buffer.py

@@ -12,6 +12,7 @@
 
 if TYPE_CHECKING:
     import numpy as np
+    import pyarrow as pa
 
 
 class PandasBuffer(Buffer):
@@ -76,3 +77,74 @@ def __repr__(self) -> str:
             )
             + ")"
         )
+
+
+class PandasBufferPyarrow(Buffer):
+    """
+    Data in the buffer is guaranteed to be contiguous in memory.
+    """
+
+    def __init__(
+        self,
+        chunked_array: pa.ChunkedArray,
+        *,
+        is_validity: bool,
+        allow_copy: bool = True,
+    ) -> None:
+        """
+        Handle pyarrow chunked arrays.
+        """
+        if len(chunked_array.chunks) == 1:
+            arr = chunked_array.chunks[0]
+        else:
+            if not allow_copy:
+                raise RuntimeError(
+                    "Found multi-chunk pyarrow array, but `allow_copy` is False"
+                )
+            arr = chunked_array.combine_chunks()
+        if is_validity:
+            self._buffer = arr.buffers()[0]
+        else:
+            self._buffer = arr.buffers()[1]
+        self._length = len(arr)
+        self._dlpack = arr.__dlpack__
+        self._is_validity = is_validity
+
+    @property
+    def bufsize(self) -> int:
+        """
+        Buffer size in bytes.
+        """
+        return self._buffer.size
+
+    @property
+    def ptr(self) -> int:
+        """
+        Pointer to start of the buffer as an integer.
+        """
+        return self._buffer.address
+
+    def __dlpack__(self) -> Any:
+        """
+        Represent this structure as DLPack interface.
+        """
+        return self._dlpack()
+
+    def __dlpack_device__(self) -> tuple[DlpackDeviceType, int | None]:
+        """
+        Device type and device ID for where the data in the buffer resides.
+        """
+        return (DlpackDeviceType.CPU, None)
+
+    def __repr__(self) -> str:
+        return (
+            "PandasBuffer[pyarrow]("
+            + str(
+                {
+                    "bufsize": self.bufsize,
+                    "ptr": self.ptr,
+                    "device": self.__dlpack_device__()[0].name,
+                }
+            )
+            + ")"
+        )

pandas/core/interchange/column.py

@@ -9,15 +9,18 @@
 from pandas.errors import NoBufferPresent
 from pandas.util._decorators import cache_readonly
 
-from pandas.core.dtypes.dtypes import (
+from pandas.core.dtypes.dtypes import BaseMaskedDtype
+
+import pandas as pd
+from pandas import (
     ArrowDtype,
-    BaseMaskedDtype,
     DatetimeTZDtype,
 )
-
-import pandas as pd
 from pandas.api.types import is_string_dtype
-from pandas.core.interchange.buffer import PandasBuffer
+from pandas.core.interchange.buffer import (
+    PandasBuffer,
+    PandasBufferPyarrow,
+)
 from pandas.core.interchange.dataframe_protocol import (
     Column,
     ColumnBuffers,
@@ -157,6 +160,9 @@ def _dtype_from_pandasdtype(self, dtype) -> tuple[DtypeKind, int, str, str]:
         else:
             byteorder = dtype.byteorder
 
+        if dtype == "bool[pyarrow]":
+            return kind, dtype.itemsize, ArrowCTypes.BOOL, byteorder
+
         return kind, dtype.itemsize * 8, dtype_to_arrow_c_fmt(dtype), byteorder
 
     @property
@@ -194,6 +200,12 @@ def describe_null(self):
             column_null_dtype = ColumnNullType.USE_BYTEMASK
             null_value = 1
             return column_null_dtype, null_value
+        if isinstance(self._col.dtype, ArrowDtype):
+            if all(
+                chunk.buffers()[0] is None for chunk in self._col.array._pa_array.chunks
+            ):
+                return ColumnNullType.NON_NULLABLE, None
+            return ColumnNullType.USE_BITMASK, 0
         kind = self.dtype[0]
         try:
             null, value = _NULL_DESCRIPTION[kind]
@@ -302,15 +314,26 @@ def _get_data_buffer(
             DtypeKind.FLOAT,
             DtypeKind.BOOL,
         ):
+            dtype = self.dtype
             arr = self._col.array
+            if isinstance(self._col.dtype, ArrowDtype):
+                arr = self._col.array
+                buffer = PandasBufferPyarrow(
+                    arr._pa_array, is_validity=False, allow_copy=self._allow_copy
+                )
+                if self.dtype[0] == DtypeKind.BOOL:
+                    dtype = (
+                        DtypeKind.BOOL,
+                        1,
+                        ArrowCTypes.BOOL,
+                        Endianness.NATIVE,
+                    )
+                return buffer, dtype
             if isinstance(self._col.dtype, BaseMaskedDtype):
                 np_arr = arr._data  # type: ignore[attr-defined]
-            elif isinstance(self._col.dtype, ArrowDtype):
-                raise NotImplementedError("ArrowDtype not handled yet")
             else:
                 np_arr = arr._ndarray  # type: ignore[attr-defined]
             buffer = PandasBuffer(np_arr, allow_copy=self._allow_copy)
-            dtype = self.dtype
         elif self.dtype[0] == DtypeKind.CATEGORICAL:
             codes = self._col.values._codes
             buffer = PandasBuffer(codes, allow_copy=self._allow_copy)
@@ -351,6 +374,16 @@ def _get_validity_buffer(self) -> tuple[PandasBuffer, Any]:
         """
         null, invalid = self.describe_null
 
+        if isinstance(self._col.dtype, ArrowDtype):
+            arr = self._col.array
+            dtype = (DtypeKind.BOOL, 1, ArrowCTypes.BOOL, Endianness.NATIVE)
+            if all(chunk.buffers()[0] is None for chunk in arr._pa_array.chunks):
+                return None
+            buffer = PandasBufferPyarrow(
+                arr._pa_array, is_validity=True, allow_copy=self._allow_copy
+            )
+            return buffer, dtype
+
         if isinstance(self._col.dtype, BaseMaskedDtype):
             mask = self._col.array._mask  # type: ignore[attr-defined]
             buffer = PandasBuffer(mask)

pandas/core/interchange/from_dataframe.py

@@ -298,13 +298,14 @@ def string_column_to_ndarray(col: Column) -> tuple[np.ndarray, Any]:
 
     null_pos = None
     if null_kind in (ColumnNullType.USE_BITMASK, ColumnNullType.USE_BYTEMASK):
-        assert buffers["validity"], "Validity buffers cannot be empty for masks"
-        valid_buff, valid_dtype = buffers["validity"]
-        null_pos = buffer_to_ndarray(
-            valid_buff, valid_dtype, offset=col.offset, length=col.size()
-        )
-        if sentinel_val == 0:
-            null_pos = ~null_pos
+        validity = buffers["validity"]
+        if validity is not None:
+            valid_buff, valid_dtype = validity
+            null_pos = buffer_to_ndarray(
+                valid_buff, valid_dtype, offset=col.offset, length=col.size()
+            )
+            if sentinel_val == 0:
+                null_pos = ~null_pos
 
     # Assemble the strings from the code units
     str_list: list[None | float | str] = [None] * col.size()
@@ -516,19 +517,21 @@ def set_nulls(
     np.ndarray or pd.Series
         Data with the nulls being set.
     """
+    if validity is None:
+        return data
     null_kind, sentinel_val = col.describe_null
     null_pos = None
 
     if null_kind == ColumnNullType.USE_SENTINEL:
         null_pos = pd.Series(data) == sentinel_val
     elif null_kind in (ColumnNullType.USE_BITMASK, ColumnNullType.USE_BYTEMASK):
-        assert validity, "Expected to have a validity buffer for the mask"
         valid_buff, valid_dtype = validity
-        null_pos = buffer_to_ndarray(
-            valid_buff, valid_dtype, offset=col.offset, length=col.size()
-        )
-        if sentinel_val == 0:
-            null_pos = ~null_pos
+        if valid_buff is not None:
+            null_pos = buffer_to_ndarray(
+                valid_buff, valid_dtype, offset=col.offset, length=col.size()
+            )
+            if sentinel_val == 0:
+                null_pos = ~null_pos
     elif null_kind in (ColumnNullType.NON_NULLABLE, ColumnNullType.USE_NAN):
         pass
     else: