Skip to content

Fix unsanitized nulls output from percentile_approx with partial empty input #24095

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

Open
davidwendt wants to merge 3 commits into
base: main
Choose a base branch
from
+173 −122
Open

Fix unsanitized nulls output from percentile_approx with partial empty input #24095

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a34ceb2
Fix unsanitized nulls output from percentile_approx with partial empt…
davidwendt Sep 9, 2026
34021c8
Merge branch 'main' into empty-tdigests
davidwendt Sep 10, 2026
2fce386
add overflow check
davidwendt Sep 10, 2026
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
267 changes: 145 additions & 122 deletions cpp/src/quantiles/tdigest/tdigest.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
#include <thrust/reduce.h>
#include <thrust/scan.h>

#include <limits>

using namespace cudf::tdigest;

namespace cudf {
Expand Down Expand Up @@ -59,102 +61,106 @@ struct make_centroid {
// kernel for computing percentiles on input tdigest (mean, weight) centroid data.
template <typename CentroidIter>
CUDF_KERNEL void compute_percentiles_kernel(device_span<int32_t const> tdigest_offsets,
device_span<int32_t const> output_offsets,
bool output_is_dense,

@davidwendt davidwendt Sep 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Recommend using "Hide whitespace" when reviewing this section.

column_device_view percentiles,
CentroidIter centroids_,
double const* min_,
double const* max_,
double const* cumulative_weight_,
double* output)
{
auto const tid = cudf::detail::grid_1d::global_thread_id();

auto const num_tdigests = tdigest_offsets.size() - 1;
auto const tdigest_index = tid / percentiles.size();
if (tdigest_index >= num_tdigests) { return; }
auto const pindex = tid % percentiles.size();

// size of the digest we're querying
auto const tdigest_size = tdigest_offsets[tdigest_index + 1] - tdigest_offsets[tdigest_index];
// no work to do. values will be set to null
if (tdigest_size == 0 || !percentiles.is_valid(pindex)) { return; }

output[tid] = [&]() {
double const percentage = percentiles.element<double>(pindex);
double const* cumulative_weight = cumulative_weight_ + tdigest_offsets[tdigest_index];

// centroids for this particular tdigest
CentroidIter centroids = centroids_ + tdigest_offsets[tdigest_index];

// min and max for the digest
double const* min_val = min_ + tdigest_index;
double const* max_val = max_ + tdigest_index;

double const total_weight = cumulative_weight[tdigest_size - 1];

// The following Arrow code serves as a basis for this computation
// https://github.com/apache/arrow/blob/master/cpp/src/arrow/util/tdigest.cc#L280
double const weighted_q = percentage * total_weight;
if (weighted_q <= 1) {
return *min_val;
} else if (weighted_q > total_weight - 1) {
return *max_val;
}

// determine what centroid this weighted quantile falls within.
size_type const centroid_index = static_cast<size_type>(cuda::std::distance(
cumulative_weight,
thrust::lower_bound(
thrust::seq, cumulative_weight, cumulative_weight + tdigest_size, weighted_q)));
centroid c = centroids[centroid_index];

// diff == how far from the "center" of the centroid we are,
// in unit weights.
// visually:
//
// centroid of weight 7
// C <-- center of the centroid
// |-------|
// | | |
// X Y Z
// X has a diff of -2 (2 units to the left of the center of the centroid)
// Y has a diff of 0 (directly in the middle of the centroid)
// Z has a diff of 3 (3 units to the right of the center of the centroid)
double const diff = weighted_q + c.weight / 2 - cumulative_weight[centroid_index];

// if we're completely within a centroid of weight 1, just return that.
if (c.weight == 1 && cuda::std::abs(diff) <= 0.5) { return c.mean; }

// otherwise, interpolate between two centroids.

// get the two centroids we want to interpolate between
auto const look_left = diff < 0;
auto const [lhs, rhs] = [&]() {
if (look_left) {
// if we're at the first centroid, "left" of us is the min value
auto const first_centroid = centroid_index == 0;
auto const lhs = first_centroid ? centroid{*min_val, 0} : centroids[centroid_index - 1];
auto const rhs = c;
return cuda::std::pair<centroid, centroid>{lhs, rhs};
} else {
// if we're at the last centroid, "right" of us is the max value
auto const last_centroid = (centroid_index == tdigest_size - 1);
auto const lhs = c;
auto const rhs = last_centroid ? centroid{*max_val, 0} : centroids[centroid_index + 1];
return cuda::std::pair<centroid, centroid>{lhs, rhs};
auto const num_tdigests = tdigest_offsets.size() - 1;
auto const num_pairs = static_cast<int64_t>(num_tdigests) * percentiles.size();
for (auto tid = static_cast<int64_t>(cudf::detail::grid_1d::global_thread_id()); tid < num_pairs;
tid += cudf::detail::grid_1d::grid_stride()) {
auto const tdigest_index = static_cast<size_type>(tid / percentiles.size());
auto const pindex = static_cast<size_type>(tid % percentiles.size());

// size of the digest we're querying
auto const tdigest_size = tdigest_offsets[tdigest_index + 1] - tdigest_offsets[tdigest_index];
// no work to do. values will be set to null
if (tdigest_size == 0 || !percentiles.is_valid(pindex)) { continue; }

auto const output_index = output_is_dense ? tid : output_offsets[tdigest_index] + pindex;
output[output_index] = [&]() {
double const percentage = percentiles.element<double>(pindex);
double const* cumulative_weight = cumulative_weight_ + tdigest_offsets[tdigest_index];

// centroids for this particular tdigest
CentroidIter centroids = centroids_ + tdigest_offsets[tdigest_index];

// min and max for the digest
double const* min_val = min_ + tdigest_index;
double const* max_val = max_ + tdigest_index;

double const total_weight = cumulative_weight[tdigest_size - 1];

// The following Arrow code serves as a basis for this computation
// https://github.com/apache/arrow/blob/master/cpp/src/arrow/util/tdigest.cc#L280
double const weighted_q = percentage * total_weight;
if (weighted_q <= 1) {
return *min_val;
} else if (weighted_q > total_weight - 1) {
return *max_val;
}
}();

// compute interpolation value t

// total interpolation range. the total range of "space" between the lhs and rhs centroids.
auto const tip = lhs.weight / 2 + rhs.weight / 2;
// if we're looking left, diff is negative, so shift it so that we are interpolating
// from lhs -> rhs.
auto const t = (look_left) ? (diff + tip) / tip : diff / tip;

// interpolate
return lerp(lhs.mean, rhs.mean, t);
}();
// determine what centroid this weighted quantile falls within.
size_type const centroid_index = static_cast<size_type>(cuda::std::distance(
cumulative_weight,
thrust::lower_bound(
thrust::seq, cumulative_weight, cumulative_weight + tdigest_size, weighted_q)));
centroid c = centroids[centroid_index];

// diff == how far from the "center" of the centroid we are,
// in unit weights.
// visually:
//
// centroid of weight 7
// C <-- center of the centroid
// |-------|
// | | |
// X Y Z
// X has a diff of -2 (2 units to the left of the center of the centroid)
// Y has a diff of 0 (directly in the middle of the centroid)
// Z has a diff of 3 (3 units to the right of the center of the centroid)
double const diff = weighted_q + c.weight / 2 - cumulative_weight[centroid_index];

// if we're completely within a centroid of weight 1, just return that.
if (c.weight == 1 && cuda::std::abs(diff) <= 0.5) { return c.mean; }

// otherwise, interpolate between two centroids.

// get the two centroids we want to interpolate between
auto const look_left = diff < 0;
auto const [lhs, rhs] = [&]() {
if (look_left) {
// if we're at the first centroid, "left" of us is the min value
auto const first_centroid = centroid_index == 0;
auto const lhs = first_centroid ? centroid{*min_val, 0} : centroids[centroid_index - 1];
auto const rhs = c;
return cuda::std::pair<centroid, centroid>{lhs, rhs};
} else {
// if we're at the last centroid, "right" of us is the max value
auto const last_centroid = (centroid_index == tdigest_size - 1);
auto const lhs = c;
auto const rhs = last_centroid ? centroid{*max_val, 0} : centroids[centroid_index + 1];
return cuda::std::pair<centroid, centroid>{lhs, rhs};
}
}();

// compute interpolation value t

// total interpolation range. the total range of "space" between the lhs and rhs centroids.
auto const tip = lhs.weight / 2 + rhs.weight / 2;
// if we're looking left, diff is negative, so shift it so that we are interpolating
// from lhs -> rhs.
auto const t = (look_left) ? (diff + tip) / tip : diff / tip;

// interpolate
return lerp(lhs.mean, rhs.mean, t);
}();
}
}

/**
Expand All @@ -174,6 +180,9 @@ CUDF_KERNEL void compute_percentiles_kernel(device_span<int32_t const> tdigest_o
*/
std::unique_ptr<column> compute_approx_percentiles(tdigest_column_view const& input,
column_view const& percentiles,
device_span<int32_t const> output_offsets,
size_type num_output_values,
bool output_is_dense,
cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
Expand Down Expand Up @@ -207,9 +216,6 @@ std::unique_ptr<column> compute_approx_percentiles(tdigest_column_view const& in

auto percentiles_cdv = column_device_view::create(percentiles, stream);

// leaf is a column of size input.size() * percentiles.size()
auto const num_output_values = input.size() * percentiles.size();

// null percentiles become null results.
auto [null_mask, null_count] = [&]() {
return percentiles.null_count() != 0
Expand All @@ -230,10 +236,17 @@ std::unique_ptr<column> compute_approx_percentiles(tdigest_column_view const& in
auto centroids = cudf::detail::make_counting_transform_iterator(
0, make_centroid{tdv.means().begin<double>(), tdv.weights().begin<double>()});

// Bound the grid because the dense input pair count can exceed the compact output size.
constexpr size_type block_size = 256;
cudf::detail::grid_1d const grid(percentiles.size() * input.size(), block_size);
constexpr size_type max_blocks = 65535;
cudf::detail::grid_1d const grid(
std::min<int64_t>(static_cast<int64_t>(percentiles.size()) * input.size(),
static_cast<int64_t>(max_blocks) * block_size),
block_size);
compute_percentiles_kernel<<<grid.num_blocks, block_size, 0, stream.get()>>>(
{offsets.begin<size_type>(), static_cast<size_t>(offsets.size())},
output_offsets,
output_is_dense,
*percentiles_cdv,
centroids,
tdv.min_begin(),
Expand Down Expand Up @@ -346,32 +359,7 @@ std::unique_ptr<column> percentile_approx(tdigest_column_view const& input,
CUDF_EXPECTS(percentiles.type().id() == type_id::FLOAT64,
"percentile_approx expects float64 percentile inputs");

// output is a list column with each row containing percentiles.size() percentile values
auto offsets = cudf::make_fixed_width_column(
data_type{type_id::INT32}, input.size() + 1, mask_state::UNALLOCATED, stream, mr);
auto const all_empty_rows = cudf::detail::count_if(
detail::size_begin(input),
detail::size_begin(input) + input.size(),
[] __device__(auto const x) { return x == 0; },
stream) == static_cast<std::size_t>(input.size());
auto row_size_iter = cuda::make_constant_iterator(all_empty_rows ? 0 : percentiles.size());
thrust::exclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
row_size_iter,
row_size_iter + input.size() + 1,
offsets->mutable_view().begin<int32_t>());

if (percentiles.size() == 0 || all_empty_rows) {
return cudf::make_lists_column(
input.size(),
std::move(offsets),
cudf::make_empty_column(type_id::FLOAT64),
input.size(),
cudf::detail::create_null_mask(
input.size(), mask_state::ALL_NULL, cuda::stream_ref(stream), mr));
}

// if any of the input digests are empty, nullify the corresponding output rows (values will be
// uninitialized)

@davidwendt davidwendt Sep 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The above code is mostly just moved below the bitmask, null_count statement.

// Empty input digests become null output rows with zero-length list ranges.
auto [bitmask, null_count] = [stream, mr, &tdv]() {
auto tdigest_is_empty = cuda::transform_iterator(
detail::size_begin(tdv),
Expand All @@ -389,11 +377,46 @@ std::unique_ptr<column> percentile_approx(tdigest_column_view const& input,
tdigest_is_empty, tdigest_is_empty + tdv.size(), cuda::std::logical_not{}, stream, mr);
}();

return cudf::make_lists_column(input.size(),
std::move(offsets),
detail::compute_approx_percentiles(input, percentiles, stream, mr),
null_count,
std::move(bitmask));
auto const compact_child_size =
static_cast<int64_t>(input.size() - null_count) * percentiles.size();
CUDF_EXPECTS(compact_child_size <= std::numeric_limits<size_type>::max(),
"The percentile_approx output exceeds the maximum column size");
auto const num_output_values = static_cast<size_type>(compact_child_size);

auto offsets = cudf::make_fixed_width_column(
data_type{type_id::INT32}, input.size() + 1, mask_state::UNALLOCATED, stream, mr);
auto const tdigest_sizes = detail::size_begin(tdv);
auto row_size_iter = cudf::detail::make_counting_transform_iterator(
0,
cuda::proclaim_return_type<size_type>(
[tdigest_sizes, num_rows = input.size(), num_percentiles = percentiles.size()] __device__(
size_type row) -> size_type {
if (row == num_rows) { return 0; }
return tdigest_sizes[row] == 0 ? 0 : num_percentiles;
}));

@davidwendt davidwendt Sep 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fancy iterator helps ensure we do not read past the end of the input memory.

thrust::exclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
row_size_iter,
row_size_iter + input.size() + 1,
offsets->mutable_view().begin<int32_t>());
if (percentiles.size() == 0 || null_count == input.size()) {
return cudf::make_lists_column(
input.size(),
std::move(offsets),
cudf::make_empty_column(type_id::FLOAT64),
input.size(),
cudf::detail::create_null_mask(
input.size(), mask_state::ALL_NULL, cuda::stream_ref(stream), mr));
}

auto const output_offsets = device_span<int32_t const>{offsets->view().begin<int32_t>(),
static_cast<size_t>(offsets->size())};
return cudf::make_lists_column(
input.size(),
std::move(offsets),
detail::compute_approx_percentiles(
input, percentiles, output_offsets, num_output_values, null_count == 0, stream, mr),
null_count,
std::move(bitmask));
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}

} // namespace tdigest
Expand Down
28 changes: 28 additions & 0 deletions cpp/tests/quantiles/percentile_approx_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -466,6 +466,34 @@ TEST_F(PercentileApproxTest, EmptyInput)
CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, *expected);
}

TEST_F(PercentileApproxTest, MixedEmptyInput)
{
auto const values =
cudf::test::fixed_width_column_wrapper<double>{{1, 0, 3}, {true, false, true}};
auto const keys = cudf::test::fixed_width_column_wrapper<int32_t>{0, 1, 2};
auto const percentiles = cudf::test::fixed_width_column_wrapper<double>{0.0, 0.5, 1.0};

cudf::groupby::groupby gb(
cudf::table_view{{keys}}, cudf::null_policy::EXCLUDE, cudf::sorted::YES);
std::vector<cudf::groupby::aggregation_request> requests;
std::vector<std::unique_ptr<cudf::groupby_aggregation>> aggregations;
aggregations.push_back(cudf::make_tdigest_aggregation<cudf::groupby_aggregation>(1000));
requests.push_back({values, std::move(aggregations)});
auto const tdigest_column = gb.aggregate(requests);

cudf::tdigest::tdigest_column_view tdv(*tdigest_column.second[0].results[0]);
auto const result = cudf::percentile_approx(tdv, percentiles);

cudf::test::fixed_width_column_wrapper<cudf::size_type> offsets{0, 3, 3, 6};
cudf::test::fixed_width_column_wrapper<double> child{1, 1, 1, 3, 3, 3};
std::vector<bool> nulls{true, false, true};
auto [null_mask, null_count] = cudf::test::detail::make_null_mask(nulls.begin(), nulls.end());
auto expected = cudf::make_lists_column(
3, offsets.release(), child.release(), null_count, std::move(null_mask));

CUDF_TEST_EXPECT_COLUMNS_EQUAL(*result, *expected);
}

TEST_F(PercentileApproxTest, EmptyPercentiles)
{
auto const delta = 1000;
Expand Down
Loading