combine chunks

Author: MarcoGorelli

pandas/core/interchange/buffer.py

@@ -12,6 +12,7 @@
 
 if TYPE_CHECKING:
     import numpy as np
+    import pyarrow as pa
 
 
 class PandasBuffer(Buffer):
@@ -83,14 +84,21 @@ class PandasPyarrowBackedBuffer(Buffer):
     Data in the buffer is guaranteed to be contiguous in memory.
     """
 
-    def __init__(self, arr: Any, allow_copy: bool = True) -> None:
+    def __init__(
+        self, chunked_array: pa.ChunkedArray, *, allow_copy: bool = True
+    ) -> None:
         """
-        Handle only regular columns (= numpy arrays) for now.
+        Handle pyarrow chunked arrays.
         """
-
-        # Store the numpy array in which the data resides as a private
-        # attribute, so we can use it to retrieve the public attributes
-        self._buffer = arr.chunks[0].buffers()[1]
+        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()
+        self._buffer = arr.buffers()[1]
         self._length = len(arr)
 
     @property