From eb5a6eaf5cd4727b4f003d052d5bf71eca18a280 Mon Sep 17 00:00:00 2001
From: Srinivas Yadav Singanaboina <ssinganaboin@nvidia.com>
Date: Fri, 2 Aug 2024 18:03:21 +0000
Subject: [PATCH 1/8] Refactor distinct using insert_or_apply

---
 cpp/src/stream_compaction/distinct.cu         | 167 +++++++++++++++++-
 cpp/src/stream_compaction/distinct_helpers.cu | 118 -------------
 .../stream_compaction/distinct_helpers.hpp    |  62 -------
 3 files changed, 165 insertions(+), 182 deletions(-)

diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu
index e5cf29f3ebff..6fb7926f901d 100644
--- a/cpp/src/stream_compaction/distinct.cu
+++ b/cpp/src/stream_compaction/distinct.cu
@@ -32,6 +32,9 @@
 #include <rmm/mr/device/per_device_resource.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuco/static_map.cuh>
+
+#include <functional>
 #include <utility>
 #include <vector>
 
@@ -72,8 +75,146 @@ rmm::device_uvector<cudf::size_type> dipatch_row_equal(
     return func(d_equal);
   }
 }
+
+struct plus_op {
+  template <typename T, cuda::thread_scope Scope>
+  __device__ void operator()(cuda::atomic_ref<T, Scope> ref, T val)
+  {
+    ref.fetch_add(1, cuda::memory_order_relaxed);
+  }
+};
+
+struct min_op {
+  template <typename T, cuda::thread_scope Scope>
+  __device__ void operator()(cuda::atomic_ref<T, Scope> ref, T val)
+  {
+    ref.fetch_min(val, cuda::memory_order_relaxed);
+  }
+};
+
+struct max_op {
+  template <typename T, cuda::thread_scope Scope>
+  __device__ void operator()(cuda::atomic_ref<T, Scope> ref, T val)
+  {
+    ref.fetch_max(val, cuda::memory_order_relaxed);
+  }
+};
+
+template <typename Map>
+rmm::device_uvector<size_type> process_keep(Map& map,
+                                            size_type num_rows,
+                                            duplicate_keep_option keep,
+                                            rmm::cuda_stream_view stream,
+                                            rmm::device_async_resource_ref mr)
+{
+  if ((keep == duplicate_keep_option::KEEP_FIRST) or (keep == duplicate_keep_option::KEEP_LAST)) {
+    auto pairs = thrust::make_transform_iterator(
+      thrust::counting_iterator<size_type>(0),
+      cuda::proclaim_return_type<cuco::pair<size_type, size_type>>([] __device__(size_type i) {
+        return cuco::pair<size_type, size_type>{i, i};
+      }));
+
+    if (keep == duplicate_keep_option::KEEP_FIRST) {
+      map.insert_or_apply(pairs, pairs + num_rows, min_op{}, stream.value());
+    } else {
+      map.insert_or_apply(pairs, pairs + num_rows, max_op{}, stream.value());
+    }
+    int map_size = map.size(stream.value());
+    auto keys    = rmm::device_uvector<size_type>(map_size, stream, mr);
+    auto values  = rmm::device_uvector<size_type>(map_size, stream, mr);
+
+    map.retrieve_all(keys.begin(), values.begin(), stream);
+    return values;
+  }
+
+  auto pairs = thrust::make_transform_iterator(
+    thrust::counting_iterator<size_type>(0),
+    cuda::proclaim_return_type<cuco::pair<size_type, size_type>>([] __device__(size_type i) {
+      return cuco::pair<size_type, size_type>{i, 1};
+    }));
+
+  auto plusop = plus_op{};
+  map.insert_or_apply(pairs, pairs + num_rows, plusop, stream.value());
+
+  int map_size = map.size(stream.value());
+  auto keys    = rmm::device_uvector<size_type>(map_size, stream, mr);
+  auto values  = rmm::device_uvector<size_type>(map_size, stream, mr);
+  map.retrieve_all(keys.begin(), values.begin(), stream.value());
+
+  auto output_indices          = rmm::device_uvector<size_type>(map_size, stream, mr);
+  auto output_indices_filtered = rmm::device_uvector<size_type>(map_size, stream, mr);
+
+  thrust::for_each(
+    rmm::exec_policy(stream),
+    thrust::make_counting_iterator(0),
+    thrust::make_counting_iterator(map_size),
+    [values         = values.begin(),
+     keys           = keys.begin(),
+     output_indices = output_indices.begin()] __device__(size_type const idx) mutable {
+      if (values[idx] == size_type{1}) {
+        output_indices[idx] = keys[idx];
+      } else {
+        output_indices[idx] = -1;
+      }
+    });
+
+  auto const map_end = thrust::copy_if(
+    rmm::exec_policy(stream),
+    output_indices.begin(),
+    output_indices.end(),
+    output_indices_filtered.begin(),
+    cuda::proclaim_return_type<bool>([] __device__(auto const idx) { return idx != -1; }));
+
+  output_indices_filtered.resize(thrust::distance(output_indices_filtered.begin(), map_end),
+                                 stream);
+  return output_indices_filtered;
+}
+
 }  // namespace
 
