Reduce distinct_helpers build time

cpp/CMakeLists.txt

@@ -1008,6 +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_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

cpp/src/hash/murmurhash3_x86_32.cu

@@ -4,17 +4,26 @@
  */
 #include "murmurhash3_x86_32.cuh"
 
+#include <cudf/column/column.hpp>
 #include <cudf/column/column_factories.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/row_operator/hashing.cuh>
+#include <cudf/detail/row_operator/preprocessed_table.cuh>
+#include <cudf/hashing.hpp>
 #include <cudf/hashing/detail/hashing.hpp>
 #include <cudf/hashing/detail/murmurhash3_x86_32.cuh>
+#include <cudf/table/table_view.hpp>
+#include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <cub/device/device_for.cuh>
 
+#include <cstdint>
+#include <memory>
+
 namespace cudf {
 namespace hashing {
 namespace detail {

cpp/src/stream_compaction/distinct.cu

@@ -1,27 +1,36 @@
 /*
- * 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
  */
 
 #include "distinct_helpers.hpp"
+#include "hash/murmurhash3_x86_32.cuh"
 
-#include <cudf/column/column_view.hpp>
+#include <cudf/column/column.hpp>
+#include <cudf/copying.hpp>
 #include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/detail/gather.hpp>
-#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/row_operator/equality.cuh>
 #include <cudf/detail/row_operator/hashing.cuh>
 #include <cudf/detail/stream_compaction.hpp>
+#include <cudf/hashing.hpp>
+#include <cudf/stream_compaction.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_buffer.hpp>
 #include <rmm/device_uvector.hpp>
 #include <rmm/mr/polymorphic_allocator.hpp>
+#include <rmm/resource_ref.hpp>
 
+#include <cuco/types.cuh>
+
+#include <memory>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -85,24 +94,54 @@ rmm::device_uvector<size_type> 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<decltype(d_equal)>;
-    auto set       = distinct_set_t<RowEqual>{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<char>{},
-                                              stream.value()};
-    return detail::reduce_by_row(set, num_rows, keep, stream, mr);
+    using RowHash  = std::decay_t<decltype(d_hash)>;
+    auto set =
+      distinct_set_t<RowEqual, RowHash>{num_rows,
+                                        0.5,  // desired load factor
+                                        cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
+                                        d_equal,
+                                        d_hash,
+                                        {},
+                                        {},
+                                        rmm::mr::polymorphic_allocator<char>{},
+                                        stream.value()};
+    return reduce_func(set);
   };
 
-  if (cudf::detail::has_nested_columns(input)) {
-    return dispatch_row_equal<true>(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<hash_value_type>()};
+      return dispatch_row_equal<true>(
+        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<true>(
+      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<false>(nulls_equal, nans_equal, has_nulls, row_equal, helper_func);
+    auto const d_hash = row_hash.device_hasher(has_nulls);
+    return dispatch_row_equal<false>(
+      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);
+        });
+      });
   }
 }
 

cpp/src/stream_compaction/distinct_helpers.cu

@@ -1,136 +1,67 @@
 /*
- * 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 <cudf/detail/algorithms/copy_if.cuh>
+#include <cudf/stream_compaction.hpp>
+#include <cudf/types.hpp>
+#include <cudf/utilities/memory_resource.hpp>
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <cuda/functional>
 #include <cuda/iterator>
-#include <cuda/std/atomic>
 #include <cuda/std/iterator>
+#include <thrust/uninitialized_fill.h>
 
 namespace cudf::detail {
 
-template <typename RowEqual>
-rmm::device_uvector<size_type> reduce_by_row(distinct_set_t<RowEqual>& 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<size_type>(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<cudf::size_type>{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<size_type>(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<cudf::size_type>{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<size_type, cuda::thread_scope_device>{
-                       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<size_type>{0},
         cuda::counting_iterator<size_type>{num_rows},
-        output_indices.begin(),
+        output,
         cuda::proclaim_return_type<bool>(
-          [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()`.
+    // KEEP_FIRST and KEEP_LAST 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<bool>([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<size_type> reduce_by_row(
-  distinct_set_t<cudf::detail::row::equality::device_row_comparator<
-    false,
-    cudf::nullate::DYNAMIC,
-    cudf::detail::row::equality::nan_equal_physical_equality_comparator>>& set,
-  size_type num_rows,
-  duplicate_keep_option keep,
-  rmm::cuda_stream_view stream,
-  rmm::device_async_resource_ref mr);
-
-template rmm::device_uvector<size_type> reduce_by_row(
-  distinct_set_t<cudf::detail::row::equality::device_row_comparator<
-    true,
-    cudf::nullate::DYNAMIC,
-    cudf::detail::row::equality::nan_equal_physical_equality_comparator>>& set,
-  size_type num_rows,
-  duplicate_keep_option keep,
-  rmm::cuda_stream_view stream,
-  rmm::device_async_resource_ref mr);
-
-template rmm::device_uvector<size_type> reduce_by_row(
-  distinct_set_t<cudf::detail::row::equality::device_row_comparator<
-    false,
-    cudf::nullate::DYNAMIC,
-    cudf::detail::row::equality::physical_equality_comparator>>& set,
-  size_type num_rows,
-  duplicate_keep_option keep,
-  rmm::cuda_stream_view stream,
-  rmm::device_async_resource_ref mr);
-
-template rmm::device_uvector<size_type> reduce_by_row(
-  distinct_set_t<cudf::detail::row::equality::device_row_comparator<
-    true,
-    cudf::nullate::DYNAMIC,
-    cudf::detail::row::equality::physical_equality_comparator>>& set,
-  size_type num_rows,
-  duplicate_keep_option keep,
-  rmm::cuda_stream_view stream,
-  rmm::device_async_resource_ref mr);
-
 }  // namespace cudf::detail