[REVIEW] Refactor the Buffer class

python/cudf/cudf/_lib/column.pyi

@@ -5,16 +5,16 @@ from __future__ import annotations
 from typing import Dict, Optional, Tuple, TypeVar
 
 from cudf._typing import Dtype, DtypeObj, ScalarLike
-from cudf.core.buffer import Buffer
+from cudf.core.buffer import DeviceBufferLike
 from cudf.core.column import ColumnBase
 
 T = TypeVar("T")
 
 class Column:
-    _data: Optional[Buffer]
-    _mask: Optional[Buffer]
-    _base_data: Optional[Buffer]
-    _base_mask: Optional[Buffer]
+    _data: Optional[DeviceBufferLike]
+    _mask: Optional[DeviceBufferLike]
+    _base_data: Optional[DeviceBufferLike]
+    _base_mask: Optional[DeviceBufferLike]
     _dtype: DtypeObj
     _size: int
     _offset: int
@@ -25,10 +25,10 @@ class Column:
 
     def __init__(
         self,
-        data: Optional[Buffer],
+        data: Optional[DeviceBufferLike],
         size: int,
         dtype: Dtype,
-        mask: Optional[Buffer] = None,
+        mask: Optional[DeviceBufferLike] = None,
         offset: int = None,
         null_count: int = None,
         children: Tuple[ColumnBase, ...] = (),
@@ -40,27 +40,27 @@ class Column:
     @property
     def size(self) -> int: ...
     @property
-    def base_data(self) -> Optional[Buffer]: ...
+    def base_data(self) -> Optional[DeviceBufferLike]: ...
     @property
     def base_data_ptr(self) -> int: ...
     @property
-    def data(self) -> Optional[Buffer]: ...
+    def data(self) -> Optional[DeviceBufferLike]: ...
     @property
     def data_ptr(self) -> int: ...
-    def set_base_data(self, value: Buffer) -> None: ...
+    def set_base_data(self, value: DeviceBufferLike) -> None: ...
     @property
     def nullable(self) -> bool: ...
     def has_nulls(self, include_nan: bool = False) -> bool: ...
     @property
-    def base_mask(self) -> Optional[Buffer]: ...
+    def base_mask(self) -> Optional[DeviceBufferLike]: ...
     @property
     def base_mask_ptr(self) -> int: ...
     @property
-    def mask(self) -> Optional[Buffer]: ...
+    def mask(self) -> Optional[DeviceBufferLike]: ...
     @property
     def mask_ptr(self) -> int: ...
-    def set_base_mask(self, value: Optional[Buffer]) -> None: ...
-    def set_mask(self: T, value: Optional[Buffer]) -> T: ...
+    def set_base_mask(self, value: Optional[DeviceBufferLike]) -> None: ...
+    def set_mask(self: T, value: Optional[DeviceBufferLike]) -> T: ...
     @property
     def null_count(self) -> int: ...
     @property

python/cudf/cudf/_lib/column.pyx

@@ -9,7 +9,7 @@ import rmm
 import cudf
 import cudf._lib as libcudf
 from cudf.api.types import is_categorical_dtype, is_list_dtype, is_struct_dtype
-from cudf.core.buffer import Buffer
+from cudf.core.buffer import Buffer, DeviceBufferLike, as_device_buffer_like
 
 from cpython.buffer cimport PyObject_CheckBuffer
 from libc.stdint cimport uintptr_t
@@ -56,9 +56,9 @@ cdef class Column:
     A Column stores columnar data in device memory.
     A Column may be composed of:
 
-    * A *data* Buffer
+    * A *data* DeviceBufferLike
     * One or more (optional) *children* Columns
-    * An (optional) *mask* Buffer representing the nullmask
+    * An (optional) *mask* DeviceBufferLike representing the nullmask
 
     The *dtype* indicates the Column's element type.
     """
@@ -110,18 +110,9 @@ cdef class Column:
         if self.base_data is None:
             return None
         if self._data is None:
-            itemsize = self.dtype.itemsize
-            size = self.size * itemsize
-            offset = self.offset * itemsize if self.size else 0
-            if offset == 0 and self.base_data.size == size:
-                # `data` spans all of `base_data`
-                self._data = self.base_data
-            else:
-                self._data = Buffer.from_buffer(
-                    buffer=self.base_data,
-                    size=size,
-                    offset=offset
-                )
+            start = self.offset * self.dtype.itemsize
+            end = start + self.size * self.dtype.itemsize
+            self._data = self.base_data[start:end]
         return self._data
 
     @property
@@ -132,9 +123,11 @@ cdef class Column:
             return self.data.ptr
 
     def set_base_data(self, value):
-        if value is not None and not isinstance(value, Buffer):
-            raise TypeError("Expected a Buffer or None for data, got " +
-                            type(value).__name__)
+        if value is not None and not isinstance(value, DeviceBufferLike):
+            raise TypeError(
+                "Expected a DeviceBufferLike or None for data, "
+                f"got {type(value).__name__}"
+            )
 
         self._data = None
         self._base_data = value
@@ -179,17 +172,18 @@ cdef class Column:
         modify size or offset in any way, so the passed mask is expected to be
         compatible with the current offset.
         """
-        if value is not None and not isinstance(value, Buffer):
-            raise TypeError("Expected a Buffer or None for mask, got " +
-                            type(value).__name__)
+        if value is not None and not isinstance(value, DeviceBufferLike):
+            raise TypeError(
+                "Expected a DeviceBufferLike or None for mask, "
+                f"got {type(value).__name__}"
+            )
 
         if value is not None:
             required_size = bitmask_allocation_size_bytes(self.base_size)
             if value.size < required_size:
                 error_msg = (
-                    "The Buffer for mask is smaller than expected, got " +
-                    str(value.size) + " bytes, expected " +
-                    str(required_size) + " bytes."
+                    "The DeviceBufferLike for mask is smaller than expected, "
+                    f"got {value.size} bytes, expected {required_size} bytes."
                 )
                 if self.offset > 0 or self.size < self.base_size:
                     error_msg += (
@@ -233,31 +227,31 @@ cdef class Column:
                 if isinstance(value, Column):
                     value = value.data_array_view
                 value = cp.asarray(value).view('|u1')
-            mask = Buffer(value)
+            mask = as_device_buffer_like(value)
             if mask.size < required_num_bytes:
                 raise ValueError(error_msg.format(str(value.size)))
             if mask.size < mask_size:
                 dbuf = rmm.DeviceBuffer(size=mask_size)
                 dbuf.copy_from_device(value)
-                mask = Buffer(dbuf)
+                mask = as_device_buffer_like(dbuf)
         elif hasattr(value, "__array_interface__"):
             value = np.asarray(value).view("u1")[:mask_size]
             if value.size < required_num_bytes:
                 raise ValueError(error_msg.format(str(value.size)))
             dbuf = rmm.DeviceBuffer(size=mask_size)
             dbuf.copy_from_host(value)
-            mask = Buffer(dbuf)
+            mask = as_device_buffer_like(dbuf)
         elif PyObject_CheckBuffer(value):
             value = np.asarray(value).view("u1")[:mask_size]
             if value.size < required_num_bytes:
                 raise ValueError(error_msg.format(str(value.size)))
             dbuf = rmm.DeviceBuffer(size=mask_size)
             dbuf.copy_from_host(value)
-            mask = Buffer(dbuf)
+            mask = as_device_buffer_like(dbuf)
         else:
             raise TypeError(
-                "Expected a Buffer-like object or None for mask, got "
-                + type(value).__name__
+                "Expected a DeviceBufferLike object or None for mask, "
+                f"got {type(value).__name__}"
             )
 
         return cudf.core.column.build_column(
@@ -455,11 +449,11 @@ cdef class Column:
         cdef column_contents contents = move(c_col.get()[0].release())
 
         data = DeviceBuffer.c_from_unique_ptr(move(contents.data))
-        data = Buffer(data)
+        data = as_device_buffer_like(data)
 
         if null_count > 0:
             mask = DeviceBuffer.c_from_unique_ptr(move(contents.null_mask))
-            mask = Buffer(mask)
+            mask = as_device_buffer_like(mask)
         else:
             mask = None
 
@@ -484,9 +478,10 @@ cdef class Column:
         Given a ``cudf::column_view``, constructs a ``cudf.Column`` from it,
         along with referencing an ``owner`` Python object that owns the memory
         lifetime. If ``owner`` is a ``cudf.Column``, we reach inside of it and
-        make the owner of each newly created ``Buffer`` the respective
-        ``Buffer`` from the ``owner`` ``cudf.Column``. If ``owner`` is
-        ``None``, we allocate new memory for the resulting ``cudf.Column``.
+        make the owner of each newly created ``DeviceBufferLike`` the
+        respective ``DeviceBufferLike`` from the ``owner`` ``cudf.Column``.
+        If ``owner`` is ``None``, we allocate new memory for the resulting
+        ``cudf.Column``.
         """
         column_owner = isinstance(owner, Column)
         mask_owner = owner
@@ -509,7 +504,7 @@ cdef class Column:
 
         if data_ptr:
             if data_owner is None:
-                data = Buffer(
+                data = as_device_buffer_like(
                     rmm.DeviceBuffer(ptr=data_ptr,
                                      size=(size+offset) * dtype.itemsize)
                 )
@@ -520,7 +515,7 @@ cdef class Column:
                     owner=data_owner
                 )
         else:
-            data = Buffer(
+            data = as_device_buffer_like(
                 rmm.DeviceBuffer(ptr=data_ptr, size=0)
             )
 
@@ -550,7 +545,7 @@ cdef class Column:
                     # result:
                     mask = None
                 else:
-                    mask = Buffer(
+                    mask = as_device_buffer_like(
                         rmm.DeviceBuffer(
                             ptr=mask_ptr,
                             size=bitmask_allocation_size_bytes(size+offset)

python/cudf/cudf/_lib/concat.pyx

@@ -1,4 +1,4 @@
-# Copyright (c) 2020, NVIDIA CORPORATION.
+# Copyright (c) 2020-2022, NVIDIA CORPORATION.
 
 from libcpp cimport bool
 from libcpp.memory cimport make_unique, unique_ptr
@@ -19,7 +19,7 @@ from cudf._lib.utils cimport (
     table_view_from_table,
 )
 
-from cudf.core.buffer import Buffer
+from cudf.core.buffer import as_device_buffer_like
 
 from rmm._lib.device_buffer cimport DeviceBuffer, device_buffer
 
@@ -31,7 +31,9 @@ cpdef concat_masks(object columns):
     with nogil:
         c_result = move(libcudf_concatenate_masks(c_views))
         c_unique_result = make_unique[device_buffer](move(c_result))
-    return Buffer(DeviceBuffer.c_from_unique_ptr(move(c_unique_result)))
+    return as_device_buffer_like(
+        DeviceBuffer.c_from_unique_ptr(move(c_unique_result))
+    )
 
 
 cpdef concat_columns(object columns):

python/cudf/cudf/_lib/copying.pyx

@@ -718,7 +718,11 @@ cdef class _CPackedColumns:
         header = {}
         frames = []
 
-        gpu_data = Buffer(self.gpu_data_ptr, self.gpu_data_size, self)
+        gpu_data = Buffer(
+            data=self.gpu_data_ptr,
+            size=self.gpu_data_size,
+            owner=self
+        )
         data_header, data_frames = gpu_data.serialize()
         header["data"] = data_header
         frames.extend(data_frames)

python/cudf/cudf/_lib/null_mask.pyx

@@ -17,7 +17,7 @@ from cudf._lib.cpp.null_mask cimport (
 )
 from cudf._lib.cpp.types cimport mask_state, size_type
 
-from cudf.core.buffer import Buffer
+from cudf.core.buffer import as_device_buffer_like
 
 
 class MaskState(Enum):
@@ -47,7 +47,7 @@ def copy_bitmask(Column col):
         up_db = make_unique[device_buffer](move(db))
 
     rmm_db = DeviceBuffer.c_from_unique_ptr(move(up_db))
-    buf = Buffer(rmm_db)
+    buf = as_device_buffer_like(rmm_db)
     return buf
 
 
@@ -93,5 +93,5 @@ def create_null_mask(size_type size, state=MaskState.UNINITIALIZED):
         up_db = make_unique[device_buffer](move(db))
 
     rmm_db = DeviceBuffer.c_from_unique_ptr(move(up_db))
-    buf = Buffer(rmm_db)
+    buf = as_device_buffer_like(rmm_db)
     return buf

python/cudf/cudf/_lib/transform.pyx

@@ -6,7 +6,7 @@ from numba.np import numpy_support
 import cudf
 from cudf._lib.types import SUPPORTED_NUMPY_TO_LIBCUDF_TYPES
 from cudf.core._internals.expressions import parse_expression
-from cudf.core.buffer import Buffer
+from cudf.core.buffer import as_device_buffer_like
 from cudf.utils import cudautils
 
 from cython.operator cimport dereference
@@ -40,7 +40,7 @@ from cudf._lib.utils cimport (
 def bools_to_mask(Column col):
     """
     Given an int8 (boolean) column, compress the data from booleans to bits and
-    return a Buffer
+    return a DeviceBufferLike
     """
     cdef column_view col_view = col.view()
     cdef pair[unique_ptr[device_buffer], size_type] cpp_out
@@ -52,7 +52,7 @@ def bools_to_mask(Column col):
         up_db = move(cpp_out.first)
 
     rmm_db = DeviceBuffer.c_from_unique_ptr(move(up_db))
-    buf = Buffer(rmm_db)
+    buf = as_device_buffer_like(rmm_db)
     return buf
 
 
@@ -61,9 +61,9 @@ def mask_to_bools(object mask_buffer, size_type begin_bit, size_type end_bit):
     Given a mask buffer, returns a boolean column representng bit 0 -> False
     and 1 -> True within range of [begin_bit, end_bit),
     """
-    if not isinstance(mask_buffer, cudf.core.buffer.Buffer):
+    if not isinstance(mask_buffer, cudf.core.buffer.DeviceBufferLike):
         raise TypeError("mask_buffer is not an instance of "
-                        "cudf.core.buffer.Buffer")
+                        "cudf.core.buffer.DeviceBufferLike")
     cdef bitmask_type* bit_mask = <bitmask_type*><uintptr_t>(mask_buffer.ptr)
 
     cdef unique_ptr[column] result
@@ -88,7 +88,7 @@ def nans_to_nulls(Column input):
         return None
 
     buffer = DeviceBuffer.c_from_unique_ptr(move(c_buffer))
-    buffer = Buffer(buffer)
+    buffer = as_device_buffer_like(buffer)
     return buffer
 
 

python/cudf/cudf/_lib/utils.pyx

@@ -341,8 +341,8 @@ cdef data_from_table_view(
     along with referencing an ``owner`` Python object that owns the memory
     lifetime. If ``owner`` is a Frame we reach inside of it and
     reach inside of each ``cudf.Column`` to make the owner of each newly
-    created ``Buffer`` underneath the ``cudf.Column`` objects of the
-    created Frame the respective ``Buffer`` from the relevant
+    created ``DeviceBufferLike`` underneath the ``cudf.Column`` objects of the
+    created Frame the respective ``DeviceBufferLike`` from the relevant
     ``cudf.Column`` of the ``owner`` Frame
     """
     cdef size_type column_idx = 0