Skip to content

Migrate T-digest test helpers to memory_resources #23608

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
bdice wants to merge 1 commit into
base: main
Choose a base branch
from
+208 −68
Draft

Migrate T-digest test helpers to memory_resources #23608

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
152 changes: 109 additions & 43 deletions cpp/include/cudf_test/tdigest_utilities.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/

Expand Down Expand Up @@ -65,12 +65,24 @@ inline T rand_range(T min, T max)
return min + static_cast<T>(frand<T>() * (max - min));
}

/**
* @brief Generate a typed column from a bucketed percentile distribution.
*
* @param buckets Upper bound for each generated bucket
* @param sizes Number of values generated for each bucket
* @param t Data type of the returned column
* @param sorted Whether to sort the generated values before conversion
* @param mr Memory resources used for returned and temporary allocations
* @return Generated column
*/
inline std::unique_ptr<column> generate_typed_percentile_distribution(
std::vector<double> const& buckets,
std::vector<int> const& sizes,
data_type t,
bool sorted = false)
bool sorted = false,
cudf::memory_resources mr = cudf::get_current_device_resource_ref())
{
auto const temporary_mr = mr.get_temporary_mr();
srand(0);

std::vector<double> values;
Expand All @@ -87,40 +99,62 @@ inline std::unique_ptr<column> generate_typed_percentile_distribution(

if (sorted) { std::sort(values.begin(), values.end()); }

cudf::test::fixed_width_column_wrapper<double> src(values.begin(), values.end());
return cudf::cast(src, t);
cudf::test::fixed_width_column_wrapper<double> src(values.begin(), values.end(), temporary_mr);
return cudf::cast(src, t, cudf::get_default_stream(), mr.get_output_mr());
}

// "standardized" means the parameters sent into generate_typed_percentile_distribution. the intent
// is to provide a standardized set of inputs for use with tdigest generation tests and
// percentile_approx tests. std::vector<double>
// buckets{10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0}; std::vector<int>
// sizes{50000, 50000, 50000, 50000, 50000, 100000, 100000, 100000, 100000, 100000};
/**
* @brief Generate the standardized percentile distribution used by T-digest tests.
*
* @param t Data type of the returned column
* @param sorted Whether to sort the generated values before conversion
* @param mr Memory resources used for returned and temporary allocations
* @return Generated column
*/
inline std::unique_ptr<column> generate_standardized_percentile_distribution(
data_type t = data_type{type_id::FLOAT64}, bool sorted = false)
data_type t = data_type{type_id::FLOAT64},
bool sorted = false,
cudf::memory_resources mr = cudf::get_current_device_resource_ref())
{
std::vector<double> buckets{10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f, 70.0f, 80.0, 90.0f, 100.0f};
std::vector<int> b_sizes{
50000, 50000, 50000, 50000, 50000, 100000, 100000, 100000, 100000, 100000};
return generate_typed_percentile_distribution(buckets, b_sizes, t, sorted);
return generate_typed_percentile_distribution(buckets, b_sizes, t, sorted, mr);
}

/**
* @brief Compare a tdigest column against a sampling of expected values.
*
* @param tdv T-digest column to validate
* @param h_expected Expected centroid index, mean, and weight tuples
* @param mr Memory resources used for temporary device allocations
*/
void tdigest_sample_compare(cudf::tdigest::tdigest_column_view const& tdv,
std::vector<expected_value> const& h_expected);
std::vector<expected_value> const& h_expected,
cudf::memory_resources mr = cudf::get_current_device_resource_ref());

