def test_radixsort(nelem, descending, dtype):
def expected_fn(key):
# Use mergesort for stable sort
# Negate the key for descending
if issubclass(dtype, np.integer):
def negate_values(v):
return ~key
else:
# Note: this doesn't work on the smallest value of integer
# i.e. -((int8)-128) -> -128
def negate_values(v):
return -key
sorted_idx = np.argsort(negate_values(key) if descending else key,
kind='mergesort')
sorted_keys = key[sorted_idx]
# Returns key, vals
return sorted_keys, sorted_idx
# Make data
key = gen_rand(dtype, nelem)
d_key = rmm.to_device(key)
col_key = new_column()
libgdf.gdf_column_view(col_key, unwrap_devary(d_key), ffi.NULL, nelem,
get_dtype(d_key.dtype))
val = np.arange(nelem, dtype=np.int64)
d_val = rmm.to_device(val)
col_val = new_column()
libgdf.gdf_column_view(col_val, unwrap_devary(d_val), ffi.NULL, nelem,
get_dtype(d_val.dtype))
sizeof_key = d_key.dtype.itemsize
sizeof_val = d_val.dtype.itemsize
begin_bit = 0
end_bit = sizeof_key * 8
# Setup plan
plan = libgdf.gdf_radixsort_plan(nelem, descending, begin_bit, end_bit)
libgdf.gdf_radixsort_plan_setup(plan, sizeof_key, sizeof_val)
# Sort
libgdf.gdf_radixsort_generic(plan, col_key, col_val)
# Cleanup
libgdf.gdf_radixsort_plan_free(plan)
# Check
got_keys = d_key.copy_to_host()
got_vals = d_val.copy_to_host()
sorted_keys, sorted_vals = expected_fn(key)
np.testing.assert_array_equal(sorted_keys, got_keys)
np.testing.assert_array_equal(sorted_vals, got_vals)
def apply_sort(sr_keys, sr_vals, ascending=True):
nelem = len(sr_keys)
begin_bit = 0
end_bit = sr_keys.dtype.itemsize * 8
plan = libgdf.gdf_radixsort_plan(nelem, not ascending, begin_bit, end_bit)
sizeof_key = sr_keys.dtype.itemsize
sizeof_val = sr_vals.dtype.itemsize
try:
libgdf.gdf_radixsort_plan_setup(plan, sizeof_key, sizeof_val)
libgdf.gdf_radixsort_generic(plan, sr_keys._cffi_view,
sr_vals._cffi_view)
finally:
libgdf.gdf_radixsort_plan_free(plan)
def apply_sort(col_keys, col_vals, ascending=True):
nelem = len(col_keys)
begin_bit = 0
end_bit = col_keys.dtype.itemsize * 8
plan = libgdf.gdf_radixsort_plan(nelem, not ascending, begin_bit, end_bit)
sizeof_key = col_keys.dtype.itemsize
sizeof_val = col_vals.dtype.itemsize
try:
libgdf.gdf_radixsort_plan_setup(plan, sizeof_key, sizeof_val)
libgdf.gdf_radixsort_generic(plan,
col_keys.cffi_view,
col_vals.cffi_view)
finally:
libgdf.gdf_radixsort_plan_free(plan)
def apply_sort(col_keys, col_vals, ascending=True):
"""Inplace sort
"""
nelem = len(col_keys)
begin_bit = 0
end_bit = col_keys.dtype.itemsize * 8
plan = libgdf.gdf_radixsort_plan(nelem, not ascending, begin_bit, end_bit)
sizeof_key = col_keys.dtype.itemsize
sizeof_val = col_vals.dtype.itemsize
try:
libgdf.gdf_radixsort_plan_setup(plan, sizeof_key, sizeof_val)
libgdf.gdf_radixsort_generic(plan, col_keys.cffi_view,
col_vals.cffi_view)
finally:
libgdf.gdf_radixsort_plan_free(plan)
