Separate buffer and column implementations

Signed-off-by: Vasily Litvinov <vasilij.n.litvinov@intel.com>

Author: vnlitvinov

pandas/core/exchange/buffer.py

@@ -0,0 +1,98 @@
+from pandas.core.exchange.dataframe_protocol import Buffer, DlpackDeviceType, DtypeKind
+import numpy as np
+from typing import Tuple
+import ctypes
+
+
+_NP_DTYPES = {
+    DtypeKind.INT: {8: np.int8, 16: np.int16, 32: np.int32, 64: np.int64},
+    DtypeKind.UINT: {8: np.uint8, 16: np.uint16, 32: np.uint32, 64: np.uint64},
+    DtypeKind.FLOAT: {32: np.float32, 64: np.float64},
+    DtypeKind.BOOL: {8: bool},
+}
+
+
+class PandasBuffer(Buffer):
+    """
+    Data in the buffer is guaranteed to be contiguous in memory.
+    """
+
+    def __init__(self, x: np.ndarray, allow_copy: bool = True) -> None:
+        """
+        Handle only regular columns (= numpy arrays) for now.
+        """
+        if not x.strides == (x.dtype.itemsize,):
+            # The protocol does not support strided buffers, so a copy is
+            # necessary. If that's not allowed, we need to raise an exception.
+            if allow_copy:
+                x = x.copy()
+            else:
+                raise RuntimeError(
+                    "Exports cannot be zero-copy in the case "
+                    "of a non-contiguous buffer"
+                )
+
+        # Store the numpy array in which the data resides as a private
+        # attribute, so we can use it to retrieve the public attributes
+        self._x = x
+
+    @property
+    def bufsize(self) -> int:
+        """
+        Buffer size in bytes.
+        """
+        return self._x.size * self._x.dtype.itemsize
+
+    @property
+    def ptr(self) -> int:
+        """
+        Pointer to start of the buffer as an integer.
+        """
+        return self._x.__array_interface__["data"][0]
+
+    def __dlpack__(self):
+        """
+        DLPack not implemented in NumPy yet, so leave it out here.
+        """
+        raise NotImplementedError("__dlpack__")
+
+    def __dlpack_device__(self) -> Tuple[DlpackDeviceType, int]:
+        """
+        Device type and device ID for where the data in the buffer resides.
+        """
+        return (DlpackDeviceType.CPU, None)
+
+    def __repr__(self) -> str:
+        return (
+            "PandasBuffer("
+            + str(
+                {
+                    "bufsize": self.bufsize,
+                    "ptr": self.ptr,
+                    "device": self.__dlpack_device__()[0].name,
+                }
+            )
+            + ")"
+        )
+
+
+def buffer_to_ndarray(_buffer: Buffer, _dtype) -> np.ndarray:
+    # Handle the dtype
+    kind = _dtype[0]
+    bitwidth = _dtype[1]
+    if kind not in _NP_DTYPES:
+        raise RuntimeError(f"Unsupported data type: {kind}")
+
+    column_dtype = _NP_DTYPES[kind][bitwidth]
+
+    # No DLPack yet, so need to construct a new ndarray from the data pointer
+    # and size in the buffer plus the dtype on the column
+    ctypes_type = np.ctypeslib.as_ctypes_type(column_dtype)
+    data_pointer = ctypes.cast(_buffer.ptr, ctypes.POINTER(ctypes_type))
+
+    # NOTE: `x` does not own its memory, so the caller of this function must
+    #       either make a copy or hold on to a reference of the column or
+    #       buffer! (not done yet, this is pretty awful ...)
+    x = np.ctypeslib.as_array(data_pointer, shape=(_buffer.bufsize // (bitwidth // 8),))
+
+    return x