def _columnview(size, data, mask, dtype, null_count):
colview = ffi.new('gdf_column*')
if null_count is None:
libgdf.gdf_column_view(
colview,
data,
mask,
size,
np_to_gdf_dtype(dtype),
)
else:
libgdf.gdf_column_view_augmented(
colview,
data,
mask,
size,
np_to_gdf_dtype(dtype),
null_count,
)
return colview
def test_prefixsum_masked(dtype, nelem):
if dtype == np.int8:
data = gen_rand(dtype, nelem, low=-2, high=2)
else:
data = gen_rand(dtype, nelem)
mask = gen_rand(np.int8, (nelem + 8 - 1) // 8)
dummy_mask = gen_rand(np.int8, (nelem + 8 - 1) // 8)
d_data = rmm.to_device(data)
d_mask = rmm.to_device(mask)
d_result = rmm.device_array(d_data.size, dtype=d_data.dtype)
d_result_mask = rmm.to_device(dummy_mask)
gdf_dtype = get_dtype(dtype)
extra_dtype_info = ffi.new('gdf_dtype_extra_info*')
extra_dtype_info.time_unit = libgdf.TIME_UNIT_NONE
col_data = new_column()
libgdf.gdf_column_view_augmented(col_data, unwrap_devary(d_data),
unwrap_devary(d_mask), nelem, gdf_dtype,
count_nulls(d_mask, nelem),
extra_dtype_info[0])
col_result = new_column()
libgdf.gdf_column_view(col_result, unwrap_devary(d_result),
unwrap_devary(d_result_mask), nelem, gdf_dtype)
inclusive = True
libgdf.gdf_prefixsum(col_data, col_result, inclusive)
boolmask = buffer_as_bits(mask)[:nelem]
expect = np.cumsum(data[boolmask])
got = d_result.copy_to_host()[boolmask]
if not inclusive:
expect = expect[:-1]
assert got[0] == 0
got = got[1:]
decimal = 4 if dtype == np.float32 else 6
np.testing.assert_array_almost_equal(expect, got, decimal=decimal)
def _columnview(size, data, mask, dtype, null_count):
colview = ffi.new('gdf_column*')
extra_dtype_info = ffi.new('gdf_dtype_extra_info*')
extra_dtype_info.time_unit = libgdf.TIME_UNIT_NONE
if null_count is None:
libgdf.gdf_column_view(
colview,
data,
mask,
size,
np_to_gdf_dtype(dtype),
)
else:
libgdf.gdf_column_view_augmented(
colview,
data,
mask,
size,
np_to_gdf_dtype(dtype),
null_count,
extra_dtype_info[0],
)
return colview
def _columnview(size, data, mask, dtype, null_count, nvcat):
colview = ffi.new('gdf_column*')
extra_dtype_info = ffi.new('gdf_dtype_extra_info*')
extra_dtype_info.time_unit = libgdf.TIME_UNIT_NONE
if nvcat is not None:
extra_dtype_info.category = ffi.cast('void*', nvcat.get_cpointer())
else:
extra_dtype_info.category = ffi.NULL
if mask is None:
null_count = 0
mask = ffi.NULL
libgdf.gdf_column_view_augmented(
colview,
data,
mask,
size,
np_to_gdf_dtype(dtype),
null_count,
extra_dtype_info[0],
)
return colview