From 95935db773c637d9a4d271bccb0c6bf8d60b109f Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Thu, 11 Aug 2022 15:40:45 +0200 Subject: [PATCH] Refactor the `Buffer` class (#11447) This PR introduces factory functions to create `Buffer` instances, which makes it possible to change the returned buffer type based on a configuration option in a follow-up PR. Beside simplifying the code base a bit, this is motivated by the spilling work in https://github.com/rapidsai/cudf/pull/10746. We would like to introduce a new spillable Buffer class that requires minimal changes to the existing code and is only used when enabled explicitly. This way, we can introduce spilling in cuDF as an experimental feature with minimal risk to the existing code. @shwina and I discussed the possibility to let `Buffer.__new__` return different class type instances instead of using factory functions but we concluded that having `Buffer()` return anything other than an instance of `Buffer` is simply too surprising :) **Notice**, this is breaking because it removes unused methods such as `Buffer.copy()` and `Buffer.nbytes`. ~~However, we still support creating a buffer directly by calling `Buffer(obj)`. AFAIK, this is the only way `Buffer` is created outside of cuDF, which [a github search seems to confirm](https://github.com/search?l=&q=cudf.core.buffer+-repo%3Arapidsai%2Fcudf&type=code).~~ This PR doesn't change the signature of `Buffer.__init__()` anymore. Authors: - Mads R. B. Kristensen (https://github.com/madsbk) Approvers: - Ashwin Srinath (https://github.com/shwina) - Lawrence Mitchell (https://github.com/wence-) - Bradley Dice (https://github.com/bdice) - https://github.com/brandon-b-miller URL: https://github.com/rapidsai/cudf/pull/11447 --- python/cudf/cudf/_lib/column.pyi | 28 +- python/cudf/cudf/_lib/column.pyx | 71 ++-- python/cudf/cudf/_lib/concat.pyx | 8 +- python/cudf/cudf/_lib/copying.pyx | 6 +- python/cudf/cudf/_lib/null_mask.pyx | 6 +- python/cudf/cudf/_lib/transform.pyx | 12 +- python/cudf/cudf/_lib/utils.pyx | 4 +- python/cudf/cudf/core/abc.py | 12 +- python/cudf/cudf/core/buffer.py | 377 +++++++++++------- python/cudf/cudf/core/column/categorical.py | 18 +- python/cudf/cudf/core/column/column.py | 86 ++-- python/cudf/cudf/core/column/datetime.py | 14 +- python/cudf/cudf/core/column/decimal.py | 14 +- python/cudf/cudf/core/column/lists.py | 3 +- python/cudf/cudf/core/column/numerical.py | 18 +- python/cudf/cudf/core/column/string.py | 6 +- python/cudf/cudf/core/column/struct.py | 4 +- python/cudf/cudf/core/column/timedelta.py | 14 +- python/cudf/cudf/core/df_protocol.py | 34 +- python/cudf/cudf/core/dtypes.py | 4 +- python/cudf/cudf/core/index.py | 2 +- python/cudf/cudf/core/series.py | 4 +- python/cudf/cudf/tests/test_buffer.py | 96 +++-- python/cudf/cudf/tests/test_column.py | 2 +- .../cudf/tests/test_cuda_array_interface.py | 6 +- python/cudf/cudf/tests/test_df_protocol.py | 3 +- python/cudf/cudf/tests/test_pickling.py | 4 +- python/cudf/cudf/tests/test_serialize.py | 4 +- python/cudf/cudf/tests/test_testing.py | 2 +- python/cudf/cudf/utils/string.py | 13 + python/cudf/cudf/utils/utils.py | 6 +- 31 files changed, 517 insertions(+), 364 deletions(-) create mode 100644 python/cudf/cudf/utils/string.py diff --git a/python/cudf/cudf/_lib/column.pyi b/python/cudf/cudf/_lib/column.pyi index c38c560b982..fd9aab038d4 100644 --- a/python/cudf/cudf/_lib/column.pyi +++ b/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 diff --git a/python/cudf/cudf/_lib/column.pyx b/python/cudf/cudf/_lib/column.pyx index 8a9a79250b9..78125c027dd 100644 --- a/python/cudf/cudf/_lib/column.pyx +++ b/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) diff --git a/python/cudf/cudf/_lib/concat.pyx b/python/cudf/cudf/_lib/concat.pyx index a7f8296bad5..ed858034032 100644 --- a/python/cudf/cudf/_lib/concat.pyx +++ b/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): diff --git a/python/cudf/cudf/_lib/copying.pyx b/python/cudf/cudf/_lib/copying.pyx index fcf70f2f69f..f1183e008f8 100644 --- a/python/cudf/cudf/_lib/copying.pyx +++ b/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) diff --git a/python/cudf/cudf/_lib/null_mask.pyx b/python/cudf/cudf/_lib/null_mask.pyx index ce83a6f0f18..b0ee28baf29 100644 --- a/python/cudf/cudf/_lib/null_mask.pyx +++ b/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 diff --git a/python/cudf/cudf/_lib/transform.pyx b/python/cudf/cudf/_lib/transform.pyx index 2d94ef2cedf..5fa45f68357 100644 --- a/python/cudf/cudf/_lib/transform.pyx +++ b/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 = (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 diff --git a/python/cudf/cudf/_lib/utils.pyx b/python/cudf/cudf/_lib/utils.pyx index 643a1adca9f..e0bdc7d8f74 100644 --- a/python/cudf/cudf/_lib/utils.pyx +++ b/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 diff --git a/python/cudf/cudf/core/abc.py b/python/cudf/cudf/core/abc.py index d3da544f8b5..dcbf96313a7 100644 --- a/python/cudf/cudf/core/abc.py +++ b/python/cudf/cudf/core/abc.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2021, NVIDIA CORPORATION. +# Copyright (c) 2020-2022, NVIDIA CORPORATION. """Common abstract base classes for cudf.""" import sys @@ -90,13 +90,15 @@ def device_serialize(self): header : dict The metadata required to reconstruct the object. frames : list - The Buffers or memoryviews that the object should contain. + The DeviceBufferLike or memoryview objects that the object + should contain. :meta private: """ header, frames = self.serialize() assert all( - (type(f) in [cudf.core.buffer.Buffer, memoryview]) for f in frames + isinstance(f, (cudf.core.buffer.DeviceBufferLike, memoryview)) + for f in frames ) header["type-serialized"] = pickle.dumps(type(self)) header["is-cuda"] = [ @@ -130,7 +132,7 @@ def device_deserialize(cls, header, frames): """ typ = pickle.loads(header["type-serialized"]) frames = [ - cudf.core.buffer.Buffer(f) if c else memoryview(f) + cudf.core.buffer.as_device_buffer_like(f) if c else memoryview(f) for c, f in zip(header["is-cuda"], frames) ] assert all( @@ -158,7 +160,7 @@ def host_serialize(self): header, frames = self.device_serialize() header["writeable"] = len(frames) * (None,) frames = [ - f.to_host_array().data if c else memoryview(f) + f.memoryview() if c else memoryview(f) for c, f in zip(header["is-cuda"], frames) ] return header, frames diff --git a/python/cudf/cudf/core/buffer.py b/python/cudf/cudf/core/buffer.py index 753be3b27e1..647e747e127 100644 --- a/python/cudf/cudf/core/buffer.py +++ b/python/cudf/cudf/core/buffer.py @@ -2,10 +2,19 @@ from __future__ import annotations -import functools -import operator +import math import pickle -from typing import Any, Dict, Tuple +from typing import ( + Any, + Dict, + List, + Mapping, + Protocol, + Sequence, + Tuple, + Union, + runtime_checkable, +) import numpy as np @@ -13,24 +22,129 @@ import cudf from cudf.core.abc import Serializable +from cudf.utils.string import format_bytes + +# Frame type for serialization and deserialization of `DeviceBufferLike` +Frame = Union[memoryview, "DeviceBufferLike"] + + +@runtime_checkable +class DeviceBufferLike(Protocol): + def __getitem__(self, key: slice) -> DeviceBufferLike: + """Create a new view of the buffer.""" + + @property + def size(self) -> int: + """Size of the buffer in bytes.""" + + @property + def nbytes(self) -> int: + """Size of the buffer in bytes.""" + + @property + def ptr(self) -> int: + """Device pointer to the start of the buffer.""" + + @property + def owner(self) -> Any: + """Object owning the memory of the buffer.""" + + @property + def __cuda_array_interface__(self) -> Mapping: + """Implementation of the CUDA Array Interface.""" + + def memoryview(self) -> memoryview: + """Read-only access to the buffer through host memory.""" + + def serialize(self) -> Tuple[dict, List[Frame]]: + """Serialize the buffer into header and frames. + + The frames can be a mixture of memoryview and device-buffer-like + objects. + + Returns + ------- + Tuple[Dict, List] + The first element of the returned tuple is a dict containing any + serializable metadata required to reconstruct the object. The + second element is a list containing the device buffers and + memoryviews of the object. + """ + + @classmethod + def deserialize( + cls, header: dict, frames: List[Frame] + ) -> DeviceBufferLike: + """Generate an buffer from a serialized representation. + + Parameters + ---------- + header : dict + The metadata required to reconstruct the object. + frames : list + The device-buffer-like and memoryview buffers that the object + should contain. + + Returns + ------- + DeviceBufferLike + A new object that implements DeviceBufferLike. + """ + + +def as_device_buffer_like(obj: Any) -> DeviceBufferLike: + """ + Factory function to wrap `obj` in a DeviceBufferLike object. + + If `obj` isn't device-buffer-like already, a new buffer that implements + DeviceBufferLike and points to the memory of `obj` is created. If `obj` + represents host memory, it is copied to a new `rmm.DeviceBuffer` device + allocation. Otherwise, the data of `obj` is **not** copied, instead the + new buffer keeps a reference to `obj` in order to retain the lifetime + of `obj`. + + Raises ValueError if the data of `obj` isn't C-contiguous. + + Parameters + ---------- + obj : buffer-like or array-like + An object that exposes either device or host memory through + `__array_interface__`, `__cuda_array_interface__`, or the + buffer protocol. If `obj` represents host memory, data will + be copied. + + Return + ------ + DeviceBufferLike + A device-buffer-like instance that represents the device memory + of `obj`. + """ + + if isinstance(obj, DeviceBufferLike): + return obj + return Buffer(obj) class Buffer(Serializable): """ - A Buffer represents a device memory allocation. + A Buffer represents device memory. + + Usually Buffers will be created using `as_device_buffer_like(obj)`, + which will make sure that `obj` is device-buffer-like and not a `Buffer` + necessarily. Parameters ---------- - data : Buffer, array_like, int - An array-like object or integer representing a - device or host pointer to pre-allocated memory. + data : int or buffer-like or array-like + An integer representing a pointer to device memory or a buffer-like + or array-like object. When not an integer, `size` and `owner` must + be None. size : int, optional - Size of memory allocation. Required if a pointer - is passed for `data`. + Size of device memory in bytes. Must be specified if `data` is an + integer. owner : object, optional - Python object to which the lifetime of the memory - allocation is tied. If provided, a reference to this - object is kept in this Buffer. + Python object to which the lifetime of the memory allocation is tied. + A reference to this object is kept in the returned Buffer. """ _ptr: int @@ -38,62 +152,61 @@ class Buffer(Serializable): _owner: object def __init__( - self, data: Any = None, size: int = None, owner: object = None + self, data: Union[int, Any], *, size: int = None, owner: object = None ): - if isinstance(data, Buffer): - self._ptr = data._ptr - self._size = data.size - self._owner = owner or data._owner - elif isinstance(data, rmm.DeviceBuffer): - self._ptr = data.ptr - self._size = data.size - self._owner = data - elif hasattr(data, "__array_interface__") or hasattr( - data, "__cuda_array_interface__" - ): - self._init_from_array_like(data, owner) - elif isinstance(data, memoryview): - self._init_from_array_like(np.asarray(data), owner) - elif isinstance(data, int): - if not isinstance(size, int): - raise TypeError("size must be integer") + if isinstance(data, int): + if size is None: + raise ValueError( + "size must be specified when `data` is an integer" + ) + if size < 0: + raise ValueError("size cannot be negative") self._ptr = data self._size = size self._owner = owner - elif data is None: - self._ptr = 0 - self._size = 0 - self._owner = None else: - try: - data = memoryview(data) - except TypeError: - raise TypeError("data must be Buffer, array-like or integer") - self._init_from_array_like(np.asarray(data), owner) + if size is not None or owner is not None: + raise ValueError( + "`size` and `owner` must be None when " + "`data` is a buffer-like object" + ) + + # `data` is a buffer-like object + buf: Any = data + if isinstance(buf, rmm.DeviceBuffer): + self._ptr = buf.ptr + self._size = buf.size + self._owner = buf + return + iface = getattr(buf, "__cuda_array_interface__", None) + if iface: + ptr, size = get_ptr_and_size(iface) + self._ptr = ptr + self._size = size + self._owner = buf + return + ptr, size = get_ptr_and_size(np.asarray(buf).__array_interface__) + buf = rmm.DeviceBuffer(ptr=ptr, size=size) + self._ptr = buf.ptr + self._size = buf.size + self._owner = buf + + def __getitem__(self, key: slice) -> Buffer: + if not isinstance(key, slice): + raise ValueError("index must be an slice") + start, stop, step = key.indices(self.size) + if step != 1: + raise ValueError("slice must be contiguous") + return self.__class__( + data=self.ptr + start, size=stop - start, owner=self.owner + ) - @classmethod - def from_buffer(cls, buffer: Buffer, size: int = None, offset: int = 0): - """ - Create a buffer from another buffer - - Parameters - ---------- - buffer : Buffer - The base buffer, which will also be set as the owner of - the memory allocation. - size : int, optional - Size of the memory allocation (default: `buffer.size`). - offset : int, optional - Start offset relative to `buffer.ptr`. - """ - - ret = cls() - ret._ptr = buffer._ptr + offset - ret._size = buffer.size if size is None else size - ret._owner = buffer - return ret + @property + def size(self) -> int: + return self._size - def __len__(self) -> int: + @property + def nbytes(self) -> int: return self._size @property @@ -101,12 +214,8 @@ def ptr(self) -> int: return self._ptr @property - def size(self) -> int: - return self._size - - @property - def nbytes(self) -> int: - return self._size + def owner(self) -> Any: + return self._owner @property def __cuda_array_interface__(self) -> dict: @@ -118,32 +227,10 @@ def __cuda_array_interface__(self) -> dict: "version": 0, } - def to_host_array(self): - data = np.empty((self.size,), "u1") - rmm._lib.device_buffer.copy_ptr_to_host(self.ptr, data) - return data - - def _init_from_array_like(self, data, owner): - - if hasattr(data, "__cuda_array_interface__"): - confirm_1d_contiguous(data.__cuda_array_interface__) - ptr, size = _buffer_data_from_array_interface( - data.__cuda_array_interface__ - ) - self._ptr = ptr - self._size = size - self._owner = owner or data - elif hasattr(data, "__array_interface__"): - confirm_1d_contiguous(data.__array_interface__) - ptr, size = _buffer_data_from_array_interface( - data.__array_interface__ - ) - dbuf = rmm.DeviceBuffer(ptr=ptr, size=size) - self._init_from_array_like(dbuf, owner) - else: - raise TypeError( - f"Cannot construct Buffer from {data.__class__.__name__}" - ) + def memoryview(self) -> memoryview: + host_buf = bytearray(self.size) + rmm._lib.device_buffer.copy_ptr_to_host(self.ptr, host_buf) + return memoryview(host_buf).toreadonly() def serialize(self) -> Tuple[dict, list]: header = {} # type: Dict[Any, Any] @@ -171,62 +258,68 @@ def deserialize(cls, header: dict, frames: list) -> Buffer: return buf - @classmethod - def empty(cls, size: int) -> Buffer: - return Buffer(rmm.DeviceBuffer(size=size)) + def __repr__(self) -> str: + return ( + f" Buffer: - """ - Create a new Buffer containing a copy of the data contained - in this Buffer. - """ - from rmm._lib.device_buffer import copy_device_to_ptr - out = Buffer.empty(size=self.size) - copy_device_to_ptr(self.ptr, out.ptr, self.size) - return out - - -def _buffer_data_from_array_interface(array_interface): - ptr = array_interface["data"][0] - if ptr is None: - ptr = 0 - itemsize = cudf.dtype(array_interface["typestr"]).itemsize - shape = ( - array_interface["shape"] if len(array_interface["shape"]) > 0 else (1,) - ) - size = functools.reduce(operator.mul, shape) - return ptr, size * itemsize +def is_c_contiguous( + shape: Sequence[int], strides: Sequence[int], itemsize: int +) -> bool: + """ + Determine if shape and strides are C-contiguous + Parameters + ---------- + shape : Sequence[int] + Number of elements in each dimension. + strides : Sequence[int] + The stride of each dimension in bytes. + itemsize : int + Size of an element in bytes. + + Return + ------ + bool + The boolean answer. + """ -def confirm_1d_contiguous(array_interface): - strides = array_interface["strides"] - shape = array_interface["shape"] - itemsize = cudf.dtype(array_interface["typestr"]).itemsize - typestr = array_interface["typestr"] - if typestr not in ("|i1", "|u1", "|b1"): - raise TypeError("Buffer data must be of uint8 type") - if not get_c_contiguity(shape, strides, itemsize): - raise ValueError("Buffer data must be 1D C-contiguous") + if any(dim == 0 for dim in shape): + return True + cumulative_stride = itemsize + for dim, stride in zip(reversed(shape), reversed(strides)): + if dim > 1 and stride != cumulative_stride: + return False + cumulative_stride *= dim + return True -def get_c_contiguity(shape, strides, itemsize): - """ - Determine if combination of array parameters represents a - c-contiguous array. +def get_ptr_and_size(array_interface: Mapping) -> Tuple[int, int]: """ - ndim = len(shape) - assert strides is None or ndim == len(strides) + Retrieve the pointer and size from an array interface. - if ndim == 0 or strides is None or (ndim == 1 and strides[0] == itemsize): - return True + Raises ValueError if array isn't C-contiguous. - # any dimension zero, trivial case - for dim in shape: - if dim == 0: - return True + Parameters + ---------- + array_interface : Mapping + The array interface metadata. + + Return + ------ + pointer : int + The pointer to device or host memory + size : int + The size in bytes + """ - for this_dim, this_stride in zip(shape, strides): - if this_stride != this_dim * itemsize: - return False - return True + shape = array_interface["shape"] or (1,) + strides = array_interface["strides"] + itemsize = cudf.dtype(array_interface["typestr"]).itemsize + if strides is None or is_c_contiguous(shape, strides, itemsize): + nelem = math.prod(shape) + ptr = array_interface["data"][0] or 0 + return ptr, nelem * itemsize + raise ValueError("Buffer data must be C-contiguous") diff --git a/python/cudf/cudf/core/column/categorical.py b/python/cudf/cudf/core/column/categorical.py index c04e2e45461..6762c0bc6c3 100644 --- a/python/cudf/cudf/core/column/categorical.py +++ b/python/cudf/cudf/core/column/categorical.py @@ -16,7 +16,7 @@ from cudf._lib.transform import bools_to_mask from cudf._typing import ColumnBinaryOperand, ColumnLike, Dtype, ScalarLike from cudf.api.types import is_categorical_dtype, is_interval_dtype -from cudf.core.buffer import Buffer +from cudf.core.buffer import DeviceBufferLike from cudf.core.column import column from cudf.core.column.methods import ColumnMethods from cudf.core.dtypes import CategoricalDtype @@ -610,7 +610,7 @@ class CategoricalColumn(column.ColumnBase): Parameters ---------- dtype : CategoricalDtype - mask : Buffer + mask : DeviceBufferLike The validity mask offset : int Data offset @@ -634,7 +634,7 @@ class CategoricalColumn(column.ColumnBase): def __init__( self, dtype: CategoricalDtype, - mask: Buffer = None, + mask: DeviceBufferLike = None, size: int = None, offset: int = 0, null_count: int = None, @@ -693,7 +693,7 @@ def _process_values_for_isin( rhs = cudf.core.column.as_column(values, dtype=self.dtype) return lhs, rhs - def set_base_mask(self, value: Optional[Buffer]): + def set_base_mask(self, value: Optional[DeviceBufferLike]): super().set_base_mask(value) self._codes = None @@ -705,16 +705,12 @@ def set_base_children(self, value: Tuple[ColumnBase, ...]): def children(self) -> Tuple[NumericalColumn]: if self._children is None: codes_column = self.base_children[0] - - buf = Buffer.from_buffer( - buffer=codes_column.base_data, - size=self.size * codes_column.dtype.itemsize, - offset=self.offset * codes_column.dtype.itemsize, - ) + start = self.offset * codes_column.dtype.itemsize + end = start + self.size * codes_column.dtype.itemsize codes_column = cast( cudf.core.column.NumericalColumn, column.build_column( - data=buf, + data=codes_column.base_data[start:end], dtype=codes_column.dtype, size=self.size, ), diff --git a/python/cudf/cudf/core/column/column.py b/python/cudf/cudf/core/column/column.py index 67d744e6690..2e75c6c2225 100644 --- a/python/cudf/cudf/core/column/column.py +++ b/python/cudf/cudf/core/column/column.py @@ -26,6 +26,8 @@ import pyarrow as pa from numba import cuda +import rmm + import cudf from cudf import _lib as libcudf from cudf._lib.column import Column @@ -61,7 +63,7 @@ is_struct_dtype, ) from cudf.core.abc import Serializable -from cudf.core.buffer import Buffer +from cudf.core.buffer import Buffer, DeviceBufferLike, as_device_buffer_like from cudf.core.dtypes import ( CategoricalDtype, IntervalDtype, @@ -351,7 +353,7 @@ def valid_count(self) -> int: return len(self) - self.null_count @property - def nullmask(self) -> Buffer: + def nullmask(self) -> DeviceBufferLike: """The gpu buffer for the null-mask""" if not self.nullable: raise ValueError("Column has no null mask") @@ -423,19 +425,9 @@ def view(self, dtype: Dtype) -> ColumnBase: # This assertion prevents mypy errors below. assert self.base_data is not None - # If the view spans all of `base_data`, we return `base_data`. - if ( - self.offset == 0 - and self.base_data.size == self.size * self.dtype.itemsize - ): - view_buf = self.base_data - else: - view_buf = Buffer.from_buffer( - buffer=self.base_data, - size=self.size * self.dtype.itemsize, - offset=self.offset * self.dtype.itemsize, - ) - return build_column(view_buf, dtype=dtype) + start = self.offset * self.dtype.itemsize + end = start + self.size * self.dtype.itemsize + return build_column(self.base_data[start:end], dtype=dtype) def element_indexing(self, index: int): """Default implementation for indexing to an element @@ -767,12 +759,12 @@ def _obtain_isin_result(self, rhs: ColumnBase) -> ColumnBase: res = res.drop_duplicates(subset="orig_order", ignore_index=True) return res._data["bool"].fillna(False) - def as_mask(self) -> Buffer: + def as_mask(self) -> DeviceBufferLike: """Convert booleans to bitmask Returns ------- - Buffer + DeviceBufferLike """ if self.has_nulls(): @@ -1277,7 +1269,11 @@ def column_empty( data = None children = ( build_column( - data=Buffer.empty(row_count * cudf.dtype("int32").itemsize), + data=as_device_buffer_like( + rmm.DeviceBuffer( + size=row_count * cudf.dtype("int32").itemsize + ) + ), dtype="int32", ), ) @@ -1286,12 +1282,18 @@ def column_empty( children = ( full(row_count + 1, 0, dtype="int32"), build_column( - data=Buffer.empty(row_count * cudf.dtype("int8").itemsize), + data=as_device_buffer_like( + rmm.DeviceBuffer( + size=row_count * cudf.dtype("int8").itemsize + ) + ), dtype="int8", ), ) else: - data = Buffer.empty(row_count * dtype.itemsize) + data = as_device_buffer_like( + rmm.DeviceBuffer(size=row_count * dtype.itemsize) + ) if masked: mask = create_null_mask(row_count, state=MaskState.ALL_NULL) @@ -1304,11 +1306,11 @@ def column_empty( def build_column( - data: Union[Buffer, None], + data: Union[DeviceBufferLike, None], dtype: Dtype, *, size: int = None, - mask: Buffer = None, + mask: DeviceBufferLike = None, offset: int = 0, null_count: int = None, children: Tuple[ColumnBase, ...] = (), @@ -1318,12 +1320,12 @@ def build_column( Parameters ---------- - data : Buffer + data : DeviceBufferLike The data buffer (can be None if constructing certain Column types like StringColumn, ListColumn, or CategoricalColumn) dtype The dtype associated with the Column to construct - mask : Buffer, optional + mask : DeviceBufferLike, optional The mask buffer size : int, optional offset : int, optional @@ -1468,7 +1470,7 @@ def build_column( def build_categorical_column( categories: ColumnBase, codes: ColumnBase, - mask: Buffer = None, + mask: DeviceBufferLike = None, size: int = None, offset: int = 0, null_count: int = None, @@ -1484,7 +1486,7 @@ def build_categorical_column( codes : Column Column of codes, the size of the resulting Column will be the size of `codes` - mask : Buffer + mask : DeviceBufferLike Null mask size : int, optional offset : int, optional @@ -1528,7 +1530,7 @@ def build_interval_column( Column of values representing the left of the interval right_col : Column Column of representing the right of the interval - mask : Buffer + mask : DeviceBufferLike Null mask size : int, optional offset : int, optional @@ -1559,7 +1561,7 @@ def build_interval_column( def build_list_column( indices: ColumnBase, elements: ColumnBase, - mask: Buffer = None, + mask: DeviceBufferLike = None, size: int = None, offset: int = 0, null_count: int = None, @@ -1573,7 +1575,7 @@ def build_list_column( Column of list indices elements : ColumnBase Column of list elements - mask: Buffer + mask: DeviceBufferLike Null mask size: int, optional offset: int, optional @@ -1597,7 +1599,7 @@ def build_struct_column( names: Sequence[str], children: Tuple[ColumnBase, ...], dtype: Optional[Dtype] = None, - mask: Buffer = None, + mask: DeviceBufferLike = None, size: int = None, offset: int = 0, null_count: int = None, @@ -1611,7 +1613,7 @@ def build_struct_column( Field names to map to children dtypes, must be strings. children : tuple - mask: Buffer + mask: DeviceBufferLike Null mask size: int, optional offset: int, optional @@ -1647,7 +1649,9 @@ def _make_copy_replacing_NaT_with_null(column): out_col = cudf._lib.replace.replace( column, build_column( - Buffer(np.array([na_value], dtype=column.dtype).view("|u1")), + as_device_buffer_like( + np.array([na_value], dtype=column.dtype).view("|u1") + ), dtype=column.dtype, ), null, @@ -1742,7 +1746,7 @@ def as_column( ): arbitrary = cupy.ascontiguousarray(arbitrary) - data = _data_from_cuda_array_interface_desc(arbitrary) + data = as_device_buffer_like(arbitrary) col = build_column(data, dtype=current_dtype, mask=mask) if dtype is not None: @@ -1890,7 +1894,7 @@ def as_column( if cast_dtype: arbitrary = arbitrary.astype(cudf.dtype("datetime64[s]")) - buffer = Buffer(arbitrary.view("|u1")) + buffer = as_device_buffer_like(arbitrary.view("|u1")) mask = None if nan_as_null is None or nan_as_null is True: data = build_column(buffer, dtype=arbitrary.dtype) @@ -1908,7 +1912,7 @@ def as_column( if cast_dtype: arbitrary = arbitrary.astype(cudf.dtype("timedelta64[s]")) - buffer = Buffer(arbitrary.view("|u1")) + buffer = as_device_buffer_like(arbitrary.view("|u1")) mask = None if nan_as_null is None or nan_as_null is True: data = build_column(buffer, dtype=arbitrary.dtype) @@ -2187,17 +2191,7 @@ def _construct_array( return arbitrary -def _data_from_cuda_array_interface_desc(obj) -> Buffer: - desc = obj.__cuda_array_interface__ - ptr = desc["data"][0] - nelem = desc["shape"][0] if len(desc["shape"]) > 0 else 1 - dtype = cudf.dtype(desc["typestr"]) - - data = Buffer(data=ptr, size=nelem * dtype.itemsize, owner=obj) - return data - - -def _mask_from_cuda_array_interface_desc(obj) -> Union[Buffer, None]: +def _mask_from_cuda_array_interface_desc(obj) -> Union[DeviceBufferLike, None]: desc = obj.__cuda_array_interface__ mask = desc.get("mask", None) diff --git a/python/cudf/cudf/core/column/datetime.py b/python/cudf/cudf/core/column/datetime.py index 375a19f5423..1419b14e8c6 100644 --- a/python/cudf/cudf/core/column/datetime.py +++ b/python/cudf/cudf/core/column/datetime.py @@ -23,7 +23,7 @@ ) from cudf.api.types import is_datetime64_dtype, is_scalar, is_timedelta64_dtype from cudf.core._compat import PANDAS_GE_120 -from cudf.core.buffer import Buffer +from cudf.core.buffer import DeviceBufferLike from cudf.core.column import ColumnBase, as_column, column, string from cudf.core.column.timedelta import _unit_to_nanoseconds_conversion from cudf.utils.utils import _fillna_natwise @@ -98,11 +98,11 @@ class DatetimeColumn(column.ColumnBase): Parameters ---------- - data : Buffer + data : DeviceBufferLike The datetime values dtype : np.dtype The data type - mask : Buffer; optional + mask : DeviceBufferLike; optional The validity mask """ @@ -121,9 +121,9 @@ class DatetimeColumn(column.ColumnBase): def __init__( self, - data: Buffer, + data: DeviceBufferLike, dtype: DtypeObj, - mask: Buffer = None, + mask: DeviceBufferLike = None, size: int = None, # TODO: make non-optional offset: int = 0, null_count: int = None, @@ -131,7 +131,9 @@ def __init__( dtype = cudf.dtype(dtype) if data.size % dtype.itemsize: - raise ValueError("Buffer size must be divisible by element size") + raise ValueError( + "DeviceBufferLike size must be divisible by element size" + ) if size is None: size = data.size // dtype.itemsize size = size - offset diff --git a/python/cudf/cudf/core/column/decimal.py b/python/cudf/cudf/core/column/decimal.py index 69009106d15..e03802e6d8c 100644 --- a/python/cudf/cudf/core/column/decimal.py +++ b/python/cudf/cudf/core/column/decimal.py @@ -16,7 +16,7 @@ ) from cudf._typing import ColumnBinaryOperand, Dtype from cudf.api.types import is_integer_dtype, is_scalar -from cudf.core.buffer import Buffer +from cudf.core.buffer import as_device_buffer_like from cudf.core.column import ColumnBase, as_column from cudf.core.dtypes import ( Decimal32Dtype, @@ -203,7 +203,7 @@ def from_arrow(cls, data: pa.Array): data_128 = cp.array(np.frombuffer(data.buffers()[1]).view("int32")) data_32 = data_128[::4].copy() return cls( - data=Buffer(data_32.view("uint8")), + data=as_device_buffer_like(data_32.view("uint8")), size=len(data), dtype=dtype, offset=data.offset, @@ -211,7 +211,7 @@ def from_arrow(cls, data: pa.Array): ) def to_arrow(self): - data_buf_32 = self.base_data.to_host_array().view("int32") + data_buf_32 = np.array(self.base_data.memoryview()).view("int32") data_buf_128 = np.empty(len(data_buf_32) * 4, dtype="int32") # use striding to set the first 32 bits of each 128-bit chunk: @@ -231,7 +231,7 @@ def to_arrow(self): mask_buf = ( self.base_mask if self.base_mask is None - else pa.py_buffer(self.base_mask.to_host_array()) + else pa.py_buffer(self.base_mask.memoryview()) ) return pa.Array.from_buffers( type=self.dtype.to_arrow(), @@ -290,7 +290,7 @@ def from_arrow(cls, data: pa.Array): data_128 = cp.array(np.frombuffer(data.buffers()[1]).view("int64")) data_64 = data_128[::2].copy() return cls( - data=Buffer(data_64.view("uint8")), + data=as_device_buffer_like(data_64.view("uint8")), size=len(data), dtype=dtype, offset=data.offset, @@ -298,7 +298,7 @@ def from_arrow(cls, data: pa.Array): ) def to_arrow(self): - data_buf_64 = self.base_data.to_host_array().view("int64") + data_buf_64 = np.array(self.base_data.memoryview()).view("int64") data_buf_128 = np.empty(len(data_buf_64) * 2, dtype="int64") # use striding to set the first 64 bits of each 128-bit chunk: @@ -312,7 +312,7 @@ def to_arrow(self): mask_buf = ( self.base_mask if self.base_mask is None - else pa.py_buffer(self.base_mask.to_host_array()) + else pa.py_buffer(self.base_mask.memoryview()) ) return pa.Array.from_buffers( type=self.dtype.to_arrow(), diff --git a/python/cudf/cudf/core/column/lists.py b/python/cudf/cudf/core/column/lists.py index 32a71a31b83..0d5b351f69e 100644 --- a/python/cudf/cudf/core/column/lists.py +++ b/python/cudf/cudf/core/column/lists.py @@ -147,8 +147,7 @@ def to_arrow(self): pa_type = pa.list_(elements.type) if self.nullable: - nbuf = self.mask.to_host_array().view("int8") - nbuf = pa.py_buffer(nbuf) + nbuf = pa.py_buffer(self.mask.memoryview()) buffers = (nbuf, offsets.buffers()[1]) else: buffers = offsets.buffers() diff --git a/python/cudf/cudf/core/column/numerical.py b/python/cudf/cudf/core/column/numerical.py index 0529c614393..4b74dde129c 100644 --- a/python/cudf/cudf/core/column/numerical.py +++ b/python/cudf/cudf/core/column/numerical.py @@ -35,7 +35,7 @@ is_integer_dtype, is_number, ) -from cudf.core.buffer import Buffer +from cudf.core.buffer import DeviceBufferLike, as_device_buffer_like from cudf.core.column import ( ColumnBase, as_column, @@ -65,10 +65,10 @@ class NumericalColumn(NumericalBaseColumn): Parameters ---------- - data : Buffer + data : DeviceBufferLike dtype : np.dtype - The dtype associated with the data Buffer - mask : Buffer, optional + The dtype associated with the data DeviceBufferLike + mask : DeviceBufferLike, optional """ _nan_count: Optional[int] @@ -76,9 +76,9 @@ class NumericalColumn(NumericalBaseColumn): def __init__( self, - data: Buffer, + data: DeviceBufferLike, dtype: DtypeObj, - mask: Buffer = None, + mask: DeviceBufferLike = None, size: int = None, # TODO: make this non-optional offset: int = 0, null_count: int = None, @@ -86,7 +86,9 @@ def __init__( dtype = cudf.dtype(dtype) if data.size % dtype.itemsize: - raise ValueError("Buffer size must be divisible by element size") + raise ValueError( + "DeviceBufferLike size must be divisible by element size" + ) if size is None: size = (data.size // dtype.itemsize) - offset self._nan_count = None @@ -266,7 +268,7 @@ def normalize_binop_value( else: ary = full(len(self), other, dtype=other_dtype) return column.build_column( - data=Buffer(ary), + data=as_device_buffer_like(ary), dtype=ary.dtype, mask=self.mask, ) diff --git a/python/cudf/cudf/core/column/string.py b/python/cudf/cudf/core/column/string.py index d591008fa9a..726985fa091 100644 --- a/python/cudf/cudf/core/column/string.py +++ b/python/cudf/cudf/core/column/string.py @@ -32,7 +32,7 @@ is_scalar, is_string_dtype, ) -from cudf.core.buffer import Buffer +from cudf.core.buffer import DeviceBufferLike from cudf.core.column import column, datetime from cudf.core.column.methods import ColumnMethods from cudf.utils.docutils import copy_docstring @@ -5051,7 +5051,7 @@ class StringColumn(column.ColumnBase): Parameters ---------- - mask : Buffer + mask : DeviceBufferLike The validity mask offset : int Data offset @@ -5085,7 +5085,7 @@ class StringColumn(column.ColumnBase): def __init__( self, - mask: Buffer = None, + mask: DeviceBufferLike = None, size: int = None, # TODO: make non-optional offset: int = 0, null_count: int = None, diff --git a/python/cudf/cudf/core/column/struct.py b/python/cudf/cudf/core/column/struct.py index a9b6d4cad12..67ff3e48dbd 100644 --- a/python/cudf/cudf/core/column/struct.py +++ b/python/cudf/cudf/core/column/struct.py @@ -47,9 +47,7 @@ def to_arrow(self): ) if self.nullable: - nbuf = self.mask.to_host_array().view("int8") - nbuf = pa.py_buffer(nbuf) - buffers = (nbuf,) + buffers = (pa.py_buffer(self.mask.memoryview()),) else: buffers = (None,) diff --git a/python/cudf/cudf/core/column/timedelta.py b/python/cudf/cudf/core/column/timedelta.py index 3dc923e7ded..e6d688014fa 100644 --- a/python/cudf/cudf/core/column/timedelta.py +++ b/python/cudf/cudf/core/column/timedelta.py @@ -13,7 +13,7 @@ from cudf import _lib as libcudf from cudf._typing import ColumnBinaryOperand, DatetimeLikeScalar, Dtype from cudf.api.types import is_scalar, is_timedelta64_dtype -from cudf.core.buffer import Buffer +from cudf.core.buffer import DeviceBufferLike from cudf.core.column import ColumnBase, column, string from cudf.utils.dtypes import np_to_pa_dtype from cudf.utils.utils import _fillna_natwise @@ -40,13 +40,13 @@ class TimeDeltaColumn(ColumnBase): """ Parameters ---------- - data : Buffer + data : DeviceBufferLike The Timedelta values dtype : np.dtype The data type size : int Size of memory allocation. - mask : Buffer; optional + mask : DeviceBufferLike; optional The validity mask offset : int Data offset @@ -78,17 +78,19 @@ class TimeDeltaColumn(ColumnBase): def __init__( self, - data: Buffer, + data: DeviceBufferLike, dtype: Dtype, size: int = None, # TODO: make non-optional - mask: Buffer = None, + mask: DeviceBufferLike = None, offset: int = 0, null_count: int = None, ): dtype = cudf.dtype(dtype) if data.size % dtype.itemsize: - raise ValueError("Buffer size must be divisible by element size") + raise ValueError( + "DeviceBufferLike size must be divisible by element size" + ) if size is None: size = data.size // dtype.itemsize size = size - offset diff --git a/python/cudf/cudf/core/df_protocol.py b/python/cudf/cudf/core/df_protocol.py index f4ce658bff3..86b2d83ceec 100644 --- a/python/cudf/cudf/core/df_protocol.py +++ b/python/cudf/cudf/core/df_protocol.py @@ -18,7 +18,7 @@ from numba.cuda import as_cuda_array import cudf -from cudf.core.buffer import Buffer +from cudf.core.buffer import Buffer, DeviceBufferLike from cudf.core.column import as_column, build_categorical_column, build_column # Implementation of interchange protocol classes @@ -64,12 +64,12 @@ class _CuDFBuffer: def __init__( self, - buf: cudf.core.buffer.Buffer, + buf: DeviceBufferLike, dtype: np.dtype, allow_copy: bool = True, ) -> None: """ - Use cudf.core.buffer.Buffer object. + Use DeviceBufferLike object. """ # Store the cudf buffer where the data resides as a private # attribute, so we can use it to retrieve the public attributes @@ -80,9 +80,9 @@ def __init__( @property def bufsize(self) -> int: """ - Buffer size in bytes. + DeviceBufferLike size in bytes. """ - return self._buf.nbytes + return self._buf.size @property def ptr(self) -> int: @@ -622,11 +622,11 @@ def __dataframe__( Notes ----- -- Interpreting a raw pointer (as in ``Buffer.ptr``) is annoying and unsafe to - do in pure Python. It's more general but definitely less friendly than - having ``to_arrow`` and ``to_numpy`` methods. So for the buffers which lack - ``__dlpack__`` (e.g., because the column dtype isn't supported by DLPack), - this is worth looking at again. +- Interpreting a raw pointer (as in ``DeviceBufferLike.ptr``) is annoying and + unsafe to do in pure Python. It's more general but definitely less friendly + than having ``to_arrow`` and ``to_numpy`` methods. So for the buffers which + lack ``__dlpack__`` (e.g., because the column dtype isn't supported by + DLPack), this is worth looking at again. """ @@ -716,7 +716,7 @@ def _protocol_to_cudf_column_numeric( _dbuffer, _ddtype = buffers["data"] _check_buffer_is_on_gpu(_dbuffer) cudfcol_num = build_column( - Buffer(_dbuffer.ptr, _dbuffer.bufsize), + Buffer(data=_dbuffer.ptr, size=_dbuffer.bufsize, owner=None), protocol_dtype_to_cupy_dtype(_ddtype), ) return _set_missing_values(col, cudfcol_num), buffers @@ -746,7 +746,10 @@ def _set_missing_values( valid_mask = protocol_col.get_buffers()["validity"] if valid_mask is not None: bitmask = cp.asarray( - Buffer(valid_mask[0].ptr, valid_mask[0].bufsize), cp.bool8 + Buffer( + data=valid_mask[0].ptr, size=valid_mask[0].bufsize, owner=None + ), + cp.bool8, ) cudf_col[~bitmask] = None @@ -784,7 +787,8 @@ def _protocol_to_cudf_column_categorical( _check_buffer_is_on_gpu(codes_buffer) cdtype = protocol_dtype_to_cupy_dtype(codes_dtype) codes = build_column( - Buffer(codes_buffer.ptr, codes_buffer.bufsize), cdtype + Buffer(data=codes_buffer.ptr, size=codes_buffer.bufsize, owner=None), + cdtype, ) cudfcol = build_categorical_column( @@ -815,7 +819,7 @@ def _protocol_to_cudf_column_string( data_buffer, data_dtype = buffers["data"] _check_buffer_is_on_gpu(data_buffer) encoded_string = build_column( - Buffer(data_buffer.ptr, data_buffer.bufsize), + Buffer(data=data_buffer.ptr, size=data_buffer.bufsize, owner=None), protocol_dtype_to_cupy_dtype(data_dtype), ) @@ -825,7 +829,7 @@ def _protocol_to_cudf_column_string( offset_buffer, offset_dtype = buffers["offsets"] _check_buffer_is_on_gpu(offset_buffer) offsets = build_column( - Buffer(offset_buffer.ptr, offset_buffer.bufsize), + Buffer(data=offset_buffer.ptr, size=offset_buffer.bufsize, owner=None), protocol_dtype_to_cupy_dtype(offset_dtype), ) diff --git a/python/cudf/cudf/core/dtypes.py b/python/cudf/cudf/core/dtypes.py index 070837c127b..62717a3c5a8 100644 --- a/python/cudf/cudf/core/dtypes.py +++ b/python/cudf/cudf/core/dtypes.py @@ -20,7 +20,7 @@ from cudf._typing import Dtype from cudf.core._compat import PANDAS_GE_130 from cudf.core.abc import Serializable -from cudf.core.buffer import Buffer +from cudf.core.buffer import DeviceBufferLike def dtype(arbitrary): @@ -370,7 +370,7 @@ def serialize(self) -> Tuple[dict, list]: header: Dict[str, Any] = {} header["type-serialized"] = pickle.dumps(type(self)) - frames: List[Buffer] = [] + frames: List[DeviceBufferLike] = [] fields: Dict[str, Union[bytes, Tuple[Any, Tuple[int, int]]]] = {} diff --git a/python/cudf/cudf/core/index.py b/python/cudf/cudf/core/index.py index 9cbaa552e48..0fdcabc0e8b 100644 --- a/python/cudf/cudf/core/index.py +++ b/python/cudf/cudf/core/index.py @@ -2796,7 +2796,7 @@ def as_index(arbitrary, nan_as_null=None, **kwargs) -> BaseIndex: Currently supported inputs are: * ``Column`` - * ``Buffer`` + * ``DeviceBufferLike`` * ``Series`` * ``Index`` * numba device array diff --git a/python/cudf/cudf/core/series.py b/python/cudf/cudf/core/series.py index 7990abbb89a..ca7919b5c40 100644 --- a/python/cudf/cudf/core/series.py +++ b/python/cudf/cudf/core/series.py @@ -1830,8 +1830,8 @@ def data(self): 3 4 dtype: int64 >>> series.data - - >>> series.data.to_host_array() + + >>> np.array(series.data.memoryview()) array([1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0], dtype=uint8) """ # noqa: E501 diff --git a/python/cudf/cudf/tests/test_buffer.py b/python/cudf/cudf/tests/test_buffer.py index 4600d932c6f..16ba18581ed 100644 --- a/python/cudf/cudf/tests/test_buffer.py +++ b/python/cudf/cudf/tests/test_buffer.py @@ -1,8 +1,10 @@ +# Copyright (c) 2020-2022, NVIDIA CORPORATION. +from typing import Callable + import cupy as cp import pytest -from cupy.testing import assert_array_equal -from cudf.core.buffer import Buffer +from cudf.core.buffer import Buffer, DeviceBufferLike, as_device_buffer_like arr_len = 10 @@ -21,38 +23,82 @@ def test_buffer_from_cuda_iface_contiguous(data): data, expect_success = data if expect_success: - buf = Buffer(data=data.view("|u1"), size=data.size) # noqa: F841 + as_device_buffer_like(data.view("|u1")) else: with pytest.raises(ValueError): - buf = Buffer(data=data.view("|u1"), size=data.size) # noqa: F841 + as_device_buffer_like(data.view("|u1")) @pytest.mark.parametrize( "data", [ - (cp.zeros(arr_len)), - (cp.zeros((1, arr_len))), - (cp.zeros((1, arr_len, 1))), - (cp.zeros((arr_len, arr_len))), - (cp.zeros((arr_len, arr_len)).reshape(arr_len * arr_len)), + cp.arange(arr_len), + cp.arange(arr_len).reshape(1, arr_len), + cp.arange(arr_len).reshape(1, arr_len, 1), + cp.arange(arr_len**2).reshape(arr_len, arr_len), ], ) @pytest.mark.parametrize("dtype", ["uint8", "int8", "float32", "int32"]) def test_buffer_from_cuda_iface_dtype(data, dtype): data = data.astype(dtype) - if dtype not in ("uint8", "int8"): - with pytest.raises( - TypeError, match="Buffer data must be of uint8 type" - ): - buf = Buffer(data=data, size=data.size) # noqa: F841 - - -@pytest.mark.parametrize("size", [0, 1, 10, 100, 1000, 10_000]) -def test_buffer_copy(size): - data = cp.random.randint(low=0, high=100, size=size, dtype="u1") - buf = Buffer(data=data) - got = buf.copy() - assert got.size == buf.size - if size > 0: - assert got.ptr != buf.ptr - assert_array_equal(cp.asarray(buf), cp.asarray(got)) + buf = as_device_buffer_like(data) + ary = cp.array(buf).flatten().view("uint8") + assert (ary == buf).all() + + +@pytest.mark.parametrize("creator", [Buffer, as_device_buffer_like]) +def test_buffer_creation_from_any(creator: Callable[[object], Buffer]): + ary = cp.arange(arr_len) + b = creator(ary) + assert isinstance(b, DeviceBufferLike) + assert ary.__cuda_array_interface__["data"][0] == b.ptr + assert ary.nbytes == b.size + + with pytest.raises( + ValueError, match="size must be specified when `data` is an integer" + ): + Buffer(42) + + +@pytest.mark.parametrize( + "size,expect", [(10, "10B"), (2**10 + 500, "1.49KiB"), (2**20, "1MiB")] +) +def test_buffer_repr(size, expect): + ary = cp.arange(size, dtype="uint8") + buf = as_device_buffer_like(ary) + assert f"size={expect}" in repr(buf) + + +@pytest.mark.parametrize( + "idx", + [ + slice(0, 0), + slice(0, 1), + slice(-2, -1), + slice(0, arr_len), + slice(2, 3), + slice(2, -1), + ], +) +def test_buffer_slice(idx): + ary = cp.arange(arr_len, dtype="uint8") + buf = as_device_buffer_like(ary) + assert (ary[idx] == buf[idx]).all() + + +@pytest.mark.parametrize( + "idx, err_msg", + [ + (1, "index must be an slice"), + (slice(3, 2), "size cannot be negative"), + (slice(1, 2, 2), "slice must be contiguous"), + (slice(1, 2, -1), "slice must be contiguous"), + (slice(3, 2, -1), "slice must be contiguous"), + ], +) +def test_buffer_slice_fail(idx, err_msg): + ary = cp.arange(arr_len, dtype="uint8") + buf = as_device_buffer_like(ary) + + with pytest.raises(ValueError, match=err_msg): + buf[idx] diff --git a/python/cudf/cudf/tests/test_column.py b/python/cudf/cudf/tests/test_column.py index 2e5b121844a..4e2a26d31bd 100644 --- a/python/cudf/cudf/tests/test_column.py +++ b/python/cudf/cudf/tests/test_column.py @@ -406,7 +406,7 @@ def test_column_view_string_slice(slc): ) def test_as_column_buffer(data, expected): actual_column = cudf.core.column.as_column( - cudf.core.buffer.Buffer(data), dtype=data.dtype + cudf.core.buffer.as_device_buffer_like(data), dtype=data.dtype ) assert_eq(cudf.Series(actual_column), cudf.Series(expected)) diff --git a/python/cudf/cudf/tests/test_cuda_array_interface.py b/python/cudf/cudf/tests/test_cuda_array_interface.py index c4eacac2017..2a62a289747 100644 --- a/python/cudf/cudf/tests/test_cuda_array_interface.py +++ b/python/cudf/cudf/tests/test_cuda_array_interface.py @@ -179,9 +179,9 @@ def test_cuda_array_interface_pytorch(): got = cudf.Series(tensor) assert_eq(got, series) - from cudf.core.buffer import Buffer - - buffer = Buffer(cupy.ones(10, dtype=np.bool_)) + buffer = cudf.core.buffer.as_device_buffer_like( + cupy.ones(10, dtype=np.bool_) + ) tensor = torch.tensor(buffer) got = cudf.Series(tensor, dtype=np.bool_) diff --git a/python/cudf/cudf/tests/test_df_protocol.py b/python/cudf/cudf/tests/test_df_protocol.py index 21e18470b2f..c88b6ac9228 100644 --- a/python/cudf/cudf/tests/test_df_protocol.py +++ b/python/cudf/cudf/tests/test_df_protocol.py @@ -25,7 +25,8 @@ def assert_buffer_equal(buffer_and_dtype: Tuple[_CuDFBuffer, Any], cudfcol): device_id = cp.asarray(cudfcol.data).device.id assert buf.__dlpack_device__() == (2, device_id) col_from_buf = build_column( - Buffer(buf.ptr, buf.bufsize), protocol_dtype_to_cupy_dtype(dtype) + Buffer(data=buf.ptr, size=buf.bufsize, owner=None), + protocol_dtype_to_cupy_dtype(dtype), ) # check that non null values are the equals as nulls are represented # by sentinel values in the buffer. diff --git a/python/cudf/cudf/tests/test_pickling.py b/python/cudf/cudf/tests/test_pickling.py index 35ebd1b77c7..1427a214a72 100644 --- a/python/cudf/cudf/tests/test_pickling.py +++ b/python/cudf/cudf/tests/test_pickling.py @@ -7,7 +7,7 @@ import pytest from cudf import DataFrame, GenericIndex, RangeIndex, Series -from cudf.core.buffer import Buffer +from cudf.core.buffer import as_device_buffer_like from cudf.testing._utils import assert_eq if sys.version_info < (3, 8): @@ -97,7 +97,7 @@ def test_pickle_index(): def test_pickle_buffer(): arr = np.arange(10).view("|u1") - buf = Buffer(arr) + buf = as_device_buffer_like(arr) assert buf.size == arr.nbytes pickled = pickle.dumps(buf) unpacked = pickle.loads(pickled) diff --git a/python/cudf/cudf/tests/test_serialize.py b/python/cudf/cudf/tests/test_serialize.py index b7d679e95d5..61eee6bba43 100644 --- a/python/cudf/cudf/tests/test_serialize.py +++ b/python/cudf/cudf/tests/test_serialize.py @@ -356,8 +356,8 @@ def test_serialize_sliced_string(): # because both should be identical for i in range(3): assert_eq( - serialized_gd_series[1][i].to_host_array(), - serialized_sliced[1][i].to_host_array(), + serialized_gd_series[1][i].memoryview(), + serialized_sliced[1][i].memoryview(), ) recreated = cudf.Series.deserialize(*sliced.serialize()) diff --git a/python/cudf/cudf/tests/test_testing.py b/python/cudf/cudf/tests/test_testing.py index e5c78b6ea9a..60f01d567ef 100644 --- a/python/cudf/cudf/tests/test_testing.py +++ b/python/cudf/cudf/tests/test_testing.py @@ -429,7 +429,7 @@ def test_assert_column_memory_slice(arrow_arrays): def test_assert_column_memory_basic_same(arrow_arrays): data = cudf.core.column.ColumnBase.from_arrow(arrow_arrays) - buf = cudf.core.buffer.Buffer(data=data.base_data, owner=data) + buf = cudf.core.buffer.as_device_buffer_like(data.base_data) left = cudf.core.column.build_column(buf, dtype=np.int32) right = cudf.core.column.build_column(buf, dtype=np.int32) diff --git a/python/cudf/cudf/utils/string.py b/python/cudf/cudf/utils/string.py new file mode 100644 index 00000000000..9c02d1d6b34 --- /dev/null +++ b/python/cudf/cudf/utils/string.py @@ -0,0 +1,13 @@ +# Copyright (c) 2022, NVIDIA CORPORATION. + + +def format_bytes(nbytes: int) -> str: + """Format `nbytes` to a human readable string""" + n = float(nbytes) + for unit in ["B", "KiB", "MiB", "GiB", "TiB"]: + if abs(n) < 1024: + if n.is_integer(): + return f"{int(n)}{unit}" + return f"{n:.2f}{unit}" + n /= 1024 + return f"{n:.2f} PiB" diff --git a/python/cudf/cudf/utils/utils.py b/python/cudf/cudf/utils/utils.py index 29a33556efc..0e080c6f078 100644 --- a/python/cudf/cudf/utils/utils.py +++ b/python/cudf/cudf/utils/utils.py @@ -14,7 +14,7 @@ import cudf from cudf.core import column -from cudf.core.buffer import Buffer +from cudf.core.buffer import as_device_buffer_like # The size of the mask in bytes mask_dtype = cudf.dtype(np.int32) @@ -277,8 +277,8 @@ def pa_mask_buffer_to_mask(mask_buf, size): if mask_buf.size < mask_size: dbuf = rmm.DeviceBuffer(size=mask_size) dbuf.copy_from_host(np.asarray(mask_buf).view("u1")) - return Buffer(dbuf) - return Buffer(mask_buf) + return as_device_buffer_like(dbuf) + return as_device_buffer_like(mask_buf) def _isnat(val):