diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 92e18a162805..eefb265464fe 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -620,6 +620,7 @@ add_library( src/groupby/sort/group_replace_nulls.cu src/groupby/sort/group_std.cu src/groupby/sort/group_sum.cu + src/groupby/sort/group_sum_overflow.cu src/groupby/sort/group_sum_scan.cu src/groupby/sort/group_topk.cu src/groupby/sort/host_udf_aggregation.cpp diff --git a/cpp/include/cudf/aggregation.hpp b/cpp/include/cudf/aggregation.hpp index 77c8836aabcf..5cbd81dc1ac8 100644 --- a/cpp/include/cudf/aggregation.hpp +++ b/cpp/include/cudf/aggregation.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -77,48 +77,50 @@ class aggregation { * @brief Possible aggregation operations. */ enum Kind : int32_t { - SUM = 0, ///< sum reduction - SUM_WITH_OVERFLOW, ///< sum reduction with overflow detection - PRODUCT, ///< product reduction - MIN, ///< min reduction - MAX, ///< max reduction - COUNT_VALID, ///< count number of valid elements - COUNT_ALL, ///< count number of elements - ANY, ///< any reduction - ALL, ///< all reduction - SUM_OF_SQUARES, ///< sum of squares reduction - MEAN, ///< arithmetic mean reduction - M2, ///< sum of squares of differences from the mean - VARIANCE, ///< variance - STD, ///< standard deviation - MEDIAN, ///< median reduction - QUANTILE, ///< compute specified quantile(s) - ARGMAX, ///< Index of max element - ARGMIN, ///< Index of min element - NUNIQUE, ///< count number of unique elements - NTH_ELEMENT, ///< get the nth element - ROW_NUMBER, ///< get row-number of current index (relative to rolling window) - EWMA, ///< get exponential weighted moving average at current index - RANK, ///< get rank of current index - COLLECT_LIST, ///< collect values into a list - COLLECT_SET, ///< collect values into a list without duplicate entries - LEAD, ///< window function, accesses row at specified offset following current row - LAG, ///< window function, accesses row at specified offset preceding current row - PTX, ///< PTX based UDF aggregation - CUDA, ///< CUDA based UDF aggregation - HOST_UDF, ///< host based UDF aggregation - MERGE_LISTS, ///< merge multiple lists values into one list - MERGE_SETS, ///< merge multiple lists values into one list then drop duplicate entries - MERGE_M2, ///< merge partial values of M2 aggregation, - COVARIANCE, ///< covariance between two sets of elements - CORRELATION, ///< correlation between two sets of elements - TDIGEST, ///< create a tdigest from a set of input values - MERGE_TDIGEST, ///< create a tdigest by merging multiple tdigests together - HISTOGRAM, ///< compute frequency of each element - MERGE_HISTOGRAM, ///< merge partial values of HISTOGRAM aggregation - BITWISE_AGG, ///< bitwise aggregation on numeric columns - TOP_K, ///< top k elements in a group - INVALID ///< invalid aggregation, used as a placeholder when default-constructed + SUM = 0, ///< sum reduction + SUM_OVERFLOW, ///< sum reduction with overflow detection + /// @deprecated Use SUM_OVERFLOW instead. + SUM_WITH_OVERFLOW = SUM_OVERFLOW, + PRODUCT, ///< product reduction + MIN, ///< min reduction + MAX, ///< max reduction + COUNT_VALID, ///< count number of valid elements + COUNT_ALL, ///< count number of elements + ANY, ///< any reduction + ALL, ///< all reduction + SUM_OF_SQUARES, ///< sum of squares reduction + MEAN, ///< arithmetic mean reduction + M2, ///< sum of squares of differences from the mean + VARIANCE, ///< variance + STD, ///< standard deviation + MEDIAN, ///< median reduction + QUANTILE, ///< compute specified quantile(s) + ARGMAX, ///< Index of max element + ARGMIN, ///< Index of min element + NUNIQUE, ///< count number of unique elements + NTH_ELEMENT, ///< get the nth element + ROW_NUMBER, ///< get row-number of current index (relative to rolling window) + EWMA, ///< get exponential weighted moving average at current index + RANK, ///< get rank of current index + COLLECT_LIST, ///< collect values into a list + COLLECT_SET, ///< collect values into a list without duplicate entries + LEAD, ///< window function, accesses row at specified offset following current row + LAG, ///< window function, accesses row at specified offset preceding current row + PTX, ///< PTX based UDF aggregation + CUDA, ///< CUDA based UDF aggregation + HOST_UDF, ///< host based UDF aggregation + MERGE_LISTS, ///< merge multiple lists values into one list + MERGE_SETS, ///< merge multiple lists values into one list then drop duplicate entries + MERGE_M2, ///< merge partial values of M2 aggregation, + COVARIANCE, ///< covariance between two sets of elements + CORRELATION, ///< correlation between two sets of elements + TDIGEST, ///< create a tdigest from a set of input values + MERGE_TDIGEST, ///< create a tdigest by merging multiple tdigests together + HISTOGRAM, ///< compute frequency of each element + MERGE_HISTOGRAM, ///< merge partial values of HISTOGRAM aggregation + BITWISE_AGG, ///< bitwise aggregation on numeric columns + TOP_K, ///< top k elements in a group + INVALID ///< invalid aggregation, used as a placeholder when default-constructed }; /** @@ -212,8 +214,14 @@ enum class ewm_history : int32_t { INFINITE, FINITE }; template std::unique_ptr make_sum_aggregation(); +/// Factory to create a SUM_OVERFLOW aggregation +/// @return A SUM_OVERFLOW aggregation object +template +std::unique_ptr make_sum_overflow_aggregation(); + /// Factory to create a SUM_WITH_OVERFLOW aggregation /// @return A SUM_WITH_OVERFLOW aggregation object +/// @deprecated Use make_sum_overflow_aggregation() instead. template std::unique_ptr make_sum_with_overflow_aggregation(); diff --git a/cpp/include/cudf/detail/aggregation/aggregation.hpp b/cpp/include/cudf/detail/aggregation/aggregation.hpp index b848f4417b57..1dbb6594ba50 100644 --- a/cpp/include/cudf/detail/aggregation/aggregation.hpp +++ b/cpp/include/cudf/detail/aggregation/aggregation.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -977,11 +977,19 @@ struct target_type_impl - requires((cudf::is_integral_not_bool() && cudf::is_signed()) || - cudf::is_fixed_point()) +concept sum_overflow_supported = + (cudf::is_integral_not_bool() && cudf::is_signed()) || + cudf::is_fixed_point(); + +// SUM_WITH_OVERFLOW outputs a struct {sum: Source, overflow: bool} where the sum matches the input +// type +template struct target_type_impl { using type = struct_view; // SUM_WITH_OVERFLOW outputs a struct with sum and overflow fields }; diff --git a/cpp/include/cudf/reduction/detail/sum_overflow.cuh b/cpp/include/cudf/reduction/detail/sum_overflow.cuh new file mode 100644 index 000000000000..87d4f656c73d --- /dev/null +++ b/cpp/include/cudf/reduction/detail/sum_overflow.cuh @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +#include + +namespace cudf { +namespace reduction::detail { + +/** + * @brief Running accumulator for a sum that detects signed-integer overflow. + * + * `wraps` is the net number of times the running sum has stepped outside [MIN, MAX]. + * A final `wraps == 0` means the true sum fits in `DeviceType`, i.e. no overflow. + */ +template +struct sum_overflow_result { + DeviceType sum; + cudf::size_type wraps; + + CUDF_HOST_DEVICE sum_overflow_result() : sum{0}, wraps{0} {} + CUDF_HOST_DEVICE sum_overflow_result(DeviceType s, cudf::size_type w) : sum{s}, wraps{w} {} +}; + +/// @brief Associative combine: wrap the sums and track the net carry direction. +template +struct overflow_sum_op { + __device__ sum_overflow_result operator()( + sum_overflow_result const& lhs, sum_overflow_result const& rhs) const + { + auto const r = cuda::add_overflow(lhs.sum, rhs.sum); + auto const carry = r.overflow ? (rhs.sum > DeviceType{0} ? 1 : -1) : 0; + return sum_overflow_result{r.value, lhs.wraps + rhs.wraps + carry}; + } +}; + +/// @brief Maps a value to a zero-wrap accumulator. +template +struct to_sum_overflow { + __device__ sum_overflow_result operator()(DeviceType value) const + { + return sum_overflow_result{value, 0}; + } +}; + +/// @brief Maps a row index to an accumulator, treating nulls as a zero contribution. +template +struct null_aware_to_sum_overflow { + cudf::column_device_view dcol; + + CUDF_HOST_DEVICE null_aware_to_sum_overflow(cudf::column_device_view const& d) : dcol{d} {} + + __device__ sum_overflow_result operator()(cudf::size_type idx) const + { + return dcol.is_valid(idx) ? sum_overflow_result{dcol.element(idx), 0} + : sum_overflow_result{DeviceType{0}, 0}; + } +}; + +} // namespace reduction::detail +} // namespace cudf diff --git a/cpp/src/aggregation/aggregation.cpp b/cpp/src/aggregation/aggregation.cpp index ab871abf94ed..e82d602e3bb5 100644 --- a/cpp/src/aggregation/aggregation.cpp +++ b/cpp/src/aggregation/aggregation.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -35,6 +35,21 @@ template CUDF_EXPORT std::unique_ptr make_sum_aggregation(); /// Factory to create a SUM_WITH_OVERFLOW aggregation +template +std::unique_ptr make_sum_overflow_aggregation() +{ + return std::make_unique(); +} +template CUDF_EXPORT std::unique_ptr make_sum_overflow_aggregation(); +template CUDF_EXPORT std::unique_ptr +make_sum_overflow_aggregation(); +template CUDF_EXPORT std::unique_ptr +make_sum_overflow_aggregation(); +template CUDF_EXPORT std::unique_ptr +make_sum_overflow_aggregation(); +template CUDF_EXPORT std::unique_ptr +make_sum_overflow_aggregation(); + template std::unique_ptr make_sum_with_overflow_aggregation() { diff --git a/cpp/src/groupby/sort/aggregate.cpp b/cpp/src/groupby/sort/aggregate.cpp index 2e9521f538f8..ae56aa2327f2 100644 --- a/cpp/src/groupby/sort/aggregate.cpp +++ b/cpp/src/groupby/sort/aggregate.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -140,6 +140,18 @@ void aggregate_result_functor::operator()(aggregation const& a get_grouped_values(), helper.num_groups(stream), helper.group_labels(stream), stream, mr)); } +template <> +void aggregate_result_functor::operator()(aggregation const& agg) +{ + if (cache.has_result(values, agg)) return; + + cache.add_result( + values, + agg, + detail::group_sum_overflow( + get_grouped_values(), helper.num_groups(stream), helper.group_labels(stream), stream, mr)); +} + template <> void aggregate_result_functor::operator()(aggregation const& agg) { @@ -878,11 +890,6 @@ std::pair, std::vector> groupby::sort auto store_functor = detail::aggregate_result_functor(request.values, helper(), cache, stream, mr); for (auto const& agg : request.aggregations) { - // SUM_WITH_OVERFLOW is only supported with hash-based groupby, not sort-based - CUDF_EXPECTS(agg->kind != aggregation::SUM_WITH_OVERFLOW, - "SUM_WITH_OVERFLOW aggregation is only supported with hash-based groupby, not " - "sort-based groupby"); - // TODO (dm): single pass compute all supported reductions cudf::detail::aggregation_dispatcher(agg->kind, store_functor, *agg); } diff --git a/cpp/src/groupby/sort/group_reductions.hpp b/cpp/src/groupby/sort/group_reductions.hpp index db764cb02f25..389a5d7fb535 100644 --- a/cpp/src/groupby/sort/group_reductions.hpp +++ b/cpp/src/groupby/sort/group_reductions.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -44,6 +44,26 @@ std::unique_ptr group_sum(column_view const& values, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); +/** + * @brief Internal API to calculate groupwise sum with overflow detection. + * + * Returns a STRUCT column with two children: the (wrapping) sum and a BOOL8 overflow flag that is + * true when the true sum does not fit in the value type. On overflow the sum value is unspecified; + * the flag is the meaningful output. A group is null only when all of its values are null. + * + * @param values Grouped values to sum + * @param num_groups Number of groups + * @param group_labels ID of group that the corresponding value belongs to + * @param stream CUDA stream used for device memory operations and kernel launches. + * @param mr Device memory resource used to allocate the returned column's device memory + */ +[[nodiscard]] std::unique_ptr group_sum_overflow( + column_view const& values, + size_type num_groups, + cudf::device_span group_labels, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + /** * @brief Internal API to calculate groupwise product * diff --git a/cpp/src/groupby/sort/group_sum_overflow.cu b/cpp/src/groupby/sort/group_sum_overflow.cu new file mode 100644 index 000000000000..a58aa3e2b37f --- /dev/null +++ b/cpp/src/groupby/sort/group_sum_overflow.cu @@ -0,0 +1,127 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "groupby/sort/group_reductions.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +namespace cudf::groupby::detail { +namespace { + +// Splits a reduced {sum, wraps} accumulator into the (sum, overflow-flag) pair of the output +// struct. +template +struct split_accumulator { + __device__ cuda::std::tuple operator()( + cudf::reduction::detail::sum_overflow_result const& acc) const + { + return {acc.sum, acc.wraps != 0}; + } +}; + +struct group_sum_overflow_fn { + template + std::unique_ptr operator()(column_view const& values, + size_type num_groups, + cudf::device_span group_labels, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const + { + using DeviceType = cudf::device_storage_type_t; + + auto const dcol = cudf::column_device_view::create(values, stream); + + auto sum_child = + cudf::make_fixed_width_column(values.type(), num_groups, mask_state::UNALLOCATED, stream, mr); + auto overflow_child = cudf::make_fixed_width_column( + cudf::data_type{type_id::BOOL8}, num_groups, mask_state::UNALLOCATED, stream, mr); + + // Segmented reduction per group, written straight into the two struct children. + auto const values_in = cudf::detail::make_counting_transform_iterator( + 0, cudf::reduction::detail::null_aware_to_sum_overflow{*dcol}); + auto const children_out = cuda::transform_output_iterator{ + cuda::make_zip_iterator(sum_child->mutable_view().begin(), + overflow_child->mutable_view().begin()), + split_accumulator{}}; + + thrust::reduce_by_key(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + group_labels.begin(), + group_labels.end(), + values_in, + cuda::make_discard_iterator(), + children_out, + cuda::std::equal_to{}, + cudf::reduction::detail::overflow_sum_op{}); + + // A group's struct entry is null only when every row in the group is null (mirrors group_sum): + // reduce per-row validity with logical-or, then build the mask from the per-group result. + auto [null_mask, null_count] = [&]() -> std::pair { + if (!values.has_nulls()) { return {rmm::device_buffer{}, size_type{0}}; } + rmm::device_uvector group_valid( + num_groups, stream, cudf::get_current_device_resource_ref()); + thrust::reduce_by_key( + rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + group_labels.begin(), + group_labels.end(), + cudf::detail::make_validity_iterator(*dcol), + cuda::make_discard_iterator(), + group_valid.begin(), + cuda::std::equal_to{}, + cuda::std::logical_or{}); + return cudf::detail::valid_if( + group_valid.begin(), group_valid.end(), cuda::std::identity{}, stream, mr); + }(); + + std::vector> children; + children.push_back(std::move(sum_child)); + children.push_back(std::move(overflow_child)); + return cudf::create_structs_hierarchy( + num_groups, std::move(children), null_count, std::move(null_mask), stream, mr); + } + + template + requires(!cudf::detail::sum_overflow_supported) + std::unique_ptr operator()(Args&&...) const + { + CUDF_FAIL("SUM_OVERFLOW is only supported for signed integral and fixed-point types"); + } +}; + +} // namespace + +std::unique_ptr group_sum_overflow(column_view const& values, + size_type num_groups, + cudf::device_span group_labels, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + return cudf::type_dispatcher( + values.type(), group_sum_overflow_fn{}, values, num_groups, group_labels, stream, mr); +} + +} // namespace cudf::groupby::detail diff --git a/cpp/src/reductions/reductions.cpp b/cpp/src/reductions/reductions.cpp index d58f145d8a7e..4d4f31e17273 100644 --- a/cpp/src/reductions/reductions.cpp +++ b/cpp/src/reductions/reductions.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -95,9 +95,7 @@ struct reduction_function : public base_reductio } }; -template - requires((cudf::is_integral_not_bool() && cudf::is_signed()) || - cudf::is_fixed_point()) +template struct reduction_function : public base_reduction_function { [[nodiscard]] std::unique_ptr reduce(reduction_parameters const& params) const diff --git a/cpp/src/reductions/sum_with_overflow.cu b/cpp/src/reductions/sum_with_overflow.cu index 34fff81317c1..31ad3e0a28ad 100644 --- a/cpp/src/reductions/sum_with_overflow.cu +++ b/cpp/src/reductions/sum_with_overflow.cu @@ -1,15 +1,17 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -27,49 +29,6 @@ namespace cudf::reduction::detail { namespace { -// `wraps` is the net number of times the running sum has stepped outside [MIN, MAX]. -// A final `wraps == 0` means the true sum fits in DeviceType, i.e. no overflow. -template -struct sum_overflow_result { - DeviceType sum; - cudf::size_type wraps; - - CUDF_HOST_DEVICE sum_overflow_result() : sum{0}, wraps{0} {} - CUDF_HOST_DEVICE sum_overflow_result(DeviceType s, cudf::size_type w) : sum{s}, wraps{w} {} -}; - -template -struct overflow_sum_op { - __device__ sum_overflow_result operator()( - sum_overflow_result const& lhs, sum_overflow_result const& rhs) const - { - auto const r = cuda::add_overflow(lhs.sum, rhs.sum); - auto const carry = r.overflow ? (rhs.sum > DeviceType{0} ? 1 : -1) : 0; - return sum_overflow_result{r.value, lhs.wraps + rhs.wraps + carry}; - } -}; - -template -struct to_sum_overflow { - __device__ sum_overflow_result operator()(DeviceType value) const - { - return sum_overflow_result{value, 0}; - } -}; - -template -struct null_aware_to_sum_overflow { - cudf::column_device_view dcol; - - CUDF_HOST_DEVICE null_aware_to_sum_overflow(cudf::column_device_view const& d) : dcol{d} {} - - __device__ sum_overflow_result operator()(cudf::size_type idx) const - { - return dcol.is_valid(idx) ? sum_overflow_result{dcol.element(idx), 0} - : sum_overflow_result{DeviceType{0}, 0}; - } -}; - template std::unique_ptr make_sum_overflow_struct_scalar( device_storage_type_t sum_value, @@ -156,9 +115,7 @@ std::unique_ptr sum_with_overflow_impl( } struct sum_with_overflow_dispatcher { - template - requires((cudf::is_integral_not_bool() && cudf::is_signed()) || - cudf::is_fixed_point()) + template std::unique_ptr operator()(column_view const& col, std::optional> init, rmm::cuda_stream_view stream, @@ -168,8 +125,7 @@ struct sum_with_overflow_dispatcher { } template - requires(!((cudf::is_integral_not_bool() && cudf::is_signed()) || - cudf::is_fixed_point())) + requires(!cudf::detail::sum_overflow_supported) std::unique_ptr operator()(column_view const&, std::optional>, rmm::cuda_stream_view, diff --git a/cpp/tests/groupby/sum_with_overflow_tests.cpp b/cpp/tests/groupby/sum_with_overflow_tests.cpp index 74d89b0945aa..0183b8864bfd 100644 --- a/cpp/tests/groupby/sum_with_overflow_tests.cpp +++ b/cpp/tests/groupby/sum_with_overflow_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -58,12 +58,9 @@ TYPED_TEST(groupby_sum_with_overflow_test, basic) auto agg = cudf::make_sum_with_overflow_aggregation(); test_single_agg(keys, vals, expect_keys, *expect_vals, std::move(agg)); - // SUM_WITH_OVERFLOW should throw with sort-based groupby auto agg_sort = cudf::make_sum_with_overflow_aggregation(); - EXPECT_THROW( - test_single_agg( - keys, vals, expect_keys, *expect_vals, std::move(agg_sort), force_use_sort_impl::YES), - cudf::logic_error); + test_single_agg( + keys, vals, expect_keys, *expect_vals, std::move(agg_sort), force_use_sort_impl::YES); } else { // For integer types cudf::test::fixed_width_column_wrapper vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; @@ -80,15 +77,62 @@ TYPED_TEST(groupby_sum_with_overflow_test, basic) auto agg = cudf::make_sum_with_overflow_aggregation(); test_single_agg(keys, vals, expect_keys, *expect_vals, std::move(agg)); - // SUM_WITH_OVERFLOW should throw with sort-based groupby auto agg_sort = cudf::make_sum_with_overflow_aggregation(); - EXPECT_THROW( - test_single_agg( - keys, vals, expect_keys, *expect_vals, std::move(agg_sort), force_use_sort_impl::YES), - cudf::logic_error); + test_single_agg( + keys, vals, expect_keys, *expect_vals, std::move(agg_sort), force_use_sort_impl::YES); } +} + +TYPED_TEST(groupby_sum_with_overflow_test, sort_path_with_tdigest) +{ + using K = int32_t; + using V = TypeParam; + + cudf::test::fixed_width_column_wrapper keys{1, 2, 3, 1, 2, 2, 1, 3, 3, 2}; + cudf::test::fixed_width_column_wrapper expect_keys{1, 2, 3}; + + // Co-request TDIGEST (a sort-only aggregation) so the whole groupby takes the sort-based path, + // then verify the SUM_WITH_OVERFLOW struct matches the hash result and TDIGEST also runs. + auto run_and_check = [&](cudf::column_view const& vals, cudf::column_view const& expect_vals) { + std::vector requests; + requests.emplace_back(); + requests[0].values = vals; + requests[0].aggregations.push_back( + cudf::make_sum_with_overflow_aggregation()); + requests[0].aggregations.push_back( + cudf::make_tdigest_aggregation(1000)); - // Note: SUM_WITH_OVERFLOW only works with hash groupby, not sort groupby + auto result = cudf::groupby::groupby(cudf::table_view{{keys}}).aggregate(requests); + + // Sort-based groupby returns keys in sorted order, aligning with expect_keys/expect_vals. + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result.first->get_column(0).view(), expect_keys); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result.second[0].results[0]->view(), expect_vals); + // TDIGEST produces one tdigest per group. + EXPECT_EQ(result.second[0].results[1]->size(), 3); + }; + + if constexpr (cudf::is_fixed_point()) { + using RepType = cudf::device_storage_type_t; + auto const scale = scale_type{0}; + auto vals = + cudf::test::fixed_point_column_wrapper{{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, scale}; + auto sum_col = cudf::test::fixed_point_column_wrapper{{9, 19, 17}, scale}; + auto overflow_col = cudf::test::fixed_width_column_wrapper{false, false, false}; + std::vector> children; + children.push_back(sum_col.release()); + children.push_back(overflow_col.release()); + auto expect_vals = cudf::create_structs_hierarchy(3, std::move(children), 0, {}); + run_and_check(vals, *expect_vals); + } else { + cudf::test::fixed_width_column_wrapper vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; + auto sum_col = cudf::test::fixed_width_column_wrapper{9, 19, 17}; + auto overflow_col = cudf::test::fixed_width_column_wrapper{false, false, false}; + std::vector> children; + children.push_back(sum_col.release()); + children.push_back(overflow_col.release()); + auto expect_vals = cudf::create_structs_hierarchy(3, std::move(children), 0, {}); + run_and_check(vals, *expect_vals); + } } TYPED_TEST(groupby_sum_with_overflow_test, empty_cols) @@ -112,8 +156,6 @@ TYPED_TEST(groupby_sum_with_overflow_test, empty_cols) auto agg = cudf::make_sum_with_overflow_aggregation(); test_single_agg(keys, vals, expect_keys, *expect_vals, std::move(agg)); - - // Note: SUM_WITH_OVERFLOW only works with hash groupby, not sort groupby } TYPED_TEST(groupby_sum_with_overflow_test, zero_valid_keys) @@ -137,8 +179,6 @@ TYPED_TEST(groupby_sum_with_overflow_test, zero_valid_keys) auto agg = cudf::make_sum_with_overflow_aggregation(); test_single_agg(keys, vals, expect_keys, *expect_vals, std::move(agg)); - - // Note: SUM_WITH_OVERFLOW only works with hash groupby, not sort groupby } TYPED_TEST(groupby_sum_with_overflow_test, zero_valid_values) @@ -167,7 +207,10 @@ TYPED_TEST(groupby_sum_with_overflow_test, zero_valid_values) auto agg = cudf::make_sum_with_overflow_aggregation(); test_single_agg(keys, vals, expect_keys, *expect_vals, std::move(agg)); - // Note: SUM_WITH_OVERFLOW only works with hash groupby, not sort groupby + // Exercise the sort-based path for an all-null group. + auto agg_sort = cudf::make_sum_with_overflow_aggregation(); + test_single_agg( + keys, vals, expect_keys, *expect_vals, std::move(agg_sort), force_use_sort_impl::YES); } TYPED_TEST(groupby_sum_with_overflow_test, null_keys_and_values) @@ -202,7 +245,10 @@ TYPED_TEST(groupby_sum_with_overflow_test, null_keys_and_values) auto agg = cudf::make_sum_with_overflow_aggregation(); test_single_agg(keys, vals, expect_keys, *expect_vals, std::move(agg)); - // Note: SUM_WITH_OVERFLOW only works with hash groupby, not sort groupby + // Exercise the sort-based path with null keys and null values. + auto agg_sort = cudf::make_sum_with_overflow_aggregation(); + test_single_agg( + keys, vals, expect_keys, *expect_vals, std::move(agg_sort), force_use_sort_impl::YES); } TYPED_TEST(groupby_sum_with_overflow_test, overflow_detection) @@ -230,6 +276,29 @@ TYPED_TEST(groupby_sum_with_overflow_test, overflow_detection) CUDF_TEST_EXPECT_COLUMNS_EQUAL(sorted->view().column(1), expect_overflow); }; + // Same check, but a co-requested sort-only aggregation (TDIGEST) forces the sort path. + auto check_overflow_flags_sort = [](cudf::column_view const& keys, + cudf::column_view const& vals, + cudf::column_view const& expect_keys, + cudf::column_view const& expect_overflow) { + std::vector requests; + requests.emplace_back(); + requests[0].values = vals; + requests[0].aggregations.push_back( + cudf::make_sum_with_overflow_aggregation()); + requests[0].aggregations.push_back( + cudf::make_tdigest_aggregation(1000)); + + auto result = cudf::groupby::groupby(cudf::table_view{{keys}}).aggregate(requests); + auto const overflow_child = + cudf::structs_column_view{result.second[0].results[0]->view()}.get_sliced_child(1); + + auto sorted = cudf::sort_by_key( + cudf::table_view{{result.first->get_column(0).view(), overflow_child}}, result.first->view()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(sorted->view().column(0), expect_keys); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(sorted->view().column(1), expect_overflow); + }; + cudf::test::fixed_width_column_wrapper keys{1, 2, 3, 4, 1, 2, 2, 1, 3, 3, 2, 4, 4}; cudf::test::fixed_width_column_wrapper expect_keys{1, 2, 3, 4}; cudf::test::fixed_width_column_wrapper expect_overflow{true, false, true, true}; @@ -267,16 +336,8 @@ TYPED_TEST(groupby_sum_with_overflow_test, overflow_detection) check_overflow_flags(keys, vals, expect_keys, expect_overflow); - // Adding nth_element forces sort-based groupby, which must throw for SUM_WITH_OVERFLOW. - std::vector sort_requests; - sort_requests.emplace_back(); - sort_requests[0].values = vals; - sort_requests[0].aggregations.push_back( - cudf::make_sum_with_overflow_aggregation()); - sort_requests[0].aggregations.push_back( - cudf::make_nth_element_aggregation(0)); - EXPECT_THROW(cudf::groupby::groupby(cudf::table_view{{keys}}).aggregate(sort_requests), - cudf::logic_error); + // Adding TDIGEST forces sort-based groupby; the overflow flags must match the hash path. + check_overflow_flags_sort(keys, vals, expect_keys, expect_overflow); } else { using DeviceType = cudf::device_storage_type_t; @@ -306,6 +367,7 @@ TYPED_TEST(groupby_sum_with_overflow_test, overflow_detection) static_cast(large_negative)}; check_overflow_flags(keys, vals, expect_keys, expect_overflow); + check_overflow_flags_sort(keys, vals, expect_keys, expect_overflow); } } diff --git a/java/src/test/java/ai/rapids/cudf/TableTest.java b/java/src/test/java/ai/rapids/cudf/TableTest.java index be4ea13e0f14..01ec01f35157 100644 --- a/java/src/test/java/ai/rapids/cudf/TableTest.java +++ b/java/src/test/java/ai/rapids/cudf/TableTest.java @@ -1,6 +1,6 @@ /* * - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * */ @@ -7848,9 +7848,31 @@ void testGroupByM2() { } } + private static void assertSumWithOverflowResult(Table results, + int[] expectedKeys, + long[] expectedSums, + boolean[] expectedOvf) { + ColumnVector structCol = results.getColumn(1); + assertEquals(DType.STRUCT, structCol.getType()); + try (ColumnView ovfChild = structCol.getChildColumnView(1); + ColumnVector ovfCol = ovfChild.copyToColumnVector(); + ColumnVector expectedKeyCol = ColumnVector.fromInts(expectedKeys); + ColumnVector expectedOvfCol = ColumnVector.fromBooleans(expectedOvf)) { + assertColumnsAreEqual(expectedKeyCol, results.getColumn(0)); + assertColumnsAreEqual(expectedOvfCol, ovfCol); + if (expectedSums != null) { + try (ColumnView sumChild = structCol.getChildColumnView(0); + ColumnVector sumCol = sumChild.copyToColumnVector(); + ColumnVector expectedSumCol = ColumnVector.fromLongs(expectedSums)) { + assertEquals(DType.INT64, sumCol.getType()); + assertColumnsAreEqual(expectedSumCol, sumCol); + } + } + } + } + @Test void testGroupByHashSumWithOverflow() { - // int64 keys 1, 2, 3 with values that fit comfortably in int64. try (Table input = new Table.TestBuilder() .column(1, 2, 3, 1, 2, 2, 1, 3, 3, 2) .column(10L, 20L, 30L, 11L, 21L, 22L, 12L, 31L, 32L, 23L) @@ -7860,28 +7882,13 @@ void testGroupByHashSumWithOverflow() { Table sorted = results.orderBy(OrderByArg.asc(0))) { assertEquals(2, sorted.getNumberOfColumns()); assertEquals(3, sorted.getRowCount()); - - ColumnVector keyCol = sorted.getColumn(0); - ColumnVector structCol = sorted.getColumn(1); - assertEquals(DType.STRUCT, structCol.getType()); - - try (ColumnView sumChild = structCol.getChildColumnView(0); - ColumnView ovfChild = structCol.getChildColumnView(1); - ColumnVector sumCol = sumChild.copyToColumnVector(); - ColumnVector ovfCol = ovfChild.copyToColumnVector(); - ColumnVector expectedKeys = ColumnVector.fromInts(1, 2, 3); - ColumnVector expectedSum = ColumnVector.fromLongs(33L, 86L, 93L); - ColumnVector expectedOvf = ColumnVector.fromBooleans(false, false, false)) { - assertColumnsAreEqual(expectedKeys, keyCol); - assertColumnsAreEqual(expectedSum, sumCol); - assertColumnsAreEqual(expectedOvf, ovfCol); - } + assertSumWithOverflowResult(sorted, + new int[]{1, 2, 3}, new long[]{33L, 86L, 93L}, new boolean[]{false, false, false}); } } @Test void testGroupByHashSumWithOverflowDetectsOverflow() { - // Group 1 overflows (max + max), group 2 stays in range. try (Table input = new Table.TestBuilder() .column(1, 1, 2, 2) .column(Long.MAX_VALUE, Long.MAX_VALUE, 3L, 4L) @@ -7889,12 +7896,7 @@ void testGroupByHashSumWithOverflowDetectsOverflow() { Table results = input.groupBy(0).aggregate( GroupByAggregation.sumWithOverflow().onColumn(1)); Table sorted = results.orderBy(OrderByArg.asc(0))) { - ColumnVector structCol = sorted.getColumn(1); - try (ColumnView ovfChild = structCol.getChildColumnView(1); - ColumnVector ovfCol = ovfChild.copyToColumnVector(); - ColumnVector expectedOvf = ColumnVector.fromBooleans(true, false)) { - assertColumnsAreEqual(expectedOvf, ovfCol); - } + assertSumWithOverflowResult(sorted, new int[]{1, 2}, null, new boolean[]{true, false}); } } @@ -7924,17 +7926,31 @@ void testGroupByHashSumWithOverflowInt32() { } @Test - void testGroupBySortSumWithOverflowThrows() { - // Sort-based groupby (keysSorted=true forces the sort impl in cudf). - // SUM_WITH_OVERFLOW is hash-only, so cudf should throw. - GroupByOptions sortOpts = GroupByOptions.builder().withKeysSorted(true).build(); + void testGroupBySortSumWithOverflow() { try (Table input = new Table.TestBuilder() .column(1, 1, 2, 2) .column(1L, 2L, 3L, 4L) - .build()) { - assertThrows(CudfException.class, () -> - input.groupBy(sortOpts, 0).aggregate( - GroupByAggregation.sumWithOverflow().onColumn(1)).close()); + .build(); + // median() is sort-only, so it forces the sort-based groupby path + Table results = input.groupBy(0).aggregate( + GroupByAggregation.sumWithOverflow().onColumn(1), + GroupByAggregation.median().onColumn(1))) { + assertSumWithOverflowResult(results, + new int[]{1, 2}, new long[]{3L, 7L}, new boolean[]{false, false}); + } + } + + @Test + void testGroupBySortSumWithOverflowDetectsOverflow() { + try (Table input = new Table.TestBuilder() + .column(1, 1, 2, 2) + .column(Long.MAX_VALUE, Long.MAX_VALUE, 3L, 4L) + .build(); + // median() is sort-only, so it forces the sort-based groupby path + Table results = input.groupBy(0).aggregate( + GroupByAggregation.sumWithOverflow().onColumn(1), + GroupByAggregation.median().onColumn(1))) { + assertSumWithOverflowResult(results, new int[]{1, 2}, null, new boolean[]{true, false}); } }