def _mask_from_cuda_array_interface_desc(obj): from cudf.utils.utils import calc_chunk_size, mask_dtype, mask_bitsize from cudf.utils.cudautils import compact_mask_bytes desc = obj.__cuda_array_interface__ mask = desc.get("mask", None) if mask is not None: desc = mask.__cuda_array_interface__ ptr = desc["data"][0] nelem = desc["shape"][0] typestr = desc["typestr"] typecode = typestr[1] if typecode == "t": nelem = calc_chunk_size(nelem, mask_bitsize) mask = Buffer( data=ptr, size=nelem * mask_dtype.itemsize, owner=obj ) elif typecode == "b": dtype = np.dtype(typestr) mask = compact_mask_bytes( rmm.device_array_from_ptr( ptr, nelem=nelem, dtype=dtype, finalizer=None ) ) mask = Buffer(mask) else: raise NotImplementedError( f"Cannot infer mask from typestr {typestr}" ) return mask
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 else: with pytest.raises(ValueError): buf = Buffer(data=data.view("|u1"), size=data.size) # noqa: F841
def deserialize(cls, header, frames): dtype = header["dtype"] data = Buffer.deserialize(header["data"], [frames[0]]) mask = None if "mask" in header: mask = Buffer.deserialize(header["mask"], [frames[1]]) return build_column(data=data, dtype=dtype, mask=mask)
def pyarrow_buffer_to_cudf_buffer(arrow_buf, mask_size=0): """ Given a PyArrow Buffer backed by either host or device memory, convert it to a cuDF Buffer """ from cudf._lib.arrow._cuda import CudaBuffer as arrowCudaBuffer # Try creating a PyArrow CudaBuffer from the PyArrow Buffer object, it # fails with an ArrowTypeError if it's a host based Buffer so we catch and # process as expected if not isinstance(arrow_buf, pa.Buffer): raise TypeError( "Expected type: {}, got type: {}".format( pa.Buffer.__name__, type(arrow_buf).__name__ ) ) try: arrow_cuda_buf = arrowCudaBuffer.from_buffer(arrow_buf) buf = Buffer( data=arrow_cuda_buf.address, size=arrow_cuda_buf.size, owner=arrow_cuda_buf, ) if buf.size < mask_size: dbuf = rmm.DeviceBuffer(size=mask_size) dbuf.copy_from_device(buf) return Buffer(dbuf) return buf except pa.ArrowTypeError: if arrow_buf.size < mask_size: dbuf = rmm.DeviceBuffer(size=mask_size) dbuf.copy_from_host(np.asarray(arrow_buf).view("u1")) return Buffer(dbuf) return Buffer(arrow_buf)
def buffers_from_pyarrow(pa_arr, dtype=None): from cudf.core.buffer import Buffer from cudf.utils.cudautils import copy_array buffers = pa_arr.buffers() if buffers[0]: mask_dev_array = make_mask(len(pa_arr)) arrow_dev_array = rmm.to_device(np.array(buffers[0]).view("int8")) copy_array(arrow_dev_array, mask_dev_array) pamask = Buffer(mask_dev_array) else: pamask = None if dtype: new_dtype = dtype else: if isinstance(pa_arr, pa.DictionaryArray): new_dtype = pa_arr.indices.type.to_pandas_dtype() else: new_dtype = pa_arr.type.to_pandas_dtype() if buffers[1]: padata = Buffer( np.array(buffers[1]).view(new_dtype)[pa_arr.offset:pa_arr.offset + len(pa_arr)]) else: padata = Buffer(np.empty(0, dtype=new_dtype)) return (pamask, padata)
def from_numpy(cls, array): cast_dtype = array.dtype.type == np.int64 if array.dtype.kind == "M": time_unit, _ = np.datetime_data(array.dtype) cast_dtype = time_unit in ("D", "W", "M", "Y") or ( len(array) > 0 and (isinstance(array[0], str) or isinstance(array[0], dt.datetime))) elif not cast_dtype: raise ValueError( ("Cannot infer datetime dtype " + "from np.array dtype `%s`") % (array.dtype)) if cast_dtype: array = array.astype(np.dtype("datetime64[s]")) assert array.dtype.itemsize == 8 mask = None if np.any(np.isnat(array)): null = cudf.core.column.column_empty_like(array, masked=True, newsize=1) col = libcudf.replace.replace( as_column(Buffer(array), dtype=array.dtype), as_column( Buffer(np.array([np.datetime64("NaT")], dtype=array.dtype)), dtype=array.dtype, ), null, ) mask = col.mask return cls(data=Buffer(array), mask=mask, dtype=array.dtype)
def deserialize(cls, header, frames): dtype = header["dtype"] data = Buffer(frames[0]) mask = None if header["frame_count"] > 1: mask = Buffer(frames[1]) return build_column(data=data, dtype=dtype, mask=mask)
def column_empty(row_count, dtype="object", masked=False): """Allocate a new column like the given row_count and dtype. """ dtype = pd.api.types.pandas_dtype(dtype) children = () if is_categorical_dtype(dtype): data = None children = (build_column( data=Buffer.empty(row_count * np.dtype("int32").itemsize), dtype="int32", ), ) elif dtype.kind in "OU": data = None children = ( build_column( data=Buffer.empty( (row_count + 1) * np.dtype("int32").itemsize), dtype="int32", ), build_column( data=Buffer.empty(row_count * np.dtype("int8").itemsize), dtype="int8", ), ) else: data = Buffer.empty(row_count * dtype.itemsize) if masked: mask = Buffer(cudautils.make_empty_mask(row_count)) else: mask = None return build_column(data, dtype, mask=mask, children=children)
def prefixsum(vals): """Compute the full prefixsum. Given the input of N. The output size is N + 1. The first value is always 0. The last value is the sum of *vals*. """ import cudf._lib as libcudf from cudf.core.column import NumericalColumn from cudf.core.buffer import Buffer # Allocate output slots = rmm.device_array(shape=vals.size + 1, dtype=vals.dtype) # Fill 0 to slot[0] gpu_fill_value[1, 1](slots[:1], 0) # Compute prefixsum on the mask in_col = NumericalColumn(data=Buffer(vals), mask=None, null_count=0, dtype=vals.dtype) out_col = NumericalColumn(data=Buffer(slots[1:]), mask=None, null_count=0, dtype=vals.dtype) libcudf.reduce.scan(in_col, out_col, "sum", inclusive=True) return slots
def column_empty(row_count, dtype="object", masked=False): """Allocate a new column like the given row_count and dtype. """ dtype = pd.api.types.pandas_dtype(dtype) children = () if is_categorical_dtype(dtype): data = None children = (build_column( data=Buffer.empty(row_count * np.dtype("int32").itemsize), dtype="int32", ), ) elif dtype.kind in "OU": data = None children = ( build_column( data=Buffer(cupy.zeros(row_count + 1, dtype="int32")), dtype="int32", ), build_column( data=Buffer.empty(row_count * np.dtype("int8").itemsize), dtype="int8", ), ) else: data = Buffer.empty(row_count * dtype.itemsize) if masked: mask = create_null_mask(row_count, state=MaskState.ALL_NULL) else: mask = None return build_column(data, dtype, mask=mask, children=children)
def _initialize_read(self): from cudf.core.buffer import Buffer from cupy.cuda.memory import UnownedMemory self._offset = 0 self._has_read_headers = False self._buffers = [] headers, buffers = _id_to_buffers[self._object_id] self._headers = headers = headers.copy() buffer_types = [] for buf in buffers: if isinstance(buf, cupy.ndarray): ptr, size = buf.data.ptr, buf.size self._buffers.append( UnownedMemory(ptr, size, Buffer(ptr, size))) buffer_types.append(["cuda", size]) elif isinstance(buf, Buffer): ptr, size = buf.ptr, buf.size if size == 0: # empty buffer cannot construct a UnownedMemory self._buffers.append(None) else: self._buffers.append( UnownedMemory(ptr, size, Buffer(ptr, size))) buffer_types.append(["cuda", size]) else: size = getattr(buf, "size", len(buf)) self._buffers.append(buf) buffer_types.append(["memory", size]) headers["buffer_types"] = buffer_types
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))
def pa_mask_buffer_to_mask(mask_buf, size): """ Convert PyArrow mask buffer to cuDF mask buffer """ mask_size = cudf._lib.null_mask.bitmask_allocation_size_bytes(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)
def _write_cuda_buffer(ptr): # pragma: no cover # copy cuda buffer to host chunk_size = CUDA_CHUNK_SIZE offset = 0 nbytes = buffer.nbytes while offset < nbytes: size = chunk_size if (offset + chunk_size) < nbytes else nbytes - offset chunk_buffer = CPBuffer(ptr + offset, size=size) # copy chunk to host memoryview writer.write(chunk_buffer.host_serialize()[1][0]) offset += size
def children(self): if self._children is None: codes_column = self.base_children[0] buf = Buffer(codes_column.base_data) buf.ptr = buf.ptr + (self.offset * codes_column.dtype.itemsize) buf.size = self.size * codes_column.dtype.itemsize codes_column = column.build_column( data=buf, dtype=codes_column.dtype, size=self.size, ) self._children = (codes_column,) return self._children
def __init__(self, data, null_count=None, name=None, **kwargs): """ Parameters ---------- data : nvstrings.nvstrings The nvstrings object null_count : int; optional The number of null values in the mask. """ from collections.abc import Sequence if isinstance(data, Sequence): data = nvstrings.to_device(data) assert isinstance(data, nvstrings.nvstrings) self._data = data self._dtype = np.dtype("object") self._name = name if null_count is None: null_count = data.null_count() self._null_count = null_count self._mask = None if self._null_count > 0: mask_size = utils.calc_chunk_size(len(self.data), utils.mask_bitsize) out_mask_arr = rmm.device_array(mask_size, dtype="int8") out_mask_ptr = libcudf.cudf.get_ctype_ptr(out_mask_arr) self.data.set_null_bitmask(out_mask_ptr, bdevmem=True) self._mask = Buffer(out_mask_arr) self._nvcategory = None self._indices = None
def __init__(self, dtype, mask=None, offset=0, children=()): """ Parameters ---------- dtype : CategoricalDtype mask : Buffer The validity mask offset : int Data offset children : Tuple[Column] Two non-null columns containing the categories and codes respectively """ data = Buffer.empty(0) size = children[0].size if isinstance(dtype, pd.api.types.CategoricalDtype): dtype = CategoricalDtype.from_pandas(dtype) if not isinstance(dtype, CategoricalDtype): raise ValueError("dtype must be instance of CategoricalDtype") super().__init__( data, size=size, dtype=dtype, mask=mask, offset=offset, children=children, ) self._codes = None
def children(self) -> Tuple[NumericalColumn]: if self._children is None: codes_column = self.base_children[0] buf = Buffer(codes_column.base_data) buf.ptr = buf.ptr + (self.offset * codes_column.dtype.itemsize) buf.size = self.size * codes_column.dtype.itemsize codes_column = cast( cudf.core.column.NumericalColumn, column.build_column( data=buf, dtype=codes_column.dtype, size=self.size, ), ) self._children = (codes_column,) return self._children
def deserialize(cls, header, frames): # Get null mask if header["null_count"] > 0: mask = Buffer(frames[-1]) else: mask = None # Deserialize child columns children = [] f = 0 for h in header["subheaders"]: fcount = h["frame_count"] child_frames = frames[f:f + fcount] column_type = pickle.loads(h["type-serialized"]) children.append(column_type.deserialize(h, child_frames)) f += fcount # Materialize list column return column.build_column( data=None, dtype=pickle.loads(header["dtype"]), mask=mask, children=tuple(children), size=header["size"], )
def __init__(self, mask=None, offset=0, children=()): """ Parameters ---------- mask : Buffer The validity mask offset : int Data offset children : Tuple[Column] Two non-null columns containing the string data and offsets respectively """ data = Buffer.empty(0) dtype = np.dtype("object") if children[0].size == 0: size = 0 else: # one less because the last element of offsets is the number of # bytes in the data buffer size = children[0].size - 1 super().__init__(data, size, dtype, mask=mask, children=children) self._nvstrings = None self._nvcategory = None self._indices = None
def buffers_from_pyarrow(pa_arr, dtype=None): """ Given a pyarrow array returns a 5 length tuple of: - size - offset - cudf.Buffer --> mask - cudf.Buffer --> data - cudf.Buffer --> string characters """ from cudf._libxx.null_mask import bitmask_allocation_size_bytes buffers = pa_arr.buffers() if pa_arr.null_count: mask_size = bitmask_allocation_size_bytes(len(pa_arr)) pamask = pyarrow_buffer_to_cudf_buffer(buffers[0], mask_size=mask_size) else: pamask = None offset = pa_arr.offset size = len(pa_arr) if buffers[1]: padata = pyarrow_buffer_to_cudf_buffer(buffers[1]) else: padata = Buffer.empty(0) pastrs = None if isinstance(pa_arr, pa.StringArray): pastrs = pyarrow_buffer_to_cudf_buffer(buffers[2]) return (size, offset, pamask, padata, pastrs)
def as_numerical(self): from cudf.core.column import numerical data = Buffer(self.data.mem.view(np.int64)) return self.view(numerical.NumericalColumn, data=data, dtype=data.dtype)
def normalize_binop_value( self, other: ScalarLike ) -> Union[ColumnBase, ScalarLike]: if other is None: return other if isinstance(other, cudf.Scalar): if self.dtype == other.dtype: return other # expensive device-host transfer just to # adjust the dtype other = other.value elif isinstance(other, np.ndarray) and other.ndim == 0: other = other.item() other_dtype = np.min_scalar_type(other) if other_dtype.kind in {"b", "i", "u", "f"}: if isinstance(other, cudf.Scalar): return other other_dtype = np.promote_types(self.dtype, other_dtype) if other_dtype == np.dtype("float16"): other_dtype = np.dtype("float32") other = other_dtype.type(other) if self.dtype.kind == "b": other_dtype = min_signed_type(other) if np.isscalar(other): other = np.dtype(other_dtype).type(other) return other else: ary = utils.scalar_broadcast_to( other, size=len(self), dtype=other_dtype ) return column.build_column( data=Buffer(ary), dtype=ary.dtype, mask=self.mask, ) else: raise TypeError(f"cannot broadcast {type(other)}")
def normalize_binop_value(self, other): if other is None: return other other_dtype = np.min_scalar_type(other) if other_dtype.kind in {"b", "i", "u", "f"}: other_dtype = np.promote_types(self.dtype, other_dtype) if other_dtype == np.dtype("float16"): other = np.dtype("float32").type(other) other_dtype = other.dtype if self.dtype.kind == "b": other_dtype = min_signed_type(other) if np.isscalar(other): other = np.dtype(other_dtype).type(other) return other else: ary = utils.scalar_broadcast_to(other, size=len(self), dtype=other_dtype) return column.build_column( data=Buffer.from_array_lik(ary), dtype=ary.dtype, mask=self.mask, ) else: raise TypeError("cannot broadcast {}".format(type(other)))
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
def deserialize(cls, header: dict, frames: list) -> CategoricalColumn: n_dtype_frames = header["dtype_frames_count"] dtype = CategoricalDtype.deserialize(header["dtype"], frames[:n_dtype_frames]) n_data_frames = header["data_frames_count"] column_type = pickle.loads(header["data"]["type-serialized"]) data = column_type.deserialize( header["data"], frames[n_dtype_frames:n_dtype_frames + n_data_frames], ) mask = None if "mask" in header: mask = Buffer.deserialize(header["mask"], [frames[n_dtype_frames + n_data_frames]]) return cast( CategoricalColumn, column.build_column( data=None, dtype=dtype, mask=mask, children=(column.as_column(data.base_data, dtype=data.dtype), ), ), )
def as_numerical_column(self, dtype, **kwargs): mem_dtype = np.dtype(dtype) str_dtype = mem_dtype out_dtype = mem_dtype if mem_dtype.type in (np.int8, np.int16): mem_dtype = np.dtype(np.int32) str_dtype = mem_dtype elif mem_dtype.type is np.datetime64: kwargs.update(units=np.datetime_data(mem_dtype)[0]) mem_dtype = np.dtype(np.int64) out_arr = rmm.device_array(shape=len(self), dtype=mem_dtype) out_ptr = libcudf.cudf.get_ctype_ptr(out_arr) kwargs.update({"devptr": out_ptr}) _str_to_numeric_typecast_functions[str_dtype](self.str(), **kwargs) out_col = column.as_column(out_arr) if self.null_count > 0: mask_size = utils.calc_chunk_size(len(self.data), utils.mask_bitsize) out_mask_arr = rmm.device_array(mask_size, dtype="int8") out_mask_ptr = libcudf.cudf.get_ctype_ptr(out_mask_arr) self.data.set_null_bitmask(out_mask_ptr, bdevmem=True) mask = Buffer(out_mask_arr) out_col = out_col.set_mask(mask) return out_col.astype(out_dtype)
def from_dlpack(pycapsule_obj): """Converts from a DLPack tensor to a cuDF object. DLPack is an open-source memory tensor structure: `dmlc/dlpack <https://github.com/dmlc/dlpack>`_. This function takes a PyCapsule object which contains a pointer to a DLPack tensor as input, and returns a cuDF object. This function deep copies the data in the DLPack tensor into a cuDF object. Parameters ---------- pycapsule_obj : PyCapsule Input DLPack tensor pointer which is encapsulated in a PyCapsule object. Returns ------- A cuDF DataFrame or Series depending on if the input DLPack tensor is 1D or 2D. """ try: res, valids = cpp_dlpack.from_dlpack(pycapsule_obj) except GDFError as err: if str(err) == "b'GDF_DATASET_EMPTY'": raise ValueError( "Cannot create a cuDF Object from a DLPack tensor of 0 size" ) else: raise err cols = [] for idx in range(len(valids)): mask = None if valids[idx]: mask = Buffer(valids[idx]) cols.append( column.build_column( Buffer(res[idx]), dtype=res[idx].dtype, mask=mask ) ) if len(cols) == 1: return Series(cols[0]) else: df = DataFrame() for idx, col in enumerate(cols): df[idx] = col return df
def test_pickle_buffer(): arr = np.arange(10) buf = Buffer(arr) assert buf.size == arr.nbytes pickled = pickle.dumps(buf) unpacked = pickle.loads(pickled) # Check that unpacked capacity equals buf.size assert unpacked.size == arr.nbytes
def _data_from_cuda_array_interface_desc(obj): desc = obj.__cuda_array_interface__ ptr = desc["data"][0] nelem = desc["shape"][0] dtype = np.dtype(desc["typestr"]) data = Buffer(data=ptr, size=nelem * dtype.itemsize, owner=obj) return data