From ec9ce8236cd91e39d938f6f74223c0fd360873a8 Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Fri, 17 Jul 2026 20:00:05 +0000 Subject: [PATCH 1/4] Reduce distinct helper build time --- cpp/CMakeLists.txt | 1 + cpp/src/stream_compaction/distinct.cu | 20 +-- cpp/src/stream_compaction/distinct_helpers.cu | 126 +++++------------- .../stream_compaction/distinct_helpers.cuh | 68 ++++++++++ .../stream_compaction/distinct_helpers.hpp | 54 ++++++-- .../distinct_helpers_nested.cu | 19 +++ 6 files changed, 166 insertions(+), 122 deletions(-) create mode 100644 cpp/src/stream_compaction/distinct_helpers.cuh create mode 100644 cpp/src/stream_compaction/distinct_helpers_nested.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 0899c3cec002..c3997b835ad4 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1008,6 +1008,7 @@ add_library( src/stream_compaction/apply_boolean_mask.cu src/stream_compaction/distinct.cu src/stream_compaction/distinct_helpers.cu + src/stream_compaction/distinct_helpers_nested.cu src/stream_compaction/drop_nans.cu src/stream_compaction/drop_nulls.cu src/stream_compaction/filter/filter.cu diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu index e371ea016b22..ea9838de2058 100644 --- a/cpp/src/stream_compaction/distinct.cu +++ b/cpp/src/stream_compaction/distinct.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -22,6 +22,8 @@ #include #include +#include +#include #include #include @@ -48,19 +50,9 @@ rmm::device_uvector dispatch_row_equal( cudf::detail::row::equality::self_comparator row_equal, Func&& func) { - if (compare_nans == nan_equality::ALL_EQUAL) { - auto const d_equal = row_equal.equal_to( - nullate::DYNAMIC{has_nulls}, - compare_nulls, - cudf::detail::row::equality::nan_equal_physical_equality_comparator{}); - return func(d_equal); - } else { - auto const d_equal = - row_equal.equal_to(nullate::DYNAMIC{has_nulls}, - compare_nulls, - cudf::detail::row::equality::physical_equality_comparator{}); - return func(d_equal); - } + auto const d_equal = row_equal.equal_to( + nullate::DYNAMIC{has_nulls}, compare_nulls, distinct_physical_equality{compare_nans}); + return func(d_equal); } } // namespace diff --git a/cpp/src/stream_compaction/distinct_helpers.cu b/cpp/src/stream_compaction/distinct_helpers.cu index 7474c39997c7..3161a7b4c74c 100644 --- a/cpp/src/stream_compaction/distinct_helpers.cu +++ b/cpp/src/stream_compaction/distinct_helpers.cu @@ -1,133 +1,69 @@ /* - * 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 */ -#include "distinct_helpers.hpp" +#include "distinct_helpers.cuh" #include +#include + +#include #include #include -#include #include +#include namespace cudf::detail { -template -rmm::device_uvector reduce_by_row(distinct_set_t& set, - size_type num_rows, - duplicate_keep_option keep, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +void initialize_reduction_results(size_type* results, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream) { - auto output_indices = rmm::device_uvector(num_rows, stream, mr); - - // If we don't care about order, just gather indices of distinct keys taken from set. - if (keep == duplicate_keep_option::KEEP_ANY) { - auto const iter = cuda::counting_iterator{0}; - set.insert_async(iter, iter + num_rows, stream.value()); - auto const output_end = set.retrieve_all(output_indices.begin(), stream.value()); - output_indices.resize(cuda::std::distance(output_indices.begin(), output_end), stream); - return output_indices; - } - - auto reduction_results = rmm::device_uvector(num_rows, stream, mr); thrust::uninitialized_fill( rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - reduction_results.begin(), - reduction_results.end(), + results, + results + num_rows, reduction_init_value(keep)); +} - auto set_ref = set.ref(cuco::op::insert_and_find); - - thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - cuda::counting_iterator{num_rows}, - [set_ref, keep, reduction_results = reduction_results.begin()] __device__( - size_type const idx) mutable { - auto const [inserted_idx_ptr, _] = set_ref.insert_and_find(idx); - - auto ref = cuda::atomic_ref{ - reduction_results[*inserted_idx_ptr]}; - if (keep == duplicate_keep_option::KEEP_FIRST) { - // Store the smallest index of all rows that are equal. - ref.fetch_min(idx, cuda::memory_order_relaxed); - } else if (keep == duplicate_keep_option::KEEP_LAST) { - // Store the greatest index of all rows that are equal. - ref.fetch_max(idx, cuda::memory_order_relaxed); - } else { - // Count the number of rows in each group of rows that are compared equal. - ref.fetch_add(size_type{1}, cuda::memory_order_relaxed); - } - }); - - auto const map_end = [&] { +size_type copy_reduction_results(size_type const* results, + size_type num_rows, + size_type* output, + duplicate_keep_option keep, + rmm::cuda_stream_view stream) +{ + auto const output_end = [&] { if (keep == duplicate_keep_option::KEEP_NONE) { - // Reduction results with `KEEP_NONE` are either group sizes of equal rows, or `0`. - // Thus, we only output index of the rows in the groups having group size of `1`. + // KEEP_NONE stores group sizes; retain only singleton groups. return cudf::detail::copy_if( cuda::counting_iterator{0}, cuda::counting_iterator{num_rows}, - output_indices.begin(), + output, cuda::proclaim_return_type( - [reduction_results = reduction_results.begin()] __device__(auto const idx) { - return reduction_results[idx] == size_type{1}; - }), + [results] __device__(auto const idx) { return results[idx] == size_type{1}; }), stream); } - // Reduction results with `KEEP_FIRST` and `KEEP_LAST` are row indices of the first/last row in - // each group of equal rows (which are the desired output indices), or the value given by - // `reduction_init_value()`. + // Other modes store desired row indices or the mode's initial marker. return cudf::detail::copy_if( - reduction_results.begin(), - reduction_results.end(), - output_indices.begin(), + results, + results + num_rows, + output, cuda::proclaim_return_type([init_value = reduction_init_value(keep)] __device__( auto const idx) { return idx != init_value; }), stream); }(); - output_indices.resize(cuda::std::distance(output_indices.begin(), map_end), stream); - return output_indices; + return cuda::std::distance(output, output_end); } template rmm::device_uvector reduce_by_row( - distinct_set_t>& set, - size_type num_rows, - duplicate_keep_option keep, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr); - -template rmm::device_uvector reduce_by_row( - distinct_set_t>& set, - size_type num_rows, - duplicate_keep_option keep, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr); - -template rmm::device_uvector reduce_by_row( - distinct_set_t>& set, - size_type num_rows, - duplicate_keep_option keep, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr); - -template rmm::device_uvector reduce_by_row( - distinct_set_t>& set, + distinct_set_t< + cudf::detail::row::equality:: + device_row_comparator>& set, size_type num_rows, duplicate_keep_option keep, rmm::cuda_stream_view stream, diff --git a/cpp/src/stream_compaction/distinct_helpers.cuh b/cpp/src/stream_compaction/distinct_helpers.cuh new file mode 100644 index 000000000000..9196582b2ee9 --- /dev/null +++ b/cpp/src/stream_compaction/distinct_helpers.cuh @@ -0,0 +1,68 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "distinct_helpers.hpp" + +#include + +#include + +#include +#include +#include +#include + +namespace cudf::detail { + +template +rmm::device_uvector reduce_by_row(distinct_set_t& set, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto output_indices = rmm::device_uvector(num_rows, stream, mr); + auto reduction_results = rmm::device_uvector(num_rows, stream, mr); + if (keep != duplicate_keep_option::KEEP_ANY) { + initialize_reduction_results(reduction_results.data(), num_rows, keep, stream); + } + + auto set_ref = set.ref(cuco::op::insert_and_find); + + thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + cuda::counting_iterator{0}, + cuda::counting_iterator{num_rows}, + [set_ref, keep, reduction_results = reduction_results.begin()] __device__( + size_type const idx) mutable { + auto const [inserted_idx_ptr, inserted] = set_ref.insert_and_find(idx); + + if (keep == duplicate_keep_option::KEEP_ANY) { + reduction_results[idx] = + inserted ? idx : cudf::detail::CUDF_SIZE_TYPE_SENTINEL; + } else { + auto ref = cuda::atomic_ref{ + reduction_results[*inserted_idx_ptr]}; + if (keep == duplicate_keep_option::KEEP_FIRST) { + // Store the smallest index of all rows that are equal. + ref.fetch_min(idx, cuda::memory_order_relaxed); + } else if (keep == duplicate_keep_option::KEEP_LAST) { + // Store the greatest index of all rows that are equal. + ref.fetch_max(idx, cuda::memory_order_relaxed); + } else { + // Count the number of rows in each group of rows that are compared equal. + ref.fetch_add(size_type{1}, cuda::memory_order_relaxed); + } + } + }); + + auto const output_size = + copy_reduction_results(reduction_results.data(), num_rows, output_indices.data(), keep, stream); + output_indices.resize(output_size, stream); + return output_indices; +} + +} // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers.hpp b/cpp/src/stream_compaction/distinct_helpers.hpp index 57cb4d5ff3d1..6c3d6dc1e9cc 100644 --- a/cpp/src/stream_compaction/distinct_helpers.hpp +++ b/cpp/src/stream_compaction/distinct_helpers.hpp @@ -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 */ @@ -17,16 +17,34 @@ #include #include -#include -#include +#include +#include + +#include +#include namespace cudf::detail { +struct distinct_physical_equality { + nan_equality compare_nans; + + template + __device__ constexpr bool operator()(Element const lhs, Element const rhs) const noexcept + { + if constexpr (cuda::std::is_floating_point_v) { + return lhs == rhs || (compare_nans == nan_equality::ALL_EQUAL && cuda::std::isnan(lhs) && + cuda::std::isnan(rhs)); + } else { + return lhs == rhs; + } + } +}; + /** - * @brief Return the reduction identity used to initialize results of `hash_reduce_by_row`. + * @brief Return the value used to initialize or mark reduction results. * - * @param keep A value of `duplicate_keep_option` type, must not be `KEEP_ANY`. - * @return The initial reduction value. + * @param keep A value of `duplicate_keep_option` type + * @return The reduction value */ auto constexpr reduction_init_value(duplicate_keep_option keep) { @@ -34,10 +52,22 @@ auto constexpr reduction_init_value(duplicate_keep_option keep) case duplicate_keep_option::KEEP_FIRST: return std::numeric_limits::max(); case duplicate_keep_option::KEEP_LAST: return std::numeric_limits::min(); case duplicate_keep_option::KEEP_NONE: return size_type{0}; - default: CUDF_UNREACHABLE("This function should not be called with KEEP_ANY"); + case duplicate_keep_option::KEEP_ANY: return CUDF_SIZE_TYPE_SENTINEL; + default: CUDF_UNREACHABLE("Invalid duplicate keep option"); } } +void initialize_reduction_results(size_type* results, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream); + +size_type copy_reduction_results(size_type const* results, + size_type num_rows, + size_type* output, + duplicate_keep_option keep, + rmm::cuda_stream_view stream); + template using distinct_set_t = cuco::static_set reduce_by_row( + distinct_set_t>& + set, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace cudf::detail From 5b14325245c35361f30a3b3abfa6cc1e45ceecee Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Sat, 18 Jul 2026 00:03:18 +0000 Subject: [PATCH 2/4] Reduce distinct helper compile times --- cpp/CMakeLists.txt | 9 +- cpp/src/hash/murmurhash3_x86_32.cu | 55 ++++++-- cpp/src/hash/murmurhash3_x86_32.cuh | 33 +++++ cpp/src/stream_compaction/distinct.cu | 76 ++++++++--- cpp/src/stream_compaction/distinct_helpers.cu | 13 +- .../stream_compaction/distinct_helpers.cuh | 60 +++++---- .../stream_compaction/distinct_helpers.hpp | 125 ++++++++++-------- .../distinct_helpers_flat_nan_equal_any.cu | 19 +++ ...distinct_helpers_flat_nan_equal_ordered.cu | 20 +++ .../distinct_helpers_flat_nan_unequal_any.cu | 19 +++ ...stinct_helpers_flat_nan_unequal_ordered.cu | 20 +++ .../distinct_helpers_nested_nan_equal_any.cu | 20 +++ ...stinct_helpers_nested_nan_equal_ordered.cu | 20 +++ ...distinct_helpers_nested_nan_unequal_any.cu | 20 +++ ...nct_helpers_nested_nan_unequal_ordered.cu} | 9 +- 15 files changed, 391 insertions(+), 127 deletions(-) create mode 100644 cpp/src/hash/murmurhash3_x86_32.cuh create mode 100644 cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_any.cu create mode 100644 cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_ordered.cu create mode 100644 cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_any.cu create mode 100644 cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_ordered.cu create mode 100644 cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_any.cu create mode 100644 cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_ordered.cu create mode 100644 cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_any.cu rename cpp/src/stream_compaction/{distinct_helpers_nested.cu => distinct_helpers_nested_nan_unequal_ordered.cu} (58%) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index c3997b835ad4..5b8946f2bbee 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1008,7 +1008,14 @@ add_library( src/stream_compaction/apply_boolean_mask.cu src/stream_compaction/distinct.cu src/stream_compaction/distinct_helpers.cu - src/stream_compaction/distinct_helpers_nested.cu + src/stream_compaction/distinct_helpers_flat_nan_equal_any.cu + src/stream_compaction/distinct_helpers_flat_nan_equal_ordered.cu + src/stream_compaction/distinct_helpers_flat_nan_unequal_any.cu + src/stream_compaction/distinct_helpers_flat_nan_unequal_ordered.cu + src/stream_compaction/distinct_helpers_nested_nan_equal_any.cu + src/stream_compaction/distinct_helpers_nested_nan_equal_ordered.cu + src/stream_compaction/distinct_helpers_nested_nan_unequal_any.cu + src/stream_compaction/distinct_helpers_nested_nan_unequal_ordered.cu src/stream_compaction/drop_nans.cu src/stream_compaction/drop_nulls.cu src/stream_compaction/filter/filter.cu diff --git a/cpp/src/hash/murmurhash3_x86_32.cu b/cpp/src/hash/murmurhash3_x86_32.cu index 3a28c1a15553..848cb71b3e5f 100644 --- a/cpp/src/hash/murmurhash3_x86_32.cu +++ b/cpp/src/hash/murmurhash3_x86_32.cu @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ +#include "murmurhash3_x86_32.cuh" + #include #include #include @@ -17,35 +19,60 @@ namespace cudf { namespace hashing { namespace detail { -std::unique_ptr murmurhash3_x86_32(table_view const& input, - uint32_t seed, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +namespace { + +template +std::unique_ptr murmurhash3_x86_32_impl( + std::shared_ptr const& input, + size_type num_rows, + uint32_t seed, + Nullate nulls, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { - auto output = make_numeric_column(data_type(type_to_id()), - input.num_rows(), - mask_state::UNALLOCATED, - stream, - mr); + auto output = make_numeric_column( + data_type(type_to_id()), num_rows, mask_state::UNALLOCATED, stream, mr); - if (input.num_rows() == 0) { return output; } + if (num_rows == 0) { return output; } - bool const nullable = has_nulls(input); - auto const row_hasher = cudf::detail::row::hash::row_hasher(input, stream); + auto const row_hasher = cudf::detail::row::hash::row_hasher(input); auto output_view = output->mutable_view(); // Compute the hash value for each row auto const output_begin = output_view.begin(); - auto const hasher = row_hasher.device_hasher(nullable, seed); + auto const hasher = row_hasher.device_hasher(nulls, seed); // thrust::tabulate is slow here, see NVIDIA/cccl#9070 CUDF_CUDA_TRY(cub::DeviceFor::Bulk( - input.num_rows(), + num_rows, [output_begin, hasher] __device__(size_type i) mutable { output_begin[i] = hasher(i); }, stream.value())); return output; } +} // namespace + +std::unique_ptr murmurhash3_x86_32(table_view const& input, + uint32_t seed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto const preprocessed_input = + cudf::detail::row::hash::preprocessed_table::create(input, stream); + return murmurhash3_x86_32_impl( + preprocessed_input, input.num_rows(), seed, nullate::DYNAMIC{has_nulls(input)}, stream, mr); +} + +std::unique_ptr murmurhash3_x86_32( + std::shared_ptr const& input, + size_type num_rows, + uint32_t seed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + return murmurhash3_x86_32_impl(input, num_rows, seed, nullate::YES{}, stream, mr); +} + } // namespace detail std::unique_ptr murmurhash3_x86_32(table_view const& input, diff --git a/cpp/src/hash/murmurhash3_x86_32.cuh b/cpp/src/hash/murmurhash3_x86_32.cuh new file mode 100644 index 000000000000..06d08bac8d5d --- /dev/null +++ b/cpp/src/hash/murmurhash3_x86_32.cuh @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include + +#include +#include + +namespace cudf { +class column; + +namespace detail::row::equality { +struct preprocessed_table; +} + +namespace hashing::detail { + +std::unique_ptr murmurhash3_x86_32( + std::shared_ptr const& input, + size_type num_rows, + uint32_t seed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace hashing::detail +} // namespace cudf diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu index ea9838de2058..c820f3b146f3 100644 --- a/cpp/src/stream_compaction/distinct.cu +++ b/cpp/src/stream_compaction/distinct.cu @@ -4,7 +4,9 @@ */ #include "distinct_helpers.hpp" +#include "hash/murmurhash3_x86_32.cuh" +#include #include #include #include @@ -50,9 +52,19 @@ rmm::device_uvector dispatch_row_equal( cudf::detail::row::equality::self_comparator row_equal, Func&& func) { - auto const d_equal = row_equal.equal_to( - nullate::DYNAMIC{has_nulls}, compare_nulls, distinct_physical_equality{compare_nans}); - return func(d_equal); + if (compare_nans == nan_equality::ALL_EQUAL) { + auto const d_equal = row_equal.equal_to( + nullate::DYNAMIC{has_nulls}, + compare_nulls, + cudf::detail::row::equality::nan_equal_physical_equality_comparator{}); + return func(d_equal); + } else { + auto const d_equal = + row_equal.equal_to(nullate::DYNAMIC{has_nulls}, + compare_nulls, + cudf::detail::row::equality::physical_equality_comparator{}); + return func(d_equal); + } } } // namespace @@ -77,24 +89,54 @@ rmm::device_uvector distinct_indices(table_view const& input, auto const row_hash = cudf::detail::row::hash::row_hasher(preprocessed_input); auto const row_equal = cudf::detail::row::equality::self_comparator(preprocessed_input); - auto const helper_func = [&](auto const& d_equal) { + auto const helper_func = [&](auto const& d_equal, auto const& d_hash, auto const& reduce_func) { using RowEqual = std::decay_t; - auto set = distinct_set_t{num_rows, - 0.5, // desired load factor - cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL}, - d_equal, - {row_hash.device_hasher(has_nulls)}, - {}, - {}, - rmm::mr::polymorphic_allocator{}, - stream.value()}; - return detail::reduce_by_row(set, num_rows, keep, stream, mr); + using RowHash = std::decay_t; + auto set = + distinct_set_t{num_rows, + 0.5, // desired load factor + cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL}, + d_equal, + d_hash, + {}, + {}, + rmm::mr::polymorphic_allocator{}, + stream.value()}; + return reduce_func(set); }; - if (cudf::detail::has_nested_columns(input)) { - return dispatch_row_equal(nulls_equal, nans_equal, has_nulls, row_equal, helper_func); + if (has_nested_columns) { + if (keep == duplicate_keep_option::KEEP_ANY) { + auto const hashes = + cudf::hashing::detail::murmurhash3_x86_32(preprocessed_input, + num_rows, + cudf::DEFAULT_HASH_SEED, + stream, + cudf::get_current_device_resource_ref()); + auto const d_hash = distinct_precomputed_hash{hashes->view().data()}; + return dispatch_row_equal( + nulls_equal, nans_equal, has_nulls, row_equal, [&](auto const& d_equal) { + return helper_func(d_equal, d_hash, [&](auto& set) { + return detail::reduce_by_row_keep_any(set, num_rows, stream, mr); + }); + }); + } + + auto const d_hash = row_hash.device_hasher(has_nulls); + return dispatch_row_equal( + nulls_equal, nans_equal, has_nulls, row_equal, [&](auto const& d_equal) { + return helper_func(d_equal, d_hash, [&](auto& set) { + return detail::reduce_by_row_keep_first_last_none(set, num_rows, keep, stream, mr); + }); + }); } else { - return dispatch_row_equal(nulls_equal, nans_equal, has_nulls, row_equal, helper_func); + auto const d_hash = row_hash.device_hasher(has_nulls); + return dispatch_row_equal( + nulls_equal, nans_equal, has_nulls, row_equal, [&](auto const& d_equal) { + return helper_func(d_equal, d_hash, [&](auto& set) { + return detail::reduce_by_row(set, num_rows, keep, stream, mr); + }); + }); } } diff --git a/cpp/src/stream_compaction/distinct_helpers.cu b/cpp/src/stream_compaction/distinct_helpers.cu index 3161a7b4c74c..09868820cbeb 100644 --- a/cpp/src/stream_compaction/distinct_helpers.cu +++ b/cpp/src/stream_compaction/distinct_helpers.cu @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "distinct_helpers.cuh" +#include "distinct_helpers.hpp" #include #include @@ -47,7 +47,7 @@ size_type copy_reduction_results(size_type const* results, stream); } - // Other modes store desired row indices or the mode's initial marker. + // KEEP_FIRST and KEEP_LAST store desired row indices or the mode's initial marker. return cudf::detail::copy_if( results, results + num_rows, @@ -60,13 +60,4 @@ size_type copy_reduction_results(size_type const* results, return cuda::std::distance(output, output_end); } -template rmm::device_uvector reduce_by_row( - distinct_set_t< - cudf::detail::row::equality:: - device_row_comparator>& set, - size_type num_rows, - duplicate_keep_option keep, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr); - } // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers.cuh b/cpp/src/stream_compaction/distinct_helpers.cuh index 9196582b2ee9..57d0ce3a0c23 100644 --- a/cpp/src/stream_compaction/distinct_helpers.cuh +++ b/cpp/src/stream_compaction/distinct_helpers.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -18,18 +18,31 @@ namespace cudf::detail { -template -rmm::device_uvector reduce_by_row(distinct_set_t& set, - size_type num_rows, - duplicate_keep_option keep, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +template +rmm::device_uvector reduce_by_row_keep_any(Set& set, + size_type num_rows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto output_indices = rmm::device_uvector(num_rows, stream, mr); + + auto const iter = cuda::counting_iterator{0}; + set.insert_async(iter, iter + num_rows, stream.value()); + auto const output_end = set.retrieve_all(output_indices.begin(), stream.value()); + output_indices.resize(cuda::std::distance(output_indices.begin(), output_end), stream); + return output_indices; +} + +template +rmm::device_uvector reduce_by_row_keep_first_last_none(Set& set, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { auto output_indices = rmm::device_uvector(num_rows, stream, mr); auto reduction_results = rmm::device_uvector(num_rows, stream, mr); - if (keep != duplicate_keep_option::KEEP_ANY) { - initialize_reduction_results(reduction_results.data(), num_rows, keep, stream); - } + initialize_reduction_results(reduction_results.data(), num_rows, keep, stream); auto set_ref = set.ref(cuco::op::insert_and_find); @@ -38,24 +51,19 @@ rmm::device_uvector reduce_by_row(distinct_set_t& set, cuda::counting_iterator{num_rows}, [set_ref, keep, reduction_results = reduction_results.begin()] __device__( size_type const idx) mutable { - auto const [inserted_idx_ptr, inserted] = set_ref.insert_and_find(idx); + auto const [inserted_idx_ptr, _] = set_ref.insert_and_find(idx); - if (keep == duplicate_keep_option::KEEP_ANY) { - reduction_results[idx] = - inserted ? idx : cudf::detail::CUDF_SIZE_TYPE_SENTINEL; + auto ref = cuda::atomic_ref{ + reduction_results[*inserted_idx_ptr]}; + if (keep == duplicate_keep_option::KEEP_FIRST) { + // Store the smallest index of all rows that are equal. + ref.fetch_min(idx, cuda::memory_order_relaxed); + } else if (keep == duplicate_keep_option::KEEP_LAST) { + // Store the greatest index of all rows that are equal. + ref.fetch_max(idx, cuda::memory_order_relaxed); } else { - auto ref = cuda::atomic_ref{ - reduction_results[*inserted_idx_ptr]}; - if (keep == duplicate_keep_option::KEEP_FIRST) { - // Store the smallest index of all rows that are equal. - ref.fetch_min(idx, cuda::memory_order_relaxed); - } else if (keep == duplicate_keep_option::KEEP_LAST) { - // Store the greatest index of all rows that are equal. - ref.fetch_max(idx, cuda::memory_order_relaxed); - } else { - // Count the number of rows in each group of rows that are compared equal. - ref.fetch_add(size_type{1}, cuda::memory_order_relaxed); - } + // Count the number of rows in each group of rows that are compared equal. + ref.fetch_add(size_type{1}, cuda::memory_order_relaxed); } }); diff --git a/cpp/src/stream_compaction/distinct_helpers.hpp b/cpp/src/stream_compaction/distinct_helpers.hpp index 6c3d6dc1e9cc..e7b17961f12a 100644 --- a/cpp/src/stream_compaction/distinct_helpers.hpp +++ b/cpp/src/stream_compaction/distinct_helpers.hpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -17,34 +18,16 @@ #include #include -#include -#include - -#include -#include +#include +#include namespace cudf::detail { -struct distinct_physical_equality { - nan_equality compare_nans; - - template - __device__ constexpr bool operator()(Element const lhs, Element const rhs) const noexcept - { - if constexpr (cuda::std::is_floating_point_v) { - return lhs == rhs || (compare_nans == nan_equality::ALL_EQUAL && cuda::std::isnan(lhs) && - cuda::std::isnan(rhs)); - } else { - return lhs == rhs; - } - } -}; - /** - * @brief Return the value used to initialize or mark reduction results. + * @brief Return the reduction identity used to initialize results of `hash_reduce_by_row`. * - * @param keep A value of `duplicate_keep_option` type - * @return The reduction value + * @param keep A value of `duplicate_keep_option` type, must not be `KEEP_ANY`. + * @return The initial reduction value. */ auto constexpr reduction_init_value(duplicate_keep_option keep) { @@ -52,34 +35,46 @@ auto constexpr reduction_init_value(duplicate_keep_option keep) case duplicate_keep_option::KEEP_FIRST: return std::numeric_limits::max(); case duplicate_keep_option::KEEP_LAST: return std::numeric_limits::min(); case duplicate_keep_option::KEEP_NONE: return size_type{0}; - case duplicate_keep_option::KEEP_ANY: return CUDF_SIZE_TYPE_SENTINEL; - default: CUDF_UNREACHABLE("Invalid duplicate keep option"); + default: CUDF_UNREACHABLE("This function should not be called with KEEP_ANY"); } } -void initialize_reduction_results(size_type* results, - size_type num_rows, - duplicate_keep_option keep, - rmm::cuda_stream_view stream); +CUDF_HIDDEN void initialize_reduction_results(size_type* results, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream); + +CUDF_HIDDEN size_type copy_reduction_results(size_type const* results, + size_type num_rows, + size_type* output, + duplicate_keep_option keep, + rmm::cuda_stream_view stream); + +struct distinct_precomputed_hash { + CUDF_HOST_DEVICE constexpr distinct_precomputed_hash(hash_value_type const* hashes) + : _hashes{hashes} + { + } + + __device__ __forceinline__ hash_value_type operator()(size_type i) const noexcept + { + return _hashes[i]; + } -size_type copy_reduction_results(size_type const* results, - size_type num_rows, - size_type* output, - duplicate_keep_option keep, - rmm::cuda_stream_view stream); + private: + hash_value_type const* _hashes; +}; -template -using distinct_set_t = - cuco::static_set, - cuda::thread_scope_device, - RowEqual, - cuco::linear_probing< - 1, - cudf::detail::row::hash::device_row_hasher>, - rmm::mr::polymorphic_allocator, - cuco::storage<1>>; +template > +using distinct_set_t = cuco::static_set, + cuda::thread_scope_device, + RowEqual, + cuco::linear_probing<1, RowHash>, + rmm::mr::polymorphic_allocator, + cuco::storage<1>>; /** * @brief Perform a reduction on groups of rows that are compared equal and returns output indices @@ -89,16 +84,18 @@ using distinct_set_t = * equal. A hash set is used to find groups of equal rows. * * Depending on the `keep` parameter, the reduction operation for each row group is: - * - If `keep == KEEP_ANY`: retain the row inserted into the set. + * - If `keep == KEEP_ANY` : order does not matter. * - If `keep == KEEP_FIRST`: min of row indices in the group. * - If `keep == KEEP_LAST`: max of row indices in the group. * - If `keep == KEEP_NONE`: count of equivalent rows (group size). * - * Except for `KEEP_ANY`, the result array is initialized with `reduction_init_value()`, then each - * row group writes its reduction at the index of an unspecified row in the group. For `KEEP_ANY`, - * each row stores either its index or the sentinel value, depending on whether it was inserted. + * Note that this function is not needed when `keep == KEEP_NONE`. + * + * At the beginning of the operation, the entire output array is filled with a value given by + * the `reduction_init_value()` function. Then, the reduction result for each row group is written + * into the output array at the index of an unspecified row in the group. * - * @tparam RowEqual The type of row equality comparator + * @tparam Set The type of the auxiliary set * * @param set The auxiliary set to perform reduction * @param num_rows The number of all input rows @@ -107,10 +104,30 @@ using distinct_set_t = * @param mr Device memory resource used to allocate the returned vector * @return A device_uvector containing the output indices */ -template -rmm::device_uvector reduce_by_row(distinct_set_t& set, +template +rmm::device_uvector reduce_by_row_keep_any(Set& set, + size_type num_rows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +template +rmm::device_uvector reduce_by_row_keep_first_last_none( + Set& set, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +template +rmm::device_uvector reduce_by_row(Set& set, size_type num_rows, duplicate_keep_option keep, rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr); + rmm::device_async_resource_ref mr) +{ + if (keep == duplicate_keep_option::KEEP_ANY) { + return reduce_by_row_keep_any(set, num_rows, stream, mr); + } + return reduce_by_row_keep_first_last_none(set, num_rows, keep, stream, mr); +} } // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_any.cu b/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_any.cu new file mode 100644 index 000000000000..61cd3d520e5e --- /dev/null +++ b/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_any.cu @@ -0,0 +1,19 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "distinct_helpers.cuh" + +namespace cudf::detail { + +template rmm::device_uvector reduce_by_row_keep_any( + distinct_set_t>& set, + size_type num_rows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_ordered.cu b/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_ordered.cu new file mode 100644 index 000000000000..c797133b96ff --- /dev/null +++ b/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_ordered.cu @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "distinct_helpers.cuh" + +namespace cudf::detail { + +template rmm::device_uvector reduce_by_row_keep_first_last_none( + distinct_set_t>& set, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_any.cu b/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_any.cu new file mode 100644 index 000000000000..c48bfd798c28 --- /dev/null +++ b/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_any.cu @@ -0,0 +1,19 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "distinct_helpers.cuh" + +namespace cudf::detail { + +template rmm::device_uvector reduce_by_row_keep_any( + distinct_set_t>& set, + size_type num_rows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_ordered.cu b/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_ordered.cu new file mode 100644 index 000000000000..dd9077704684 --- /dev/null +++ b/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_ordered.cu @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "distinct_helpers.cuh" + +namespace cudf::detail { + +template rmm::device_uvector reduce_by_row_keep_first_last_none( + distinct_set_t>& set, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_any.cu b/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_any.cu new file mode 100644 index 000000000000..416b2b6d146b --- /dev/null +++ b/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_any.cu @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "distinct_helpers.cuh" + +namespace cudf::detail { + +template rmm::device_uvector reduce_by_row_keep_any( + distinct_set_t, + distinct_precomputed_hash>& set, + size_type num_rows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_ordered.cu b/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_ordered.cu new file mode 100644 index 000000000000..79da72147233 --- /dev/null +++ b/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_ordered.cu @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "distinct_helpers.cuh" + +namespace cudf::detail { + +template rmm::device_uvector reduce_by_row_keep_first_last_none( + distinct_set_t>& set, + size_type num_rows, + duplicate_keep_option keep, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_any.cu b/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_any.cu new file mode 100644 index 000000000000..811a4088c581 --- /dev/null +++ b/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_any.cu @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "distinct_helpers.cuh" + +namespace cudf::detail { + +template rmm::device_uvector reduce_by_row_keep_any( + distinct_set_t, + distinct_precomputed_hash>& set, + size_type num_rows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace cudf::detail diff --git a/cpp/src/stream_compaction/distinct_helpers_nested.cu b/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_ordered.cu similarity index 58% rename from cpp/src/stream_compaction/distinct_helpers_nested.cu rename to cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_ordered.cu index aefc8214ee80..f57586b05693 100644 --- a/cpp/src/stream_compaction/distinct_helpers_nested.cu +++ b/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_ordered.cu @@ -7,10 +7,11 @@ namespace cudf::detail { -template rmm::device_uvector reduce_by_row( - distinct_set_t>& - set, +template rmm::device_uvector reduce_by_row_keep_first_last_none( + distinct_set_t>& set, size_type num_rows, duplicate_keep_option keep, rmm::cuda_stream_view stream, From 0412edef4b66de6ad245afb123f31c2c47ee7e9d Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Sat, 18 Jul 2026 00:35:13 +0000 Subject: [PATCH 3/4] Add direct includes to distinct helpers --- cpp/src/hash/murmurhash3_x86_32.cu | 9 +++++++++ cpp/src/stream_compaction/distinct.cu | 9 +++++++-- cpp/src/stream_compaction/distinct_helpers.cu | 4 ++++ cpp/src/stream_compaction/distinct_helpers.cuh | 6 ++++++ cpp/src/stream_compaction/distinct_helpers.hpp | 15 +++++++++++---- .../distinct_helpers_flat_nan_equal_any.cu | 8 ++++++++ .../distinct_helpers_flat_nan_equal_ordered.cu | 9 +++++++++ .../distinct_helpers_flat_nan_unequal_any.cu | 8 ++++++++ .../distinct_helpers_flat_nan_unequal_ordered.cu | 9 +++++++++ .../distinct_helpers_nested_nan_equal_any.cu | 8 ++++++++ .../distinct_helpers_nested_nan_equal_ordered.cu | 9 +++++++++ .../distinct_helpers_nested_nan_unequal_any.cu | 8 ++++++++ ...distinct_helpers_nested_nan_unequal_ordered.cu | 9 +++++++++ 13 files changed, 105 insertions(+), 6 deletions(-) diff --git a/cpp/src/hash/murmurhash3_x86_32.cu b/cpp/src/hash/murmurhash3_x86_32.cu index 848cb71b3e5f..f82d552456d7 100644 --- a/cpp/src/hash/murmurhash3_x86_32.cu +++ b/cpp/src/hash/murmurhash3_x86_32.cu @@ -4,17 +4,26 @@ */ #include "murmurhash3_x86_32.cuh" +#include #include #include #include +#include +#include #include #include +#include +#include #include #include +#include #include +#include +#include + namespace cudf { namespace hashing { namespace detail { diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu index c820f3b146f3..9db6bbd76c8c 100644 --- a/cpp/src/stream_compaction/distinct.cu +++ b/cpp/src/stream_compaction/distinct.cu @@ -7,22 +7,27 @@ #include "hash/murmurhash3_x86_32.cuh" #include -#include +#include #include #include -#include #include #include #include #include +#include +#include #include #include #include #include #include +#include #include #include +#include + +#include #include #include diff --git a/cpp/src/stream_compaction/distinct_helpers.cu b/cpp/src/stream_compaction/distinct_helpers.cu index 09868820cbeb..be5fa5da6e45 100644 --- a/cpp/src/stream_compaction/distinct_helpers.cu +++ b/cpp/src/stream_compaction/distinct_helpers.cu @@ -6,9 +6,13 @@ #include "distinct_helpers.hpp" #include +#include +#include #include +#include #include +#include #include #include diff --git a/cpp/src/stream_compaction/distinct_helpers.cuh b/cpp/src/stream_compaction/distinct_helpers.cuh index 57d0ce3a0c23..0772d9d250f1 100644 --- a/cpp/src/stream_compaction/distinct_helpers.cuh +++ b/cpp/src/stream_compaction/distinct_helpers.cuh @@ -7,13 +7,19 @@ #include "distinct_helpers.hpp" +#include +#include #include +#include +#include #include +#include #include #include #include +#include #include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers.hpp b/cpp/src/stream_compaction/distinct_helpers.hpp index e7b17961f12a..7746d58dd88b 100644 --- a/cpp/src/stream_compaction/distinct_helpers.hpp +++ b/cpp/src/stream_compaction/distinct_helpers.hpp @@ -5,21 +5,28 @@ #pragma once -#include -#include #include +#include +#include #include #include +#include #include #include #include #include #include +#include +#include +#include #include -#include -#include +#include +#include + +#include +#include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_any.cu b/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_any.cu index 61cd3d520e5e..91c112e03595 100644 --- a/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_any.cu +++ b/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_any.cu @@ -4,6 +4,14 @@ */ #include "distinct_helpers.cuh" +#include "distinct_helpers.hpp" + +#include +#include + +#include +#include +#include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_ordered.cu b/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_ordered.cu index c797133b96ff..f266862dcce6 100644 --- a/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_ordered.cu +++ b/cpp/src/stream_compaction/distinct_helpers_flat_nan_equal_ordered.cu @@ -4,6 +4,15 @@ */ #include "distinct_helpers.cuh" +#include "distinct_helpers.hpp" + +#include +#include +#include + +#include +#include +#include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_any.cu b/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_any.cu index c48bfd798c28..ce224e6d91d0 100644 --- a/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_any.cu +++ b/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_any.cu @@ -4,6 +4,14 @@ */ #include "distinct_helpers.cuh" +#include "distinct_helpers.hpp" + +#include +#include + +#include +#include +#include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_ordered.cu b/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_ordered.cu index dd9077704684..40a9a327e6c4 100644 --- a/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_ordered.cu +++ b/cpp/src/stream_compaction/distinct_helpers_flat_nan_unequal_ordered.cu @@ -4,6 +4,15 @@ */ #include "distinct_helpers.cuh" +#include "distinct_helpers.hpp" + +#include +#include +#include + +#include +#include +#include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_any.cu b/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_any.cu index 416b2b6d146b..e60ad6a2547e 100644 --- a/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_any.cu +++ b/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_any.cu @@ -4,6 +4,14 @@ */ #include "distinct_helpers.cuh" +#include "distinct_helpers.hpp" + +#include +#include + +#include +#include +#include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_ordered.cu b/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_ordered.cu index 79da72147233..3030a4c63e99 100644 --- a/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_ordered.cu +++ b/cpp/src/stream_compaction/distinct_helpers_nested_nan_equal_ordered.cu @@ -4,6 +4,15 @@ */ #include "distinct_helpers.cuh" +#include "distinct_helpers.hpp" + +#include +#include +#include + +#include +#include +#include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_any.cu b/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_any.cu index 811a4088c581..111ee665bd53 100644 --- a/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_any.cu +++ b/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_any.cu @@ -4,6 +4,14 @@ */ #include "distinct_helpers.cuh" +#include "distinct_helpers.hpp" + +#include +#include + +#include +#include +#include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_ordered.cu b/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_ordered.cu index f57586b05693..0a1b98929e81 100644 --- a/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_ordered.cu +++ b/cpp/src/stream_compaction/distinct_helpers_nested_nan_unequal_ordered.cu @@ -4,6 +4,15 @@ */ #include "distinct_helpers.cuh" +#include "distinct_helpers.hpp" + +#include +#include +#include + +#include +#include +#include namespace cudf::detail { From 0db0236369bbe9786ad93f79fd1dec66c3c65647 Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Tue, 28 Jul 2026 19:04:55 +0000 Subject: [PATCH 4/4] Address distinct helper review feedback --- cpp/src/stream_compaction/distinct_helpers.cu | 2 +- .../stream_compaction/distinct_helpers.cuh | 7 +-- .../stream_compaction/distinct_helpers.hpp | 51 +++++++++++-------- 3 files changed, 34 insertions(+), 26 deletions(-) diff --git a/cpp/src/stream_compaction/distinct_helpers.cu b/cpp/src/stream_compaction/distinct_helpers.cu index be5fa5da6e45..2596496c8774 100644 --- a/cpp/src/stream_compaction/distinct_helpers.cu +++ b/cpp/src/stream_compaction/distinct_helpers.cu @@ -17,7 +17,7 @@ #include #include #include -#include +#include namespace cudf::detail { diff --git a/cpp/src/stream_compaction/distinct_helpers.cuh b/cpp/src/stream_compaction/distinct_helpers.cuh index 0772d9d250f1..7649ea081380 100644 --- a/cpp/src/stream_compaction/distinct_helpers.cuh +++ b/cpp/src/stream_compaction/distinct_helpers.cuh @@ -46,8 +46,9 @@ rmm::device_uvector reduce_by_row_keep_first_last_none(Set& set, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - auto output_indices = rmm::device_uvector(num_rows, stream, mr); - auto reduction_results = rmm::device_uvector(num_rows, stream, mr); + auto output_indices = rmm::device_uvector(num_rows, stream, mr); + auto reduction_results = + rmm::device_uvector(num_rows, stream, cudf::get_current_device_resource_ref()); initialize_reduction_results(reduction_results.data(), num_rows, keep, stream); auto set_ref = set.ref(cuco::op::insert_and_find); @@ -56,7 +57,7 @@ rmm::device_uvector reduce_by_row_keep_first_last_none(Set& set, cuda::counting_iterator{0}, cuda::counting_iterator{num_rows}, [set_ref, keep, reduction_results = reduction_results.begin()] __device__( - size_type const idx) mutable { + size_type const idx) mutable -> void { auto const [inserted_idx_ptr, _] = set_ref.insert_and_find(idx); auto ref = cuda::atomic_ref{ diff --git a/cpp/src/stream_compaction/distinct_helpers.hpp b/cpp/src/stream_compaction/distinct_helpers.hpp index 7746d58dd88b..c7903aae7672 100644 --- a/cpp/src/stream_compaction/distinct_helpers.hpp +++ b/cpp/src/stream_compaction/distinct_helpers.hpp @@ -84,32 +84,14 @@ using distinct_set_t = cuco::static_set>; /** - * @brief Perform a reduction on groups of rows that are compared equal and returns output indices - * of the occurrences of the distinct elements based on `keep` parameter. - * - * This is essentially a reduce-by-key operation with keys are non-contiguous rows and are compared - * equal. A hash set is used to find groups of equal rows. - * - * Depending on the `keep` parameter, the reduction operation for each row group is: - * - If `keep == KEEP_ANY` : order does not matter. - * - If `keep == KEEP_FIRST`: min of row indices in the group. - * - If `keep == KEEP_LAST`: max of row indices in the group. - * - If `keep == KEEP_NONE`: count of equivalent rows (group size). - * - * Note that this function is not needed when `keep == KEEP_NONE`. - * - * At the beginning of the operation, the entire output array is filled with a value given by - * the `reduction_init_value()` function. Then, the reduction result for each row group is written - * into the output array at the index of an unspecified row in the group. + * @brief Returns one unspecified row index from each group of equal rows. * * @tparam Set The type of the auxiliary set - * - * @param set The auxiliary set to perform reduction - * @param num_rows The number of all input rows - * @param keep The parameter to determine what type of reduction to perform + * @param set The auxiliary set used to identify groups of equal rows + * @param num_rows The number of input rows * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used to allocate the returned vector - * @return A device_uvector containing the output indices + * @return A device vector containing one row index from each group */ template rmm::device_uvector reduce_by_row_keep_any(Set& set, @@ -117,6 +99,20 @@ rmm::device_uvector reduce_by_row_keep_any(Set& set, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); +/** + * @brief Returns row indices selected from groups of equal rows according to `keep`. + * + * `KEEP_FIRST` returns the smallest row index in each group, `KEEP_LAST` returns the largest, and + * `KEEP_NONE` returns indices only for singleton groups. + * + * @tparam Set The type of the auxiliary set + * @param set The auxiliary set used to identify groups of equal rows + * @param num_rows The number of input rows + * @param keep The duplicate selection mode; must not be `KEEP_ANY` + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the returned vector + * @return A device vector containing the selected row indices + */ template rmm::device_uvector reduce_by_row_keep_first_last_none( Set& set, @@ -125,6 +121,17 @@ rmm::device_uvector reduce_by_row_keep_first_last_none( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); +/** + * @brief Returns row indices selected from groups of equal rows according to `keep`. + * + * @tparam Set The type of the auxiliary set + * @param set The auxiliary set used to identify groups of equal rows + * @param num_rows The number of input rows + * @param keep The duplicate selection mode + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the returned vector + * @return A device vector containing the selected row indices + */ template rmm::device_uvector reduce_by_row(Set& set, size_type num_rows,