+/**
+ * @brief Return the reduction identity used to initialize results of `hash_reduce_by_row`.
+ *
+ * @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)
+{
+  switch (keep) {
+    case duplicate_keep_option::KEEP_FIRST: return std::numeric_limits<size_type>::max();
+    case duplicate_keep_option::KEEP_LAST: return std::numeric_limits<size_type>::min();
+    case duplicate_keep_option::KEEP_NONE: return size_type{0};
+    default: CUDF_UNREACHABLE("This function should not be called with KEEP_ANY");
+  }
+}
+
+template <typename RowHasher>
+using hash_set_type =
+  cuco::static_set<size_type,
+                   cuco::extent<int64_t>,
+                   cuda::thread_scope_device,
+                   RowHasher,
+                   cuco::linear_probing<1,
+                                        cudf::experimental::row::hash::device_row_hasher<
+                                          cudf::hashing::detail::default_hash,
+                                          cudf::nullate::DYNAMIC>>,
+                   cudf::detail::cuco_allocator,
+                   cuco::storage<1>>;
+
+template <typename RowHasher>
+using hash_map_type =
+  cuco::static_map<size_type,
+                   size_type,
+                   cuco::extent<int64_t>,
+                   cuda::thread_scope_device,
+                   RowHasher,
+                   cuco::linear_probing<1,
+                                        cudf::experimental::row::hash::device_row_hasher<
+                                          cudf::hashing::detail::default_hash,
+                                          cudf::nullate::DYNAMIC>>,
+                   cudf::detail::cuco_allocator,
+                   cuco::storage<1>>;
+
 rmm::device_uvector<size_type> distinct_indices(table_view const& input,
                                                 duplicate_keep_option keep,
                                                 null_equality nulls_equal,
@@ -97,16 +238,38 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
 
   auto const helper_func = [&](auto const& d_equal) {
     using RowHasher = std::decay_t<decltype(d_equal)>;
-    auto set        = hash_set_type<RowHasher>{num_rows,
+    // If we don't care about order, just gather indices of distinct keys taken from set.
+    if (keep == duplicate_keep_option::KEEP_ANY) {
+      auto set = hash_set_type<RowHasher>{num_rows,
+                                          0.5,  // desired load factor
+                                          cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
+                                          d_equal,
+                                          {row_hash.device_hasher(has_nulls)},
+                                          {},
+                                          {},
+                                          cudf::detail::cuco_allocator{stream},
+                                          stream.value()};
+
+      auto const iter = thrust::counting_iterator<cudf::size_type>{0};
+      set.insert_async(iter, iter + num_rows, stream.value());
+      auto output_indices   = rmm::device_uvector<size_type>(num_rows, stream, mr);
+      auto const output_end = set.retrieve_all(output_indices.begin(), stream.value());
+      output_indices.resize(thrust::distance(output_indices.begin(), output_end), stream);
+      return output_indices;
+    }
+
+    auto const init = reduction_init_value(keep);
+    auto map        = hash_map_type<RowHasher>{num_rows,
                                                0.5,  // desired load factor
                                                cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
+                                               cuco::empty_value{init},
                                                d_equal,
                                                {row_hash.device_hasher(has_nulls)},
                                                {},
                                                {},
                                                cudf::detail::cuco_allocator{stream},
                                                stream.value()};
-    return detail::reduce_by_row(set, num_rows, keep, stream, mr);
+    return process_keep(map, num_rows, keep, stream, mr);
   };
 
   if (cudf::detail::has_nested_columns(input)) {
diff --git a/cpp/src/stream_compaction/distinct_helpers.cu b/cpp/src/stream_compaction/distinct_helpers.cu
index c3a004b7f289..826638d76be8 100644
--- a/cpp/src/stream_compaction/distinct_helpers.cu
+++ b/cpp/src/stream_compaction/distinct_helpers.cu
@@ -21,122 +21,4 @@
 
 namespace cudf::detail {
 
-template <typename RowHasher>
-rmm::device_uvector<size_type> reduce_by_row(hash_set_type<RowHasher>& 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<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 = thrust::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(thrust::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),
-                             reduction_results.begin(),
-                             reduction_results.end(),
-                             reduction_init_value(keep));
-
-  auto set_ref = set.ref(cuco::op::insert_and_find);
-
-  thrust::for_each(rmm::exec_policy_nosync(stream),
-                   thrust::make_counting_iterator(0),
-                   thrust::make_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 = [&] {
-    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`.
-      return thrust::copy_if(
-        rmm::exec_policy(stream),
-        thrust::make_counting_iterator(0),
-        thrust::make_counting_iterator(num_rows),
-        output_indices.begin(),
-        cuda::proclaim_return_type<bool>(
-          [reduction_results = reduction_results.begin()] __device__(auto const idx) {
-            return reduction_results[idx] == size_type{1};
-          }));
-    }
-
-    // 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()`.
-    return thrust::copy_if(
-      rmm::exec_policy(stream),
-      reduction_results.begin(),
-      reduction_results.end(),
-      output_indices.begin(),
-      cuda::proclaim_return_type<bool>([init_value = reduction_init_value(keep)] __device__(
-                                         auto const idx) { return idx != init_value; }));
-  }();
-
-  output_indices.resize(thrust::distance(output_indices.begin(), map_end), stream);
-  return output_indices;
-}
-
-template rmm::device_uvector<size_type> reduce_by_row(
-  hash_set_type<cudf::experimental::row::equality::device_row_comparator<
-    false,
-    cudf::nullate::DYNAMIC,
-    cudf::experimental::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(
-  hash_set_type<cudf::experimental::row::equality::device_row_comparator<
-    true,
-    cudf::nullate::DYNAMIC,
-    cudf::experimental::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(
-  hash_set_type<cudf::experimental::row::equality::device_row_comparator<
-    false,
-    cudf::nullate::DYNAMIC,
-    cudf::experimental::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(
-  hash_set_type<cudf::experimental::row::equality::device_row_comparator<
-    true,
-    cudf::nullate::DYNAMIC,
-    cudf::experimental::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
diff --git a/cpp/src/stream_compaction/distinct_helpers.hpp b/cpp/src/stream_compaction/distinct_helpers.hpp
index fca67c988730..0e3ad78e787d 100644
--- a/cpp/src/stream_compaction/distinct_helpers.hpp
+++ b/cpp/src/stream_compaction/distinct_helpers.hpp
@@ -31,66 +31,4 @@
 
 namespace cudf::detail {
 
-/**
- * @brief Return the reduction identity used to initialize results of `hash_reduce_by_row`.
- *
- * @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)
-{
-  switch (keep) {
-    case duplicate_keep_option::KEEP_FIRST: return std::numeric_limits<size_type>::max();
-    case duplicate_keep_option::KEEP_LAST: return std::numeric_limits<size_type>::min();
-    case duplicate_keep_option::KEEP_NONE: return size_type{0};
-    default: CUDF_UNREACHABLE("This function should not be called with KEEP_ANY");
-  }
-}
-
-template <typename RowHasher>
-using hash_set_type =
-  cuco::static_set<size_type,
-                   cuco::extent<int64_t>,
-                   cuda::thread_scope_device,
-                   RowHasher,
-                   cuco::linear_probing<1,
-                                        cudf::experimental::row::hash::device_row_hasher<
-                                          cudf::hashing::detail::default_hash,
-                                          cudf::nullate::DYNAMIC>>,
-                   cudf::detail::cuco_allocator,
-                   cuco::storage<1>>;
-
-/**
- * @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.
- *
- * @param set The auxiliary set to perform reduction
- * @param set_size The number of elements in set
- * @param num_rows The number of all input rows
- * @param keep The parameter to determine what type of reduction to perform
- * @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
- */
-template <typename RowHasher>
-rmm::device_uvector<size_type> reduce_by_row(hash_set_type<RowHasher>& set,
-                                             size_type num_rows,
-                                             duplicate_keep_option keep,
-                                             rmm::cuda_stream_view stream,
-                                             rmm::device_async_resource_ref mr);
 }  // namespace cudf::detail