/**
* @brief Compare the min/max values of a tdigest against inputs.
*
* @tparam T Input element type
* @param tdv T-digest column to validate
* @param input_values Values whose extrema are expected in the T-digest
* @param mr Memory resources used for temporary device allocations
*/
template <typename T>
void tdigest_minmax_compare(cudf::tdigest::tdigest_column_view const& tdv,
cudf::column_view const& input_values)
cudf::column_view const& input_values,
cudf::memory_resources mr = cudf::get_current_device_resource_ref())
{
using ScalarType = cudf::scalar_type_t<T>;
auto const temporary_mr = mr.get_temporary_mr();
using ScalarType = cudf::scalar_type_t<T>;

auto [col_min, col_max] = cudf::minmax(input_values);
auto [col_min, col_max] = cudf::minmax(input_values, cudf::get_default_stream(), temporary_mr);

auto min_scalar = static_cast<ScalarType*>(col_min.get());
auto max_scalar = static_cast<ScalarType*>(col_max.get());
Expand Down Expand Up @@ -148,106 +182,135 @@ struct expected_tdigest {

/**
* @brief Create an expected tdigest column given component inputs.
*
* @param groups T-digest component values for each output row
* @param mr Memory resources used for returned and temporary allocations
* @return Expected T-digest column
*/
std::unique_ptr<column> make_expected_tdigest_column(std::vector<expected_tdigest> const& groups);
std::unique_ptr<column> make_expected_tdigest_column(
std::vector<expected_tdigest> const& groups,
cudf::memory_resources mr = cudf::get_current_device_resource_ref());

// shared test for groupby/reduction.
template <typename T, typename Func>
void tdigest_simple_aggregation(Func op)
void tdigest_simple_aggregation(Func op,
cudf::memory_resources mr = cudf::get_current_device_resource_ref())
{
auto const temporary_mr = mr.get_temporary_mr();
auto const temporary_resources = cudf::memory_resources{temporary_mr, temporary_mr};
bool is_cpu_cluster_computation_disabled[2] = {true, false};
for (int idx = 0; idx < 2; idx++) {
cudf::tdigest::detail::is_cpu_cluster_computation_disabled =
is_cpu_cluster_computation_disabled[idx];

// create a tdigest that has far fewer values in it than the delta value. this should result
// in every value remaining uncompressed
cudf::test::fixed_width_column_wrapper<T> values{126, 15, 1, 99, 67};
cudf::test::fixed_width_column_wrapper<T> values({126, 15, 1, 99, 67}, temporary_resources);
int const delta = 1000;
auto result = cudf::type_dispatcher(
static_cast<column_view>(values).type(), tdigest_gen{}, op, values, delta);

cudf::test::fixed_width_column_wrapper<T> raw_mean({1, 15, 67, 99, 126});
cudf::test::fixed_width_column_wrapper<double> weight{1, 1, 1, 1, 1};
auto mean = cudf::cast(raw_mean, data_type{type_id::FLOAT64});
cudf::test::fixed_width_column_wrapper<T> raw_mean({1, 15, 67, 99, 126}, temporary_resources);
cudf::test::fixed_width_column_wrapper<double> weight({1, 1, 1, 1, 1}, temporary_resources);
auto mean =
cudf::cast(raw_mean, data_type{type_id::FLOAT64}, cudf::get_default_stream(), temporary_mr);
double const min = 1;
double const max = 126;
auto expected = make_expected_tdigest_column({{*mean,
weight,
static_cast<double>(static_cast<T>(min)),
static_cast<double>(static_cast<T>(max))}});
static_cast<double>(static_cast<T>(max))}},
temporary_resources);

CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, *expected);
CUDF_TEST_EXPECT_COLUMNS_EQUAL(
*result, *expected, cudf::test::debug_output_level::FIRST_ERROR, mr);
}
}

// shared test for groupby/reduction.
template <typename T, typename Func>
void tdigest_simple_with_nulls_aggregation(Func op)
void tdigest_simple_with_nulls_aggregation(
Func op, cudf::memory_resources mr = cudf::get_current_device_resource_ref())
{
auto const temporary_mr = mr.get_temporary_mr();
auto const temporary_resources = cudf::memory_resources{temporary_mr, temporary_mr};
bool is_cpu_cluster_computation_disabled[2] = {true, false};
for (int idx = 0; idx < 2; idx++) {
cudf::tdigest::detail::is_cpu_cluster_computation_disabled =
is_cpu_cluster_computation_disabled[idx];

// create a tdigest that has far fewer values in it than the delta value. this should result
// in every value remaining uncompressed
cudf::test::fixed_width_column_wrapper<T> values{{122, 15, 1, 99, 67, 101, 100, 84, 44, 2},
{1, 0, 1, 0, 1, 0, 1, 0, 1, 0}};
cudf::test::fixed_width_column_wrapper<T> values({122, 15, 1, 99, 67, 101, 100, 84, 44, 2},
{1, 0, 1, 0, 1, 0, 1, 0, 1, 0},
temporary_resources);
int const delta = 1000;
auto result = cudf::type_dispatcher(
static_cast<column_view>(values).type(), tdigest_gen{}, op, values, delta);

cudf::test::fixed_width_column_wrapper<T> raw_mean({1, 44, 67, 100, 122});
cudf::test::fixed_width_column_wrapper<double> weight{1, 1, 1, 1, 1};
auto mean = cudf::cast(raw_mean, data_type{type_id::FLOAT64});
cudf::test::fixed_width_column_wrapper<T> raw_mean({1, 44, 67, 100, 122}, temporary_resources);
cudf::test::fixed_width_column_wrapper<double> weight({1, 1, 1, 1, 1}, temporary_resources);
auto mean =
cudf::cast(raw_mean, data_type{type_id::FLOAT64}, cudf::get_default_stream(), temporary_mr);
double const min = 1;
double const max = 122;
auto expected = make_expected_tdigest_column({{*mean,
weight,
static_cast<double>(static_cast<T>(min)),
static_cast<double>(static_cast<T>(max))}});
static_cast<double>(static_cast<T>(max))}},
temporary_resources);

CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, *expected);
CUDF_TEST_EXPECT_COLUMNS_EQUAL(
*result, *expected, cudf::test::debug_output_level::FIRST_ERROR, mr);
}
}

// shared test for groupby/reduction.
template <typename T, typename Func>
void tdigest_simple_all_nulls_aggregation(Func op)
void tdigest_simple_all_nulls_aggregation(
Func op, cudf::memory_resources mr = cudf::get_current_device_resource_ref())
{
auto const temporary_mr = mr.get_temporary_mr();
auto const temporary_resources = cudf::memory_resources{temporary_mr, temporary_mr};
bool is_cpu_cluster_computation_disabled[2] = {true, false};
for (int idx = 0; idx < 2; idx++) {
cudf::tdigest::detail::is_cpu_cluster_computation_disabled =
is_cpu_cluster_computation_disabled[idx];

// create a tdigest that has far fewer values in it than the delta value. this should result
// in every value remaining uncompressed
cudf::test::fixed_width_column_wrapper<T> values{{122, 15, 1, 99, 67, 101, 100, 84, 44, 2},
{0, 0, 0, 0, 0, 0, 0, 0, 0, 0}};
cudf::test::fixed_width_column_wrapper<T> values({122, 15, 1, 99, 67, 101, 100, 84, 44, 2},
{0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
temporary_resources);
int const delta = 1000;
auto result = cudf::type_dispatcher(
static_cast<column_view>(values).type(), tdigest_gen{}, op, values, delta);

// NOTE: an empty tdigest column still has 1 row.
auto expected = cudf::tdigest::detail::make_empty_tdigests_column(
1, cudf::get_default_stream(), cudf::get_current_device_resource_ref());
1, cudf::get_default_stream(), temporary_mr);

CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, *expected);
CUDF_TEST_EXPECT_COLUMNS_EQUAL(
*result, *expected, cudf::test::debug_output_level::FIRST_ERROR, mr);
}
}

// Note: there is no need to test different types here as the internals of a tdigest are always
// the same regardless of input.
template <typename Func, typename MergeFunc>
void tdigest_merge_simple(Func op, MergeFunc merge_op)
void tdigest_merge_simple(Func op,
MergeFunc merge_op,
cudf::memory_resources mr = cudf::get_current_device_resource_ref())
{
auto const temporary_mr = mr.get_temporary_mr();
auto const temporary_resources = cudf::memory_resources{temporary_mr, temporary_mr};
bool is_cpu_cluster_computation_disabled[2] = {true, false};
for (int idx = 0; idx < 2; idx++) {
cudf::tdigest::detail::is_cpu_cluster_computation_disabled =
is_cpu_cluster_computation_disabled[idx];

auto values = generate_standardized_percentile_distribution(data_type{type_id::FLOAT64});
auto values = generate_standardized_percentile_distribution(
data_type{type_id::FLOAT64}, false, temporary_resources);
CUDF_EXPECTS(values->size() == 750000, "Unexpected distribution size");

auto split_values = cudf::split(*values, {250000, 500000});
Expand All @@ -272,7 +335,7 @@ void tdigest_merge_simple(Func op, MergeFunc merge_op)
int const merge_delta = 1000;

// merge them
auto merge_input = cudf::concatenate(part_views);
auto merge_input = cudf::concatenate(part_views, cudf::get_default_stream(), temporary_mr);
auto result = merge_op(*merge_input, merge_delta);
cudf::tdigest::tdigest_column_view tdv(*result);

Expand All @@ -290,43 +353,46 @@ void tdigest_merge_simple(Func op, MergeFunc merge_op)
{625, 98.20470345147104751504, 405},
{700, 99.96818381983835877236, 56},
{711, 99.99970905482754801596, 1}};
tdigest_sample_compare(tdv, expected);
tdigest_sample_compare(tdv, expected, mr);

// verify min/max
tdigest_minmax_compare<double>(tdv, *values);
tdigest_minmax_compare<double>(tdv, *values, mr);
}
}
}

// shared test for groupby/reduction.
template <typename MergeFunc>
void tdigest_merge_empty(MergeFunc merge_op)
void tdigest_merge_empty(MergeFunc merge_op,
cudf::memory_resources mr = cudf::get_current_device_resource_ref())
{
auto const temporary_mr = mr.get_temporary_mr();
bool is_cpu_cluster_computation_disabled[2] = {true, false};
for (int idx = 0; idx < 2; idx++) {
cudf::tdigest::detail::is_cpu_cluster_computation_disabled =
is_cpu_cluster_computation_disabled[idx];

// 3 empty tdigests all in the same group
auto a = cudf::tdigest::detail::make_empty_tdigests_column(
1, cudf::get_default_stream(), cudf::get_current_device_resource_ref());
1, cudf::get_default_stream(), temporary_mr);
auto b = cudf::tdigest::detail::make_empty_tdigests_column(
1, cudf::get_default_stream(), cudf::get_current_device_resource_ref());
1, cudf::get_default_stream(), temporary_mr);
auto c = cudf::tdigest::detail::make_empty_tdigests_column(
1, cudf::get_default_stream(), cudf::get_current_device_resource_ref());
1, cudf::get_default_stream(), temporary_mr);
std::vector<column_view> cols;
cols.push_back(*a);
cols.push_back(*b);
cols.push_back(*c);
auto values = cudf::concatenate(cols);
auto values = cudf::concatenate(cols, cudf::get_default_stream(), temporary_mr);

