def _mask_from_cuda_array_interface_desc(desc): from cudf.utils.utils import calc_chunk_size, mask_dtype, mask_bitsize from cudf.utils.cudautils import compact_mask_bytes 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": mask = rmm.device_array_from_ptr( ptr, nelem=calc_chunk_size(nelem, mask_bitsize), dtype=mask_dtype, finalizer=None, ) mask = Buffer(mask) 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 apply_join(col_lhs, col_rhs, how, method='hash'): """Returns a tuple of the left and right joined indices as gpu arrays. """ if(len(col_lhs) != len(col_rhs)): msg = "Unequal #columns in list 'col_lhs' and list 'col_rhs'" raise ValueError(msg) joiner = _join_how_api[how] method_api = _join_method_api[method] gdf_context = ffi.new('gdf_context*') if method == 'hash': libgdf.gdf_context_view(gdf_context, 0, method_api, 0, 0, 0) elif method == 'sort': libgdf.gdf_context_view(gdf_context, 1, method_api, 0, 0, 0) else: msg = "method not supported" raise ValueError(msg) col_result_l = columnview(0, None, dtype=np.int32) col_result_r = columnview(0, None, dtype=np.int32) if(how in ['left', 'inner']): list_lhs = [] list_rhs = [] for i in range(len(col_lhs)): list_lhs.append(col_lhs[i].cffi_view) list_rhs.append(col_rhs[i].cffi_view) # Call libgdf joiner(len(col_lhs), list_lhs, list_rhs, col_result_l, col_result_r, gdf_context) else: joiner(col_lhs[0].cffi_view, col_rhs[0].cffi_view, col_result_l, col_result_r) # Extract result left = rmm.device_array_from_ptr(ptr=col_result_l.data, nelem=col_result_l.size, dtype=np.int32) right = rmm.device_array_from_ptr(ptr=col_result_r.data, nelem=col_result_r.size, dtype=np.int32) yield(left, right) libgdf.gdf_column_free(col_result_l) libgdf.gdf_column_free(col_result_r)
def cffi_view_to_column_mem(cffi_view): intaddr = int(ffi.cast("uintptr_t", cffi_view.data)) data = rmm.device_array_from_ptr(intaddr, nelem=cffi_view.size, dtype=gdf_to_np_dtype(cffi_view.dtype), finalizer=rmm._make_finalizer(intaddr, 0)) if cffi_view.valid: intaddr = int(ffi.cast("uintptr_t", cffi_view.valid)) mask = rmm.device_array_from_ptr( intaddr, nelem=calc_chunk_size(cffi_view.size, mask_bitsize), dtype=mask_dtype, finalizer=rmm._make_finalizer(intaddr, 0)) else: mask = None return data, mask
def _data_from_cuda_array_interface_desc(desc): ptr = desc["data"][0] nelem = desc["shape"][0] dtype = np.dtype(desc["typestr"]) data = rmm.device_array_from_ptr(ptr, nelem=nelem, dtype=dtype, finalizer=None) data = Buffer(data) return data
def cffi_view_to_column_mem(cffi_view): gdf_dtype = cffi_view.dtype if gdf_dtype == libgdf.GDF_STRING_CATEGORY: data_ptr = int(ffi.cast("uintptr_t", cffi_view.data)) # We need to create this just to make sure the memory is properly freed data = rmm.device_array_from_ptr(data_ptr, nelem=cffi_view.size, dtype='int32', finalizer=rmm._make_finalizer( data_ptr, 0)) nvcat_ptr = int(ffi.cast("uintptr_t", cffi_view.dtype_info.category)) nvcat_obj = nvcategory.bind_cpointer(nvcat_ptr) nvstr_obj = nvcat_obj.to_strings() mask = None if cffi_view.valid: mask_ptr = int(ffi.cast("uintptr_t", cffi_view.valid)) mask = rmm.device_array_from_ptr( mask_ptr, nelem=calc_chunk_size(cffi_view.size, mask_bitsize), dtype=mask_dtype, finalizer=rmm._make_finalizer(mask_ptr, 0)) return nvstr_obj, mask else: intaddr = int(ffi.cast("uintptr_t", cffi_view.data)) data = rmm.device_array_from_ptr( intaddr, nelem=cffi_view.size, dtype=gdf_to_np_dtype(cffi_view.dtype), finalizer=rmm._make_finalizer(intaddr, 0)) mask = None if cffi_view.valid: intaddr = int(ffi.cast("uintptr_t", cffi_view.valid)) mask = rmm.device_array_from_ptr( intaddr, nelem=calc_chunk_size(cffi_view.size, mask_bitsize), dtype=mask_dtype, finalizer=rmm._make_finalizer(intaddr, 0)) return data, mask
def libgdf_join(col_lhs, col_rhs, on, how, method='sort'): joiner = _join_how_api[how] method_api = _join_method_api[method] gdf_context = ffi.new('gdf_context*') libgdf.gdf_context_view(gdf_context, 0, method_api, 0, 0, 0) if how not in ['left', 'inner', 'outer']: msg = "new join api only supports left or inner" raise ValueError(msg) list_lhs = [] list_rhs = [] result_cols = [] result_col_names = [] left_idx = [] right_idx = [] # idx = 0 for name, col in col_lhs.items(): list_lhs.append(col._column.cffi_view) if name not in on: result_cols.append(columnview(0, None, dtype=col._column.dtype)) result_col_names.append(name) for name in on: result_cols.append(columnview(0, None, dtype=col_lhs[name]._column.dtype)) result_col_names.append(name) left_idx.append(list(col_lhs.keys()).index(name)) right_idx.append(list(col_rhs.keys()).index(name)) for name, col in col_rhs.items(): list_rhs.append(col._column.cffi_view) if name not in on: result_cols.append(columnview(0, None, dtype=col._column.dtype)) result_col_names.append(name) num_cols_to_join = len(on) result_num_cols = len(list_lhs) + len(list_rhs) - num_cols_to_join joiner(list_lhs, len(list_lhs), left_idx, list_rhs, len(list_rhs), right_idx, num_cols_to_join, result_num_cols, result_cols, ffi.NULL, ffi.NULL, gdf_context) res = [] valids = [] for col in result_cols: intaddr = int(ffi.cast("uintptr_t", col.data)) res.append(rmm.device_array_from_ptr(ptr=intaddr, nelem=col.size, dtype=gdf_to_np_dtype(col.dtype), finalizer=rmm._make_finalizer( intaddr, 0))) intaddr = int(ffi.cast("uintptr_t", col.valid)) valids.append(rmm.device_array_from_ptr(ptr=intaddr, nelem=calc_chunk_size( col.size, mask_bitsize), dtype=mask_dtype, finalizer=rmm._make_finalizer( intaddr, 0))) return res, valids
def gpu_view_as(buf, dtype, shape=None, strides=None): ptr = numba.cuda.cudadrv.driver.device_pointer(buf.to_numba()) return rmm.device_array_from_ptr( ptr, buf.size // dtype.itemsize, dtype=dtype )