def test_prefixsum(dtype, nelem):
if dtype == np.int8:
# to keep data in range
data = gen_rand(dtype, nelem, low=-2, high=2)
else:
data = gen_rand(dtype, nelem)
d_data = rmm.to_device(data)
d_result = rmm.device_array(d_data.size, dtype=d_data.dtype)
col_data = new_column()
gdf_dtype = get_dtype(dtype)
libgdf.gdf_column_view(col_data, unwrap_devary(d_data), ffi.NULL, nelem,
gdf_dtype)
col_result = new_column()
libgdf.gdf_column_view(col_result, unwrap_devary(d_result), ffi.NULL,
nelem, gdf_dtype)
inclusive = True
libgdf.gdf_prefixsum_generic(col_data, col_result, inclusive)
expect = np.cumsum(d_data.copy_to_host())
got = d_result.copy_to_host()
if not inclusive:
expect = expect[:-1]
assert got[0] == 0
got = got[1:]
np.testing.assert_array_equal(expect, got)
def apply_prefixsum(col_inp, col_out, inclusive):
libgdf.gdf_prefixsum_generic(col_inp, col_out, inclusive)
def apply_prefixsum(col_inp, col_out, inclusive):
libgdf.gdf_prefixsum_generic(col_inp, col_out, inclusive)