diff --git a/cpp/src/quantiles/tdigest/tdigest.cu b/cpp/src/quantiles/tdigest/tdigest.cu index 283f66bbaf8..d493308b4b8 100644 --- a/cpp/src/quantiles/tdigest/tdigest.cu +++ b/cpp/src/quantiles/tdigest/tdigest.cu @@ -32,6 +32,8 @@ #include #include +#include + using namespace cudf::tdigest; namespace cudf { @@ -59,6 +61,8 @@ struct make_centroid { // kernel for computing percentiles on input tdigest (mean, weight) centroid data. template CUDF_KERNEL void compute_percentiles_kernel(device_span tdigest_offsets, + device_span output_offsets, + bool output_is_dense, column_device_view percentiles, CentroidIter centroids_, double const* min_, @@ -66,95 +70,97 @@ CUDF_KERNEL void compute_percentiles_kernel(device_span tdigest_o 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(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(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{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{lhs, rhs}; + auto const num_tdigests = tdigest_offsets.size() - 1; + auto const num_pairs = static_cast(num_tdigests) * percentiles.size(); + for (auto tid = static_cast(cudf::detail::grid_1d::global_thread_id()); tid < num_pairs; + tid += cudf::detail::grid_1d::grid_stride()) { + auto const tdigest_index = static_cast(tid / percentiles.size()); + auto const pindex = static_cast(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(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(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{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{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); + }(); + } } /** @@ -174,6 +180,9 @@ CUDF_KERNEL void compute_percentiles_kernel(device_span tdigest_o */ std::unique_ptr compute_approx_percentiles(tdigest_column_view const& input, column_view const& percentiles, + device_span output_offsets, + size_type num_output_values, + bool output_is_dense, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { @@ -207,9 +216,6 @@ std::unique_ptr 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 @@ -230,10 +236,17 @@ std::unique_ptr compute_approx_percentiles(tdigest_column_view const& in auto centroids = cudf::detail::make_counting_transform_iterator( 0, make_centroid{tdv.means().begin(), tdv.weights().begin()}); + // 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(static_cast(percentiles.size()) * input.size(), + static_cast(max_blocks) * block_size), + block_size); compute_percentiles_kernel<<>>( {offsets.begin(), static_cast(offsets.size())}, + output_offsets, + output_is_dense, *percentiles_cdv, centroids, tdv.min_begin(), @@ -346,32 +359,7 @@ std::unique_ptr 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(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()); - - 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) + // 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), @@ -389,11 +377,46 @@ std::unique_ptr 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(input.size() - null_count) * percentiles.size(); + CUDF_EXPECTS(compact_child_size <= std::numeric_limits::max(), + "The percentile_approx output exceeds the maximum column size"); + auto const num_output_values = static_cast(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( + [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; + })); + 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()); + 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{offsets->view().begin(), + static_cast(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)); } } // namespace tdigest diff --git a/cpp/tests/quantiles/percentile_approx_test.cpp b/cpp/tests/quantiles/percentile_approx_test.cpp index 9d51b0ed7b2..d0503e5d598 100644 --- a/cpp/tests/quantiles/percentile_approx_test.cpp +++ b/cpp/tests/quantiles/percentile_approx_test.cpp @@ -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{{1, 0, 3}, {true, false, true}}; + auto const keys = cudf::test::fixed_width_column_wrapper{0, 1, 2}; + auto const percentiles = cudf::test::fixed_width_column_wrapper{0.0, 0.5, 1.0}; + + cudf::groupby::groupby gb( + cudf::table_view{{keys}}, cudf::null_policy::EXCLUDE, cudf::sorted::YES); + std::vector requests; + std::vector> aggregations; + aggregations.push_back(cudf::make_tdigest_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 offsets{0, 3, 3, 6}; + cudf::test::fixed_width_column_wrapper child{1, 1, 1, 3, 3, 3}; + std::vector 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;