auto const delta = 1000;
auto result = merge_op(*values, delta);

auto expected = cudf::tdigest::detail::make_empty_tdigests_column(
1, cudf::get_default_stream(), cudf::get_current_device_resource_ref());
1, cudf::get_default_stream(), temporary_mr);

CUDF_TEST_EXPECT_COLUMNS_EQUAL(*expected, *result);
CUDF_TEST_EXPECT_COLUMNS_EQUAL(
*expected, *result, cudf::test::debug_output_level::FIRST_ERROR, mr);
}
}

Expand Down
49 changes: 48 additions & 1 deletion cpp/tests/reductions/tdigest_tests.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/

Expand All @@ -10,6 +10,8 @@

#include <cudf/reduction.hpp>

#include <rmm/mr/statistics_resource_adaptor.hpp>

template <typename T>
struct ReductionTDigestAllTypes : public cudf::test::BaseFixture {};
TYPED_TEST_SUITE(ReductionTDigestAllTypes, cudf::test::NumericTypes);
Expand Down Expand Up @@ -77,6 +79,51 @@ TEST_F(ReductionTDigestMerge, Simple)
cudf::test::tdigest_merge_simple(reduce_op{}, reduce_merge_op{});
}

TEST_F(ReductionTDigestMerge, TestUtilityMemoryResourceControl)
{
auto upstream = cudf::get_current_device_resource_ref();
auto output_mr = rmm::mr::statistics_resource_adaptor(upstream);
auto temporary_mr = rmm::mr::statistics_resource_adaptor(upstream);
auto resources = cudf::memory_resources{output_mr, temporary_mr};

{
auto distribution = cudf::test::generate_typed_percentile_distribution(
{10.0}, {4}, cudf::data_type{cudf::type_id::FLOAT64}, false, resources);
cudf::test::get_default_stream().synchronize();
EXPECT_GT(output_mr.get_bytes_counter().value, 0);
EXPECT_EQ(temporary_mr.get_bytes_counter().value, 0);
EXPECT_GT(temporary_mr.get_bytes_counter().total, 0);
}
cudf::test::get_default_stream().synchronize();
EXPECT_EQ(output_mr.get_bytes_counter().value, 0);

cudf::test::fixed_width_column_wrapper<double> means{1.0, 2.0};
cudf::test::fixed_width_column_wrapper<double> weights{1.0, 1.0};
auto validation_output_mr = rmm::mr::statistics_resource_adaptor(upstream);
auto validation_temporary_mr = rmm::mr::statistics_resource_adaptor(upstream);
auto validation_resources = cudf::memory_resources{validation_output_mr, validation_temporary_mr};
auto const temporary_bytes_before = temporary_mr.get_bytes_counter().total;

{
auto expected =
cudf::test::make_expected_tdigest_column({{means, weights, 1.0, 2.0}}, resources);
cudf::tdigest::tdigest_column_view tdv(*expected);
cudf::test::tdigest_sample_compare(tdv, {{0, 1.0, 1.0}, {1, 2.0, 1.0}}, validation_resources);
cudf::test::tdigest_minmax_compare<double>(tdv, means, validation_resources);

cudf::test::get_default_stream().synchronize();
EXPECT_GT(output_mr.get_bytes_counter().value, 0);
EXPECT_EQ(temporary_mr.get_bytes_counter().value, 0);
EXPECT_GT(temporary_mr.get_bytes_counter().total, temporary_bytes_before);
EXPECT_EQ(validation_output_mr.get_bytes_counter().value, 0);
EXPECT_EQ(validation_output_mr.get_bytes_counter().total, 0);
EXPECT_EQ(validation_temporary_mr.get_bytes_counter().value, 0);
EXPECT_GT(validation_temporary_mr.get_bytes_counter().total, 0);
}
cudf::test::get_default_stream().synchronize();
EXPECT_EQ(output_mr.get_bytes_counter().value, 0);
}

// tests an issue with the cluster generating code with a small number of centroids that have large
// weights
TEST_F(ReductionTDigestMerge, FewHeavyCentroids)
Expand Down
Loading
Loading