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)