[cudf] Thread memory resources through equality/hash row preprocessing

cpp/include/cudf/detail/row_operator/equality.cuh

@@ -407,9 +407,12 @@ class self_comparator {
    * @param t The table to compare
    * @param stream The stream to construct this object on. Not the stream that will be used for
    * comparisons using this object.
+   * @param temp_mr Device memory resource used for temporary allocations
    */
-  self_comparator(table_view const& t, rmm::cuda_stream_view stream)
-    : d_t(preprocessed_table::create(t, stream))
+  self_comparator(table_view const& t,
+                  rmm::cuda_stream_view stream,
+                  rmm::device_async_resource_ref temp_mr)
+    : d_t(preprocessed_table::create(t, stream, temp_mr))
   {
   }
 
@@ -515,10 +518,12 @@ class two_table_comparator {
    * @param right The right table to compare.
    * @param stream The stream to construct this object on. Not the stream that will be used for
    * comparisons using this object.
+   * @param temp_mr Device memory resource used for temporary allocations
    */
   two_table_comparator(table_view const& left,
                        table_view const& right,
-                       rmm::cuda_stream_view stream);
+                       rmm::cuda_stream_view stream,
+                       rmm::device_async_resource_ref temp_mr);
 
   /**
    * @brief Construct an owning object for performing equality comparisons between two rows from two

cpp/include/cudf/detail/row_operator/hashing.cuh

@@ -240,9 +240,12 @@ class row_hasher {
    * @param t The table containing rows to hash
    * @param stream The stream to construct this object on. Not the stream that will be used for
    * comparisons using this object.
+   * @param temp_mr Device memory resource used for temporary allocations
    */
-  row_hasher(table_view const& t, rmm::cuda_stream_view stream)
-    : d_t(preprocessed_table::create(t, stream))
+  row_hasher(table_view const& t,
+             rmm::cuda_stream_view stream,
+             rmm::device_async_resource_ref temp_mr)
+    : d_t(preprocessed_table::create(t, stream, temp_mr))
   {
   }
 

cpp/include/cudf/detail/row_operator/preprocessed_table.cuh

@@ -10,6 +10,7 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <memory>
 #include <vector>
@@ -47,10 +48,12 @@ struct preprocessed_table {
    *
    * @param table The table to preprocess
    * @param stream The cuda stream to use while preprocessing.
+   * @param temp_mr Device memory resource used for temporary allocations
    * @return A preprocessed table as shared pointer
    */
   static std::shared_ptr<preprocessed_table> create(table_view const& table,
-                                                    rmm::cuda_stream_view stream);
+                                                    rmm::cuda_stream_view stream,
+                                                    rmm::device_async_resource_ref temp_mr);
 
   /**
    * @brief Implicit conversion operator to a `table_device_view` of the preprocessed table.

cpp/src/binaryop/compiled/struct_binary_ops.cuh

@@ -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
  */
 
@@ -146,7 +146,8 @@ void apply_struct_equality_op(mutable_column_view& out,
 
   auto tlhs             = table_view{{lhs}};
   auto trhs             = table_view{{rhs}};
-  auto table_comparator = cudf::detail::row::equality::two_table_comparator{tlhs, trhs, stream};
+  auto table_comparator = cudf::detail::row::equality::two_table_comparator{
+    tlhs, trhs, stream, cudf::get_current_device_resource_ref()};
 
   auto outd = column_device_view::create(out, stream);
   auto optional_iter =

cpp/src/dictionary/decode.cu

@@ -41,8 +41,9 @@ std::unique_ptr<column> decode(dictionary_column_view const& source,
   if (source.is_empty()) return make_empty_column(type_id::EMPTY);
 
   // annotated indices include the offset, size and bitmask from it's parent
-  auto const indices       = source.get_indices_annotated();
-  auto const d_indices     = column_device_view::create(indices, stream);
+  auto const indices = source.get_indices_annotated();
+  auto const d_indices =
+    column_device_view::create(indices, stream, cudf::get_current_device_resource_ref());
   auto const d_iterator    = cudf::detail::indexalator_factory::make_input_iterator(indices);
   auto const indices_begin = cudf::detail::make_counting_transform_iterator(
     0, indices_handler_fn{d_iterator, *d_indices, source.keys().size()});

cpp/src/dictionary/detail/concatenate.cu

@@ -187,8 +187,9 @@ std::unique_ptr<column> concatenate(host_span<column_view const> columns,
     cudf::detail::row::hash::device_row_hasher<cudf::hashing::detail::default_hash,
                                                cudf::nullate::NO>>;
   auto const tv         = cudf::table_view({all_keys->view()});
-  auto const row_hash   = cudf::detail::row::hash::row_hasher(tv, stream);
-  auto const row_equal  = cudf::detail::row::equality::self_comparator(tv, stream);
+  auto const temp_mr    = cudf::get_current_device_resource_ref();
+  auto const row_hash   = cudf::detail::row::hash::row_hasher(tv, stream, temp_mr);
+  auto const row_equal  = cudf::detail::row::equality::self_comparator(tv, stream, temp_mr);
   auto const comparator = cudf::detail::row::equality::nan_equal_physical_equality_comparator{};
   auto const d_equal =
     row_equal.equal_to<false>(cudf::nullate::NO{}, null_equality::EQUAL, comparator);

cpp/src/dictionary/encode.cu

@@ -82,36 +82,35 @@ std::unique_ptr<column> encode(column_view const& input,
 
   auto const has_nulls  = nullate::DYNAMIC{input.has_nulls()};
   auto const tv         = cudf::table_view({input});
-  auto const row_hash   = cudf::detail::row::hash::row_hasher(tv, stream);
-  auto const row_equal  = cudf::detail::row::equality::self_comparator(tv, stream);
+  auto const temp_mr    = cudf::get_current_device_resource_ref();
+  auto const row_hash   = cudf::detail::row::hash::row_hasher(tv, stream, temp_mr);
+  auto const row_equal  = cudf::detail::row::equality::self_comparator(tv, stream, temp_mr);
   auto const comparator = cudf::detail::row::equality::nan_equal_physical_equality_comparator{};
   auto const d_equal    = row_equal.equal_to<false>(has_nulls, null_equality::EQUAL, comparator);
   auto const empty_key  = cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL};
   auto probe            = encode_probe_t{row_hash.device_hasher(has_nulls)};
-  auto allocator        = rmm::mr::polymorphic_allocator<char>{};
+  auto allocator        = rmm::mr::polymorphic_allocator<char>{temp_mr};
   auto set =
     cuco::static_set{input.size(), 0.5, empty_key, d_equal, probe, {}, {}, allocator, stream.get()};
   auto set_ref    = set.ref(cuco::insert_and_find);
   using set_ref_t = decltype(set_ref);
 
   // build a static_set of the input values
   // and keep track of the indices of the unique values
-  auto d_indices = rmm::device_uvector<size_type>(input.size(), stream);
-  auto d_input   = column_device_view::create(input, stream);
-  thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+  auto d_indices = rmm::device_uvector<size_type>(input.size(), stream, temp_mr);
+  auto d_input   = column_device_view::create(input, stream, temp_mr);
+  thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
                     cuda::counting_iterator<size_type>{0},
                     cuda::counting_iterator<size_type>{input.size()},
                     d_indices.begin(),
                     encode_fn{set_ref, *d_input});
 
-  auto keys_indices = rmm::device_uvector<size_type>(input.size(), stream);
+  auto keys_indices = rmm::device_uvector<size_type>(input.size(), stream, temp_mr);
   auto keys_end     = set.retrieve_all(keys_indices.begin(), stream.get());
   keys_indices.resize(cuda::std::distance(keys_indices.begin(), keys_end), stream);
 
   // sort the keys_indices so we can use lower-bound on them
-  thrust::sort(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
-               keys_indices.begin(),
-               keys_indices.end());
+  thrust::sort(rmm::exec_policy_nosync(stream, temp_mr), keys_indices.begin(), keys_indices.end());
 
   // use keys_indices to retrieve the keys
   auto const oob_policy   = cudf::out_of_bounds_policy::DONT_CHECK;
@@ -124,7 +123,7 @@ std::unique_ptr<column> encode(column_view const& input,
   // call lower-bound with keys_indices and d_indices to get the output indices_column
   auto d_result =
     cudf::detail::indexalator_factory::make_output_iterator(indices_column->mutable_view());
-  thrust::lower_bound(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+  thrust::lower_bound(rmm::exec_policy_nosync(stream, temp_mr),
                       keys_indices.begin(),
                       keys_indices.end(),
                       d_indices.begin(),

cpp/src/dictionary/match_keys.cu

@@ -50,8 +50,9 @@ struct unique_keys_dispatch_fn {
 
     auto const has_nulls  = nullate::DYNAMIC{false};
     auto const keys_tv    = table_view({all_keys});
-    auto const row_hash   = cudf::detail::row::hash::row_hasher(keys_tv, stream);
-    auto const row_equal  = cudf::detail::row::equality::self_comparator(keys_tv, stream);
+    auto const temp_mr    = cudf::get_current_device_resource_ref();
+    auto const row_hash   = cudf::detail::row::hash::row_hasher(keys_tv, stream, temp_mr);
+    auto const row_equal  = cudf::detail::row::equality::self_comparator(keys_tv, stream, temp_mr);
     auto const comparator = cudf::detail::row::equality::nan_equal_physical_equality_comparator{};
     auto const d_equal    = row_equal.equal_to<false>(has_nulls, null_equality::EQUAL, comparator);
     auto const empty_key  = cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL};

cpp/src/groupby/hash/groupby.cu

@@ -42,10 +42,11 @@ std::unique_ptr<table> dispatch_groupby(table_view const& keys,
   auto const has_null             = nullate::DYNAMIC{cudf::has_nested_nulls(keys)};
   auto const skip_rows_with_nulls = keys_have_nulls and include_null_keys == null_policy::EXCLUDE;
 
-  auto preprocessed_keys = cudf::detail::row::hash::preprocessed_table::create(keys, stream);
-  auto const comparator  = cudf::detail::row::equality::self_comparator{preprocessed_keys};
-  auto const row_hash    = cudf::detail::row::hash::row_hasher{std::move(preprocessed_keys)};
-  auto const d_row_hash  = row_hash.device_hasher(has_null);
+  auto preprocessed_keys = cudf::detail::row::hash::preprocessed_table::create(
+    keys, stream, cudf::get_current_device_resource_ref());
+  auto const comparator = cudf::detail::row::equality::self_comparator{preprocessed_keys};
+  auto const row_hash   = cudf::detail::row::hash::row_hasher{std::move(preprocessed_keys)};
+  auto const d_row_hash = row_hash.device_hasher(has_null);
 
   if (cudf::detail::has_nested_columns(keys)) {
     auto const d_row_equal = comparator.equal_to<true>(has_null, null_keys_are_equal);

cpp/src/groupby/sort/group_nunique.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -79,7 +79,8 @@ std::unique_ptr<column> group_nunique(column_view const& values,
   if (num_groups == 0) { return result; }
 
   auto const values_view = table_view{{values}};
-  auto const comparator  = cudf::detail::row::equality::self_comparator{values_view, stream};
+  auto const comparator  = cudf::detail::row::equality::self_comparator{
+    values_view, stream, cudf::get_current_device_resource_ref()};
 
   auto const d_values_view = column_device_view::create(values, stream);
 

cpp/src/groupby/sort/group_rank_scan.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -94,7 +94,9 @@ std::unique_ptr<column> rank_generator(column_view const& grouped_values,
                                        rmm::device_async_resource_ref mr)
 {
   auto const grouped_values_view = table_view{{grouped_values}};
-  auto const comparator = cudf::detail::row::equality::self_comparator{grouped_values_view, stream};
+  auto const temp_mr             = cudf::get_current_device_resource_ref();
+  auto const comparator =
+    cudf::detail::row::equality::self_comparator{grouped_values_view, stream, temp_mr};
 
   auto ranks = make_fixed_width_column(
     data_type{type_to_id<size_type>()}, grouped_values.size(), mask_state::UNALLOCATED, stream, mr);
@@ -104,7 +106,7 @@ std::unique_ptr<column> rank_generator(column_view const& grouped_values,
     auto const permuted_equal =
       permuted_row_equality_comparator(d_equal, value_order.begin<size_type>());
 
-    thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+    thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
                       cuda::counting_iterator<size_type>(0),
                       cuda::counting_iterator<size_type>(grouped_values.size()),
                       mutable_ranks.begin<size_type>(),
@@ -130,14 +132,13 @@ std::unique_ptr<column> rank_generator(column_view const& grouped_values,
                              cuda::std::reverse_iterator(mutable_ranks.end<size_type>())};
     }
   }();
-  thrust::inclusive_scan_by_key(
-    rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
-    group_labels_begin,
-    group_labels_begin + group_labels.size(),
-    mutable_rank_begin,
-    mutable_rank_begin,
-    cuda::std::equal_to{},
-    scan_op);
+  thrust::inclusive_scan_by_key(rmm::exec_policy_nosync(stream, temp_mr),
+                                group_labels_begin,
+                                group_labels_begin + group_labels.size(),
+                                mutable_rank_begin,
+                                mutable_rank_begin,
+                                cuda::std::equal_to{},
+                                scan_op);
   return ranks;
 }
 }  // namespace

cpp/src/groupby/sort/sort_helper_group_offsets.cuh

@@ -38,22 +38,19 @@ size_type compute_group_offsets(table_view const& keys,
                                 rmm::device_uvector<size_type>& group_offsets,
                                 rmm::cuda_stream_view stream)
 {
-  auto const comparator  = cudf::detail::row::equality::self_comparator{keys, stream};
+  auto const temp_mr     = cudf::get_current_device_resource_ref();
+  auto const comparator  = cudf::detail::row::equality::self_comparator{keys, stream, temp_mr};
   auto const d_key_equal = comparator.equal_to<HasNested>(
     cudf::nullate::DYNAMIC{cudf::has_nested_nulls(keys)}, null_equality::EQUAL);
   // Using a temporary buffer for intermediate transform results from the iterator containing
   // the comparator speeds up compile-time significantly without much degradation in
   // runtime performance over using the comparator directly in thrust::unique_copy.
-  auto result       = rmm::device_uvector<bool>(size, stream);
+  auto result       = rmm::device_uvector<bool>(size, stream, temp_mr);
   auto const itr    = cuda::counting_iterator<size_type>{0};
   auto const row_eq = permuted_row_equality_comparator(d_key_equal, sorted_order);
   auto const ufn    = cudf::detail::unique_copy_fn<decltype(itr), decltype(row_eq)>{
     itr, duplicate_keep_option::KEEP_FIRST, row_eq, size - 1};
-  thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
-                    itr,
-                    itr + size,
-                    result.begin(),
-                    ufn);
+  thrust::transform(rmm::exec_policy_nosync(stream, temp_mr), itr, itr + size, result.begin(), ufn);
   auto const result_end = cudf::detail::copy_if(
     itr, itr + size, result.begin(), group_offsets.begin(), cuda::std::identity{}, stream);
   return cuda::std::distance(group_offsets.begin(), result_end);

cpp/src/groupby/streaming_groupby/insert.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -34,7 +34,8 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_
   auto const has_null   = cudf::nullate::DYNAMIC{_has_nullable_keys};
 
   // Preprocess batch for row operators.
-  auto preprocessed_batch = cudf::detail::row::hash::preprocessed_table::create(batch_keys, stream);
+  auto preprocessed_batch =
+    cudf::detail::row::hash::preprocessed_table::create(batch_keys, stream, temp_mr);
   auto const batch_hasher_obj = cudf::detail::row::hash::row_hasher{preprocessed_batch};
   auto const d_batch_hash     = batch_hasher_obj.device_hasher(has_null);
 
@@ -110,7 +111,7 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_
                                           temp_mr);
 
     auto preprocessed_compacted =
-      cudf::detail::row::hash::preprocessed_table::create(compacted->view(), stream);
+      cudf::detail::row::hash::preprocessed_table::create(compacted->view(), stream, temp_mr);
 
     // Store the compacted batch.
     auto const new_batch_id    = static_cast<size_type>(_compacted_batches.size());

cpp/src/hash/murmurhash3_x86_32.cu

@@ -15,6 +15,7 @@
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/resource_ref.hpp>
@@ -66,8 +67,8 @@ std::unique_ptr<column> murmurhash3_x86_32(table_view const& input,
                                            rmm::cuda_stream_view stream,
                                            rmm::device_async_resource_ref mr)
 {
-  auto const preprocessed_input =
-    cudf::detail::row::hash::preprocessed_table::create(input, stream);
+  auto const preprocessed_input = cudf::detail::row::hash::preprocessed_table::create(
+    input, stream, cudf::get_current_device_resource_ref());
   return murmurhash3_x86_32_impl(
     preprocessed_input, input.num_rows(), seed, nullate::DYNAMIC{has_nulls(input)}, stream, mr);
 }

cpp/src/hash/xxhash_32.cu

@@ -30,9 +30,10 @@ std::unique_ptr<column> xxhash_32(table_view const& input,
 
   if (input.num_rows() == 0) { return output; }
 
-  bool const nullable   = has_nulls(input);
-  auto const row_hasher = cudf::detail::row::hash::row_hasher(input, stream);
-  auto output_view      = output->mutable_view();
+  bool const nullable = has_nulls(input);
+  auto const row_hasher =
+    cudf::detail::row::hash::row_hasher(input, stream, cudf::get_current_device_resource_ref());
+  auto output_view = output->mutable_view();
 
   // Compute the hash value for each row
   auto const output_begin = output_view.begin<hash_value_type>();

cpp/src/hash/xxhash_64.cu

@@ -32,9 +32,10 @@ std::unique_ptr<column> xxhash_64(table_view const& input,
 
   if (input.num_rows() == 0) { return output; }
 
-  bool const nullable   = has_nulls(input);
-  auto const row_hasher = cudf::detail::row::hash::row_hasher(input, stream);
-  auto output_view      = output->mutable_view();
+  bool const nullable = has_nulls(input);
+  auto const row_hasher =
+    cudf::detail::row::hash::row_hasher(input, stream, cudf::get_current_device_resource_ref());
+  auto output_view = output->mutable_view();
 
   // Compute the hash value for each row
   auto const output_begin = output_view.begin<hash_value_type>();