From 86c2a227824c8a2d728356e12c5f9a0d627ab3a1 Mon Sep 17 00:00:00 2001
From: Srinivas Yadav Singanaboina <ssinganaboin@nvidia.com>
Date: Fri, 2 Aug 2024 22:29:29 +0000
Subject: [PATCH 2/8] clean up process_keep function

---
 cpp/src/stream_compaction/distinct.cu | 97 +++++++++++++--------------
 1 file changed, 47 insertions(+), 50 deletions(-)

diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu
index 6fb7926f901d..c6ab2dd4b602 100644
--- a/cpp/src/stream_compaction/distinct.cu
+++ b/cpp/src/stream_compaction/distinct.cu
@@ -33,6 +33,8 @@
 #include <rmm/resource_ref.hpp>
 
 #include <cuco/static_map.cuh>
+#include <cuco/static_set.cuh>
+#include <thrust/iterator/discard_iterator.h>
 
 #include <functional>
 #include <utility>
@@ -77,24 +79,24 @@ rmm::device_uvector<cudf::size_type> dipatch_row_equal(
 }
 
 struct plus_op {
-  template <typename T, cuda::thread_scope Scope>
-  __device__ void operator()(cuda::atomic_ref<T, Scope> ref, T val)
+  template <cuda::thread_scope Scope>
+  __device__ void operator()(cuda::atomic_ref<size_type, Scope> ref, size_type const val)
   {
-    ref.fetch_add(1, cuda::memory_order_relaxed);
+    ref.fetch_add(static_cast<size_type>(1), cuda::memory_order_relaxed);
   }
 };
 
 struct min_op {
-  template <typename T, cuda::thread_scope Scope>
-  __device__ void operator()(cuda::atomic_ref<T, Scope> ref, T val)
+  template <cuda::thread_scope Scope>
+  __device__ void operator()(cuda::atomic_ref<size_type, Scope> ref, size_type const val)
   {
     ref.fetch_min(val, cuda::memory_order_relaxed);
   }
 };
 
 struct max_op {
-  template <typename T, cuda::thread_scope Scope>
-  __device__ void operator()(cuda::atomic_ref<T, Scope> ref, T val)
+  template <cuda::thread_scope Scope>
+  __device__ void operator()(cuda::atomic_ref<size_type, Scope> ref, size_type const val)
   {
     ref.fetch_max(val, cuda::memory_order_relaxed);
   }
@@ -108,66 +110,61 @@ rmm::device_uvector<size_type> process_keep(Map& map,
                                             rmm::device_async_resource_ref mr)
 {
   if ((keep == duplicate_keep_option::KEEP_FIRST) or (keep == duplicate_keep_option::KEEP_LAST)) {
-    auto pairs = thrust::make_transform_iterator(
-      thrust::counting_iterator<size_type>(0),
-      cuda::proclaim_return_type<cuco::pair<size_type, size_type>>([] __device__(size_type i) {
-        return cuco::pair<size_type, size_type>{i, i};
-      }));
+    auto output_indices = rmm::device_uvector<size_type>(num_rows, stream, mr);
+
+    auto pairs =
+      thrust::make_transform_iterator(thrust::counting_iterator<size_type>(0),
+                                      cuda::proclaim_return_type<cuco::pair<size_type, size_type>>(
+                                        [] __device__(size_type const i) {
+                                          return cuco::pair<size_type, size_type>{i, i};
+                                        }));
 
     if (keep == duplicate_keep_option::KEEP_FIRST) {
-      map.insert_or_apply(pairs, pairs + num_rows, min_op{}, stream.value());
+      map.insert_or_apply_async(pairs, pairs + num_rows, min_op{}, stream.value());
     } else {
-      map.insert_or_apply(pairs, pairs + num_rows, max_op{}, stream.value());
+      map.insert_or_apply_async(pairs, pairs + num_rows, max_op{}, stream.value());
     }
-    int map_size = map.size(stream.value());
-    auto keys    = rmm::device_uvector<size_type>(map_size, stream, mr);
-    auto values  = rmm::device_uvector<size_type>(map_size, stream, mr);
 
-    map.retrieve_all(keys.begin(), values.begin(), stream);
-    return values;
+    auto const [_, output_end] =
+      map.retrieve_all(thrust::make_discard_iterator(), output_indices.begin(), stream.value());
+    output_indices.resize(thrust::distance(output_indices.begin(), output_end), stream);
+    return output_indices;
   }
 
+  auto keys   = rmm::device_uvector<size_type>(num_rows, stream, mr);
+  auto values = rmm::device_uvector<size_type>(num_rows, stream, mr);
+
   auto pairs = thrust::make_transform_iterator(
     thrust::counting_iterator<size_type>(0),
-    cuda::proclaim_return_type<cuco::pair<size_type, size_type>>([] __device__(size_type i) {
+    cuda::proclaim_return_type<cuco::pair<size_type, size_type>>([] __device__(size_type const i) {
       return cuco::pair<size_type, size_type>{i, 1};
     }));
 
-  auto plusop = plus_op{};
-  map.insert_or_apply(pairs, pairs + num_rows, plusop, stream.value());
-
-  int map_size = map.size(stream.value());
-  auto keys    = rmm::device_uvector<size_type>(map_size, stream, mr);
-  auto values  = rmm::device_uvector<size_type>(map_size, stream, mr);
-  map.retrieve_all(keys.begin(), values.begin(), stream.value());
-
-  auto output_indices          = rmm::device_uvector<size_type>(map_size, stream, mr);
-  auto output_indices_filtered = rmm::device_uvector<size_type>(map_size, stream, mr);
-
-  thrust::for_each(
-    rmm::exec_policy(stream),
-    thrust::make_counting_iterator(0),
-    thrust::make_counting_iterator(map_size),
-    [values         = values.begin(),
-     keys           = keys.begin(),
-     output_indices = output_indices.begin()] __device__(size_type const idx) mutable {
-      if (values[idx] == size_type{1}) {
-        output_indices[idx] = keys[idx];
-      } else {
-        output_indices[idx] = -1;
-      }
-    });
+  map.insert_or_apply_async(pairs, pairs + num_rows, plus_op{}, stream.value());
+  auto const [keys_end, _] = map.retrieve_all(keys.begin(), values.begin(), stream.value());
+
+  auto num_distinct_keys = thrust::distance(keys.begin(), keys_end);
+  keys.resize(num_distinct_keys, stream);
+  values.resize(num_distinct_keys, stream);
+
+  auto output_indices = rmm::device_uvector<size_type>(num_distinct_keys, stream, mr);
+
+  auto const output_iter = cudf::detail::make_counting_transform_iterator(
+    size_type(0),
+    cuda::proclaim_return_type<size_type>(
+      [keys = keys.begin(), values = values.begin()] __device__(auto const idx) {
+        return values[idx] == size_type{1} ? keys[idx] : -1;
+      }));
 
   auto const map_end = thrust::copy_if(
-    rmm::exec_policy(stream),
+    rmm::exec_policy_nosync(stream),
+    output_iter,
+    output_iter + num_distinct_keys,
     output_indices.begin(),
-    output_indices.end(),
-    output_indices_filtered.begin(),
     cuda::proclaim_return_type<bool>([] __device__(auto const idx) { return idx != -1; }));
 
-  output_indices_filtered.resize(thrust::distance(output_indices_filtered.begin(), map_end),
-                                 stream);
-  return output_indices_filtered;
+  output_indices.resize(thrust::distance(output_indices.begin(), map_end), stream);
+  return output_indices;
 }
 
 }  // namespace

From f2d8eef265bfd2dda8aa57acd65b701e59e5660c Mon Sep 17 00:00:00 2001
From: Srinivas Yadav Singanaboina <ssinganaboin@nvidia.com>
Date: Fri, 2 Aug 2024 23:02:34 +0000
Subject: [PATCH 3/8] minor cleanups

---
 cpp/src/stream_compaction/distinct.cu | 73 +++++++++------------------
 1 file changed, 25 insertions(+), 48 deletions(-)

diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu
index c6ab2dd4b602..e3ab96822b15 100644
--- a/cpp/src/stream_compaction/distinct.cu
+++ b/cpp/src/stream_compaction/distinct.cu
@@ -185,33 +185,6 @@ auto constexpr reduction_init_value(duplicate_keep_option keep)
   }
 }
 
-template <typename RowHasher>
-using hash_set_type =
-  cuco::static_set<size_type,
-                   cuco::extent<int64_t>,
-                   cuda::thread_scope_device,
-                   RowHasher,
-                   cuco::linear_probing<1,
-                                        cudf::experimental::row::hash::device_row_hasher<
-                                          cudf::hashing::detail::default_hash,
-                                          cudf::nullate::DYNAMIC>>,
-                   cudf::detail::cuco_allocator,
-                   cuco::storage<1>>;
-
-template <typename RowHasher>
-using hash_map_type =
-  cuco::static_map<size_type,
-                   size_type,
-                   cuco::extent<int64_t>,
-                   cuda::thread_scope_device,
-                   RowHasher,
-                   cuco::linear_probing<1,
-                                        cudf::experimental::row::hash::device_row_hasher<
-                                          cudf::hashing::detail::default_hash,
-                                          cudf::nullate::DYNAMIC>>,
-                   cudf::detail::cuco_allocator,
-                   cuco::storage<1>>;
-
 rmm::device_uvector<size_type> distinct_indices(table_view const& input,
                                                 duplicate_keep_option keep,
                                                 null_equality nulls_equal,
@@ -233,19 +206,24 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
   auto const row_hash  = cudf::experimental::row::hash::row_hasher(preprocessed_input);
   auto const row_equal = cudf::experimental::row::equality::self_comparator(preprocessed_input);
 
+  auto const probing_scheme = cuco::linear_probing<
+    1,
+    cudf::experimental::row::hash::device_row_hasher<cudf::hashing::detail::default_hash,
+                                                     cudf::nullate::DYNAMIC>>{
+    row_hash.device_hasher(has_nulls)};
+
   auto const helper_func = [&](auto const& d_equal) {
-    using RowHasher = std::decay_t<decltype(d_equal)>;
     // If we don't care about order, just gather indices of distinct keys taken from set.
     if (keep == duplicate_keep_option::KEEP_ANY) {
-      auto set = hash_set_type<RowHasher>{num_rows,
-                                          0.5,  // desired load factor
-                                          cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
-                                          d_equal,
-                                          {row_hash.device_hasher(has_nulls)},
-                                          {},
-                                          {},
-                                          cudf::detail::cuco_allocator{stream},
-                                          stream.value()};
+      auto set = cuco::static_set{num_rows,
+                                  0.5,  // desired load factor
+                                  cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
+                                  d_equal,
+                                  probing_scheme,
+                                  {},
+                                  {},
+                                  cudf::detail::cuco_allocator{stream},
+                                  stream.value()};
 
       auto const iter = thrust::counting_iterator<cudf::size_type>{0};
       set.insert_async(iter, iter + num_rows, stream.value());
@@ -255,17 +233,16 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
       return output_indices;
     }
 
-    auto const init = reduction_init_value(keep);
-    auto map        = hash_map_type<RowHasher>{num_rows,
-                                               0.5,  // desired load factor
-                                               cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
-                                               cuco::empty_value{init},
-                                               d_equal,
-                                               {row_hash.device_hasher(has_nulls)},
-                                               {},
-                                               {},
-                                               cudf::detail::cuco_allocator{stream},
-                                               stream.value()};
+    auto map = cuco::static_map{num_rows,
+                                0.5,  // desired load factor
+                                cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
+                                cuco::empty_value{reduction_init_value(keep)},
+                                d_equal,
+                                probing_scheme,
+                                {},
+                                {},
+                                cudf::detail::cuco_allocator{stream},
+                                stream.value()};
     return process_keep(map, num_rows, keep, stream, mr);
   };
 

From 5b662406c00e38890da4d39869e0ceacaf241dd0 Mon Sep 17 00:00:00 2001
From: Srinivas Yadav Singanaboina <ssinganaboin@nvidia.com>
Date: Fri, 2 Aug 2024 23:20:30 +0000
Subject: [PATCH 4/8] remove helper files

---
 cpp/CMakeLists.txt                            |  1 -
 cpp/src/stream_compaction/distinct.cu         |  5 +--
 cpp/src/stream_compaction/distinct_helpers.cu | 24 -------------
 .../stream_compaction/distinct_helpers.hpp    | 34 -------------------
 4 files changed, 3 insertions(+), 61 deletions(-)
 delete mode 100644 cpp/src/stream_compaction/distinct_helpers.cu
 delete mode 100644 cpp/src/stream_compaction/distinct_helpers.hpp

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 65347bd6689a..c1be38df98ea 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -563,7 +563,6 @@ add_library(
   src/stream_compaction/apply_boolean_mask.cu
   src/stream_compaction/distinct.cu
   src/stream_compaction/distinct_count.cu
-  src/stream_compaction/distinct_helpers.cu
   src/stream_compaction/drop_nans.cu
   src/stream_compaction/drop_nulls.cu
   src/stream_compaction/stable_distinct.cu
diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu
index e3ab96822b15..3c394dfc963e 100644
--- a/cpp/src/stream_compaction/distinct.cu
+++ b/cpp/src/stream_compaction/distinct.cu
@@ -14,8 +14,6 @@
  * limitations under the License.
  */
 
-#include "distinct_helpers.hpp"
-
 #include <cudf/column/column_view.hpp>
 #include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/detail/gather.hpp>
@@ -34,6 +32,9 @@
 
 #include <cuco/static_map.cuh>
 #include <cuco/static_set.cuh>
+#include <thrust/copy.h>
+#include <thrust/distance.h>
+#include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/discard_iterator.h>
 
 #include <functional>
diff --git a/cpp/src/stream_compaction/distinct_helpers.cu b/cpp/src/stream_compaction/distinct_helpers.cu
deleted file mode 100644
index 826638d76be8..000000000000
--- a/cpp/src/stream_compaction/distinct_helpers.cu
+++ /dev/null
@@ -1,24 +0,0 @@
-/*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "distinct_helpers.hpp"
-
-#include <cuda/functional>
-#include <cuda/std/atomic>
-
-namespace cudf::detail {
-
-}  // namespace cudf::detail
diff --git a/cpp/src/stream_compaction/distinct_helpers.hpp b/cpp/src/stream_compaction/distinct_helpers.hpp
deleted file mode 100644
index 0e3ad78e787d..000000000000
--- a/cpp/src/stream_compaction/distinct_helpers.hpp
+++ /dev/null
@@ -1,34 +0,0 @@
-/*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <cudf/detail/cuco_helpers.hpp>
-#include <cudf/stream_compaction.hpp>
-#include <cudf/table/experimental/row_operators.cuh>
-#include <cudf/types.hpp>
-
-#include <rmm/cuda_stream_view.hpp>
-#include <rmm/device_uvector.hpp>
-#include <rmm/resource_ref.hpp>
-
-#include <cuco/static_set.cuh>
-#include <cuda/functional>
-#include <thrust/copy.h>
-#include <thrust/distance.h>
-#include <thrust/iterator/counting_iterator.h>
-
-namespace cudf::detail {
-
-}  // namespace cudf::detail

From 213e1ce046f2fc9ccc88d0baf23bc56c51fdb45f Mon Sep 17 00:00:00 2001
From: Srinivas Yadav Singanaboina <ssinganaboin@nvidia.com>
Date: Sat, 3 Aug 2024 04:41:54 +0000
Subject: [PATCH 5/8] Revert "remove helper files"

reverting back this commit, as compilation time of distinct.cu
is almost 6 mins.

This reverts commit 5b662406c00e38890da4d39869e0ceacaf241dd0.
---
 cpp/CMakeLists.txt                            |  1 +
 cpp/src/stream_compaction/distinct.cu         |  5 ++-
 cpp/src/stream_compaction/distinct_helpers.cu | 24 +++++++++++++
 .../stream_compaction/distinct_helpers.hpp    | 34 +++++++++++++++++++
 4 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 cpp/src/stream_compaction/distinct_helpers.cu
 create mode 100644 cpp/src/stream_compaction/distinct_helpers.hpp

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 349cc95423a9..310bc99b2791 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -565,6 +565,7 @@ add_library(
   src/stream_compaction/apply_boolean_mask.cu
   src/stream_compaction/distinct.cu
   src/stream_compaction/distinct_count.cu
+  src/stream_compaction/distinct_helpers.cu
   src/stream_compaction/drop_nans.cu
   src/stream_compaction/drop_nulls.cu
   src/stream_compaction/stable_distinct.cu
diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu
index 33278a0554bd..e30443bb6569 100644
--- a/cpp/src/stream_compaction/distinct.cu
+++ b/cpp/src/stream_compaction/distinct.cu
@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+#include "distinct_helpers.hpp"
+
 #include <cudf/column/column_view.hpp>
 #include <cudf/detail/cuco_helpers.hpp>
 #include <cudf/detail/gather.hpp>
@@ -32,9 +34,6 @@
 
 #include <cuco/static_map.cuh>
 #include <cuco/static_set.cuh>
-#include <thrust/copy.h>
-#include <thrust/distance.h>
-#include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/discard_iterator.h>
 
 #include <functional>
diff --git a/cpp/src/stream_compaction/distinct_helpers.cu b/cpp/src/stream_compaction/distinct_helpers.cu
new file mode 100644
index 000000000000..826638d76be8
--- /dev/null
+++ b/cpp/src/stream_compaction/distinct_helpers.cu
@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "distinct_helpers.hpp"
+
+#include <cuda/functional>
+#include <cuda/std/atomic>
+
+namespace cudf::detail {
+
+}  // namespace cudf::detail
diff --git a/cpp/src/stream_compaction/distinct_helpers.hpp b/cpp/src/stream_compaction/distinct_helpers.hpp
new file mode 100644
index 000000000000..0e3ad78e787d
--- /dev/null
+++ b/cpp/src/stream_compaction/distinct_helpers.hpp
@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cudf/detail/cuco_helpers.hpp>
+#include <cudf/stream_compaction.hpp>
+#include <cudf/table/experimental/row_operators.cuh>
+#include <cudf/types.hpp>
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_uvector.hpp>
+#include <rmm/resource_ref.hpp>
+
+#include <cuco/static_set.cuh>
+#include <cuda/functional>
+#include <thrust/copy.h>
+#include <thrust/distance.h>
+#include <thrust/iterator/counting_iterator.h>
+
+namespace cudf::detail {
+
+}  // namespace cudf::detail

From 9f7db16b08579e641bd97f01d88fa9f193e6c1b4 Mon Sep 17 00:00:00 2001
From: Srinivas Yadav Singanaboina <ssinganaboin@nvidia.com>
Date: Sun, 4 Aug 2024 02:03:49 +0000
Subject: [PATCH 6/8] Split the implementation due to high compile times

---
 cpp/src/stream_compaction/distinct.cu         | 156 +++---------------
 cpp/src/stream_compaction/distinct_helpers.cu | 105 ++++++++++++
 .../stream_compaction/distinct_helpers.hpp    | 103 ++++++++++++
 3 files changed, 230 insertions(+), 134 deletions(-)

diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu
index e30443bb6569..5067a9ea7a18 100644
--- a/cpp/src/stream_compaction/distinct.cu
+++ b/cpp/src/stream_compaction/distinct.cu
@@ -32,9 +32,7 @@
 #include <rmm/mr/device/per_device_resource.hpp>
 #include <rmm/resource_ref.hpp>
 
-#include <cuco/static_map.cuh>
 #include <cuco/static_set.cuh>
-#include <thrust/iterator/discard_iterator.h>
 
 #include <functional>
 #include <utility>
@@ -77,114 +75,8 @@ rmm::device_uvector<cudf::size_type> dispatch_row_equal(
     return func(d_equal);
   }
 }
-
-struct plus_op {
-  template <cuda::thread_scope Scope>
-  __device__ void operator()(cuda::atomic_ref<size_type, Scope> ref, size_type const val)
-  {
-    ref.fetch_add(static_cast<size_type>(1), cuda::memory_order_relaxed);
-  }
-};
-
-struct min_op {
-  template <cuda::thread_scope Scope>
-  __device__ void operator()(cuda::atomic_ref<size_type, Scope> ref, size_type const val)
-  {
-    ref.fetch_min(val, cuda::memory_order_relaxed);
-  }
-};
-
-struct max_op {
-  template <cuda::thread_scope Scope>
-  __device__ void operator()(cuda::atomic_ref<size_type, Scope> ref, size_type const val)
-  {
-    ref.fetch_max(val, cuda::memory_order_relaxed);
-  }
-};
-
-template <typename Map>
-rmm::device_uvector<size_type> process_keep(Map& map,
-                                            size_type num_rows,
-                                            duplicate_keep_option keep,
-                                            rmm::cuda_stream_view stream,
-                                            rmm::device_async_resource_ref mr)
-{
-  if ((keep == duplicate_keep_option::KEEP_FIRST) or (keep == duplicate_keep_option::KEEP_LAST)) {
-    auto output_indices = rmm::device_uvector<size_type>(num_rows, stream, mr);
-
-    auto pairs =
-      thrust::make_transform_iterator(thrust::counting_iterator<size_type>(0),
-                                      cuda::proclaim_return_type<cuco::pair<size_type, size_type>>(
-                                        [] __device__(size_type const i) {
-                                          return cuco::pair<size_type, size_type>{i, i};
-                                        }));
-
-    if (keep == duplicate_keep_option::KEEP_FIRST) {
-      map.insert_or_apply_async(pairs, pairs + num_rows, min_op{}, stream.value());
-    } else {
-      map.insert_or_apply_async(pairs, pairs + num_rows, max_op{}, stream.value());
-    }
-
-    auto const [_, output_end] =
-      map.retrieve_all(thrust::make_discard_iterator(), output_indices.begin(), stream.value());
-    output_indices.resize(thrust::distance(output_indices.begin(), output_end), stream);
-    return output_indices;
-  }
-
-  auto keys   = rmm::device_uvector<size_type>(num_rows, stream, mr);
-  auto values = rmm::device_uvector<size_type>(num_rows, stream, mr);
-
-  auto pairs = thrust::make_transform_iterator(
-    thrust::counting_iterator<size_type>(0),
-    cuda::proclaim_return_type<cuco::pair<size_type, size_type>>([] __device__(size_type const i) {
-      return cuco::pair<size_type, size_type>{i, 1};
-    }));
-
-  map.insert_or_apply_async(pairs, pairs + num_rows, plus_op{}, stream.value());
-  auto const [keys_end, _] = map.retrieve_all(keys.begin(), values.begin(), stream.value());
-
-  auto num_distinct_keys = thrust::distance(keys.begin(), keys_end);
-  keys.resize(num_distinct_keys, stream);
-  values.resize(num_distinct_keys, stream);
-
-  auto output_indices = rmm::device_uvector<size_type>(num_distinct_keys, stream, mr);
-
-  auto const output_iter = cudf::detail::make_counting_transform_iterator(
-    size_type(0),
-    cuda::proclaim_return_type<size_type>(
-      [keys = keys.begin(), values = values.begin()] __device__(auto const idx) {
-        return values[idx] == size_type{1} ? keys[idx] : -1;
-      }));
-
-  auto const map_end = thrust::copy_if(
-    rmm::exec_policy_nosync(stream),
-    output_iter,
-    output_iter + num_distinct_keys,
-    output_indices.begin(),
-    cuda::proclaim_return_type<bool>([] __device__(auto const idx) { return idx != -1; }));
-
-  output_indices.resize(thrust::distance(output_indices.begin(), map_end), stream);
-  return output_indices;
-}
-
 }  // namespace
 
-/**
- * @brief Return the reduction identity used to initialize results of `hash_reduce_by_row`.
- *
- * @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)
-{
-  switch (keep) {
-    case duplicate_keep_option::KEEP_FIRST: return std::numeric_limits<size_type>::max();
-    case duplicate_keep_option::KEEP_LAST: return std::numeric_limits<size_type>::min();
-    case duplicate_keep_option::KEEP_NONE: return size_type{0};
-    default: CUDF_UNREACHABLE("This function should not be called with KEEP_ANY");
-  }
-}
-
 rmm::device_uvector<size_type> distinct_indices(table_view const& input,
                                                 duplicate_keep_option keep,
                                                 null_equality nulls_equal,
@@ -206,24 +98,20 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
   auto const row_hash  = cudf::experimental::row::hash::row_hasher(preprocessed_input);
   auto const row_equal = cudf::experimental::row::equality::self_comparator(preprocessed_input);
 
-  auto const probing_scheme = cuco::linear_probing<
-    1,
-    cudf::experimental::row::hash::device_row_hasher<cudf::hashing::detail::default_hash,
-                                                     cudf::nullate::DYNAMIC>>{
-    row_hash.device_hasher(has_nulls)};
-
   auto const helper_func = [&](auto const& d_equal) {
+    using RowHasher = cuda::std::decay_t<decltype(d_equal)>;
+
     // If we don't care about order, just gather indices of distinct keys taken from set.
     if (keep == duplicate_keep_option::KEEP_ANY) {
-      auto set = cuco::static_set{num_rows,
-                                  0.5,  // desired load factor
-                                  cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
-                                  d_equal,
-                                  probing_scheme,
-                                  {},
-                                  {},
-                                  cudf::detail::cuco_allocator{stream},
-                                  stream.value()};
+      auto set = hash_set_type<RowHasher>{num_rows,
+                                          0.5,  // desired load factor
+                                          cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
+                                          d_equal,
+                                          {row_hash.device_hasher(has_nulls)},
+                                          {},
+                                          {},
+                                          cudf::detail::cuco_allocator{stream},
+                                          stream.value()};
 
       auto const iter = thrust::counting_iterator<cudf::size_type>{0};
       set.insert_async(iter, iter + num_rows, stream.value());
@@ -233,17 +121,17 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
       return output_indices;
     }
 
-    auto map = cuco::static_map{num_rows,
-                                0.5,  // desired load factor
-                                cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
-                                cuco::empty_value{reduction_init_value(keep)},
-                                d_equal,
-                                probing_scheme,
-                                {},
-                                {},
-                                cudf::detail::cuco_allocator{stream},
-                                stream.value()};
-    return process_keep(map, num_rows, keep, stream, mr);
+    auto map = hash_map_type<RowHasher>{num_rows,
+                                        0.5,  // desired load factor
+                                        cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
+                                        cuco::empty_value{reduction_init_value(keep)},
+                                        d_equal,
+                                        {row_hash.device_hasher(has_nulls)},
+                                        {},
+                                        {},
+                                        cudf::detail::cuco_allocator{stream},
+                                        stream.value()};
+    return reduce_by_row(map, num_rows, keep, stream, mr);
   };
 
   if (cudf::detail::has_nested_columns(input)) {
diff --git a/cpp/src/stream_compaction/distinct_helpers.cu b/cpp/src/stream_compaction/distinct_helpers.cu
index 826638d76be8..882be4532dc5 100644
--- a/cpp/src/stream_compaction/distinct_helpers.cu
+++ b/cpp/src/stream_compaction/distinct_helpers.cu
@@ -21,4 +21,109 @@
 
 namespace cudf::detail {
 
+template <typename RowHasher>
+rmm::device_uvector<size_type> reduce_by_row(hash_map_type<RowHasher>& map,
+                                             size_type num_rows,
+                                             duplicate_keep_option keep,
+                                             rmm::cuda_stream_view stream,
+                                             rmm::device_async_resource_ref mr)
+{
+  if ((keep == duplicate_keep_option::KEEP_FIRST) or (keep == duplicate_keep_option::KEEP_LAST)) {
+    auto output_indices = rmm::device_uvector<size_type>(num_rows, stream, mr);
+
+    auto pairs =
+      thrust::make_transform_iterator(thrust::counting_iterator<size_type>(0),
+                                      cuda::proclaim_return_type<cuco::pair<size_type, size_type>>(
+                                        [] __device__(size_type const i) {
+                                          return cuco::pair<size_type, size_type>{i, i};
+                                        }));
+
+    if (keep == duplicate_keep_option::KEEP_FIRST) {
+      map.insert_or_apply_async(pairs, pairs + num_rows, min_op{}, stream.value());
+    } else {
+      map.insert_or_apply_async(pairs, pairs + num_rows, max_op{}, stream.value());
+    }
+
+    auto const [_, output_end] =
+      map.retrieve_all(thrust::make_discard_iterator(), output_indices.begin(), stream.value());
+    output_indices.resize(thrust::distance(output_indices.begin(), output_end), stream);
+    return output_indices;
+  }
+
+  auto keys   = rmm::device_uvector<size_type>(num_rows, stream);
+  auto values = rmm::device_uvector<size_type>(num_rows, stream);
+
+  auto pairs = thrust::make_transform_iterator(
+    thrust::counting_iterator<size_type>(0),
+    cuda::proclaim_return_type<cuco::pair<size_type, size_type>>([] __device__(size_type const i) {
+      return cuco::pair<size_type, size_type>{i, 1};
+    }));
+
+  map.insert_or_apply_async(pairs, pairs + num_rows, plus_op{}, stream.value());
+  auto const [keys_end, _] = map.retrieve_all(keys.begin(), values.begin(), stream.value());
+
+  auto num_distinct_keys = thrust::distance(keys.begin(), keys_end);
+  keys.resize(num_distinct_keys, stream);
+  values.resize(num_distinct_keys, stream);
+
+  auto output_indices = rmm::device_uvector<size_type>(num_distinct_keys, stream, mr);
+
+  auto const output_iter = cudf::detail::make_counting_transform_iterator(
+    size_type(0),
+    cuda::proclaim_return_type<size_type>(
+      [keys = keys.begin(), values = values.begin()] __device__(auto const idx) {
+        return values[idx] == size_type{1} ? keys[idx] : -1;
+      }));
+
+  auto const map_end = thrust::copy_if(
+    rmm::exec_policy_nosync(stream),
+    output_iter,
+    output_iter + num_distinct_keys,
+    output_indices.begin(),
+    cuda::proclaim_return_type<bool>([] __device__(auto const idx) { return idx != -1; }));
+
+  output_indices.resize(thrust::distance(output_indices.begin(), map_end), stream);
+  return output_indices;
+}
+
+template rmm::device_uvector<size_type> reduce_by_row(
+  hash_map_type<cudf::experimental::row::equality::device_row_comparator<
+    false,
+    cudf::nullate::DYNAMIC,
+    cudf::experimental::row::equality::nan_equal_physical_equality_comparator>>& map,
+  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(
+  hash_map_type<cudf::experimental::row::equality::device_row_comparator<
+    true,
+    cudf::nullate::DYNAMIC,
+    cudf::experimental::row::equality::nan_equal_physical_equality_comparator>>& map,
+  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(
+  hash_map_type<cudf::experimental::row::equality::device_row_comparator<
+    false,
+    cudf::nullate::DYNAMIC,
+    cudf::experimental::row::equality::physical_equality_comparator>>& map,
+  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(
+  hash_map_type<cudf::experimental::row::equality::device_row_comparator<
+    true,
+    cudf::nullate::DYNAMIC,
+    cudf::experimental::row::equality::physical_equality_comparator>>& map,
+  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.hpp b/cpp/src/stream_compaction/distinct_helpers.hpp
index 0e3ad78e787d..64efbeec200d 100644
--- a/cpp/src/stream_compaction/distinct_helpers.hpp
+++ b/cpp/src/stream_compaction/distinct_helpers.hpp
@@ -23,12 +23,115 @@
 #include <rmm/device_uvector.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuco/static_map.cuh>
 #include <cuco/static_set.cuh>
 #include <cuda/functional>
 #include <thrust/copy.h>
 #include <thrust/distance.h>
 #include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/discard_iterator.h>
 
 namespace cudf::detail {
 
+/**
+ * @brief Return the reduction identity used to initialize results of `hash_reduce_by_row`.
+ *
+ * @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)
+{
+  switch (keep) {
+    case duplicate_keep_option::KEEP_FIRST: return std::numeric_limits<size_type>::max();
+    case duplicate_keep_option::KEEP_LAST: return std::numeric_limits<size_type>::min();
+    case duplicate_keep_option::KEEP_NONE: return size_type{0};
+    default: CUDF_UNREACHABLE("This function should not be called with KEEP_ANY");
+  }
+}
+
+struct plus_op {
+  template <cuda::thread_scope Scope>
+  __device__ void operator()(cuda::atomic_ref<size_type, Scope> ref, size_type const val)
+  {
+    ref.fetch_add(static_cast<size_type>(1), cuda::memory_order_relaxed);
+  }
+};
+
+struct min_op {
+  template <cuda::thread_scope Scope>
+  __device__ void operator()(cuda::atomic_ref<size_type, Scope> ref, size_type const val)
+  {
+    ref.fetch_min(val, cuda::memory_order_relaxed);
+  }
+};
+
+struct max_op {
+  template <cuda::thread_scope Scope>
+  __device__ void operator()(cuda::atomic_ref<size_type, Scope> ref, size_type const val)
+  {
+    ref.fetch_max(val, cuda::memory_order_relaxed);
+  }
+};
+
+// The static_set type used to process `keep_any` option
+template <typename RowHasher>
+using hash_set_type =
+  cuco::static_set<size_type,
+                   cuco::extent<int64_t>,
+                   cuda::thread_scope_device,
+                   RowHasher,
+                   cuco::linear_probing<1,
+                                        cudf::experimental::row::hash::device_row_hasher<
+                                          cudf::hashing::detail::default_hash,
+                                          cudf::nullate::DYNAMIC>>,
+                   cudf::detail::cuco_allocator,
+                   cuco::storage<1>>;
+
+// The static_map type used to process `keep_first`, `keep_last` and `keep_none` option
+template <typename RowHasher>
+using hash_map_type =
+  cuco::static_map<size_type,
+                   size_type,
+                   cuco::extent<int64_t>,
+                   cuda::thread_scope_device,
+                   RowHasher,
+                   cuco::linear_probing<1,
+                                        cudf::experimental::row::hash::device_row_hasher<
+                                          cudf::hashing::detail::default_hash,
+                                          cudf::nullate::DYNAMIC>>,
+                   cudf::detail::cuco_allocator,
+                   cuco::storage<1>>;
+
+/**
+ * @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 map 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.
+ *
+ * @param map The auxiliary map 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 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
+ */
+template <typename RowHasher>
+rmm::device_uvector<size_type> reduce_by_row(hash_map_type<RowHasher>& map,
+                                             size_type num_rows,
+                                             duplicate_keep_option keep,
+                                             rmm::cuda_stream_view stream,
+                                             rmm::device_async_resource_ref mr);
 }  // namespace cudf::detail

From c10e0474393ab8914f2322d95e37e6f0d994cc36 Mon Sep 17 00:00:00 2001
From: Srinivas Yadav Singanaboina <ssinganaboin@nvidia.com>
Date: Thu, 8 Aug 2024 18:26:56 +0000
Subject: [PATCH 7/8] fix formatting

---
 cpp/src/stream_compaction/distinct.cu | 40 ++++++++++++++-------------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/cpp/src/stream_compaction/distinct.cu b/cpp/src/stream_compaction/distinct.cu
index 7f0e96b72e59..abd4bb41e8e7 100644
--- a/cpp/src/stream_compaction/distinct.cu
+++ b/cpp/src/stream_compaction/distinct.cu
@@ -103,15 +103,16 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
 
     // If we don't care about order, just gather indices of distinct keys taken from set.
     if (keep == duplicate_keep_option::KEEP_ANY) {
-      auto set = hash_set_type<RowHasher>{num_rows,
-                                          0.5,  // desired load factor
-                                          cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
-                                          d_equal,
-                                          {row_hash.device_hasher(has_nulls)},
-                                          {},
-                                          {},
-                                          cudf::detail::cuco_allocator<char>{rmm::mr::polymorphic_allocator<char>{}, stream},
-                                          stream.value()};
+      auto set = hash_set_type<RowHasher>{
+        num_rows,
+        0.5,  // desired load factor
+        cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
+        d_equal,
+        {row_hash.device_hasher(has_nulls)},
+        {},
+        {},
+        cudf::detail::cuco_allocator<char>{rmm::mr::polymorphic_allocator<char>{}, stream},
+        stream.value()};
 
       auto const iter = thrust::counting_iterator<cudf::size_type>{0};
       set.insert_async(iter, iter + num_rows, stream.value());
@@ -121,16 +122,17 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
       return output_indices;
     }
 
-    auto map = hash_map_type<RowHasher>{num_rows,
-                                        0.5,  // desired load factor
-                                        cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
-                                        cuco::empty_value{reduction_init_value(keep)},
-                                        d_equal,
-                                        {row_hash.device_hasher(has_nulls)},
-                                        {},
-                                        {},
-                                        cudf::detail::cuco_allocator<char>{rmm::mr::polymorphic_allocator<char>{}, stream},
-                                        stream.value()};
+    auto map = hash_map_type<RowHasher>{
+      num_rows,
+      0.5,  // desired load factor
+      cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL},
+      cuco::empty_value{reduction_init_value(keep)},
+      d_equal,
+      {row_hash.device_hasher(has_nulls)},
+      {},
+      {},
+      cudf::detail::cuco_allocator<char>{rmm::mr::polymorphic_allocator<char>{}, stream},
+      stream.value()};
     return reduce_by_row(map, num_rows, keep, stream, mr);
   };
 

From 78561d103ed9ce72fa74041a5346cd42e5381ffb Mon Sep 17 00:00:00 2001
From: Srinivas Yadav Singanaboina <ssinganaboin@nvidia.com>
Date: Thu, 8 Aug 2024 18:45:15 +0000
Subject: [PATCH 8/8] fix missing template parameter for allocator

---
 cpp/src/stream_compaction/distinct_helpers.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/src/stream_compaction/distinct_helpers.hpp b/cpp/src/stream_compaction/distinct_helpers.hpp
index 765b3b206c21..de5933f67769 100644
--- a/cpp/src/stream_compaction/distinct_helpers.hpp
+++ b/cpp/src/stream_compaction/distinct_helpers.hpp
@@ -99,7 +99,7 @@ using hash_map_type =
                                         cudf::experimental::row::hash::device_row_hasher<
                                           cudf::hashing::detail::default_hash,
                                           cudf::nullate::DYNAMIC>>,
-                   cudf::detail::cuco_allocator,
+                   cudf::detail::cuco_allocator<char>,
                    cuco::storage<1>>;
 
 /**