Add memory resource parameter to cudf::distinct_hash_join constructors

cpp/include/cudf/detail/join/distinct_hash_join.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -111,20 +111,24 @@ class distinct_hash_join {
    * @param right The right table, from which the hash table is built
    * @param compare_nulls Controls whether null join-key values should match or not.
    * @param stream CUDA stream used for device memory operations and kernel launches.
+   * @param mr Device memory resource used to allocate the internal hash table.
    */
   distinct_hash_join(cudf::table_view const& right,
                      cudf::null_equality compare_nulls,
-                     rmm::cuda_stream_view stream);
+                     rmm::cuda_stream_view stream,
+                     cuda::mr::any_resource<cuda::mr::device_accessible> mr);
 
   /**
-   * @copydoc distinct_hash_join(cudf::table_view const&, null_equality, rmm::cuda_stream_view)
+   * @copydoc distinct_hash_join(cudf::table_view const&, null_equality, rmm::cuda_stream_view,
+   * cuda::mr::any_resource<cuda::mr::device_accessible>)
    *
    * @param load_factor The hash table occupancy ratio in (0,1]. A value of 0.5 means 50% occupancy.
    */
   distinct_hash_join(cudf::table_view const& right,
                      cudf::null_equality compare_nulls,
                      double load_factor,
-                     rmm::cuda_stream_view stream);
+                     rmm::cuda_stream_view stream,
+                     cuda::mr::any_resource<cuda::mr::device_accessible> mr);
 
   /**
    * @copydoc cudf::distinct_hash_join::inner_join

cpp/include/cudf/join/distinct_hash_join.hpp

@@ -68,11 +68,14 @@ class distinct_hash_join {
    * in range (0,1]. For example, 0.5 indicates a target of 50% occupancy. Note that the actual
    * occupancy achieved may be slightly lower than the specified value.
    * @param stream CUDA stream used for device memory operations and kernel launches
+   * @param mr Device memory resource used to allocate the internal hash table
    */
   distinct_hash_join(cudf::table_view const& right,
                      null_equality compare_nulls  = null_equality::EQUAL,
                      double load_factor           = 0.5,
-                     rmm::cuda_stream_view stream = cudf::get_default_stream());
+                     rmm::cuda_stream_view stream = cudf::get_default_stream(),
+                     cuda::mr::any_resource<cuda::mr::device_accessible> mr =
+                       cudf::get_current_device_resource_ref());
 
   /**
    * @brief Returns the row indices that can be used to construct the result of performing

cpp/src/join/distinct_hash_join.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #include "join_common_utils.cuh"
@@ -151,15 +151,17 @@ void find_matches_in_hash_table(HashTableType const& hash_table,
 
 distinct_hash_join::distinct_hash_join(cudf::table_view const& right,
                                        cudf::null_equality compare_nulls,
-                                       rmm::cuda_stream_view stream)
-  : distinct_hash_join{right, compare_nulls, CUCO_DESIRED_LOAD_FACTOR, stream}
+                                       rmm::cuda_stream_view stream,
+                                       cuda::mr::any_resource<cuda::mr::device_accessible> mr)
+  : distinct_hash_join{right, compare_nulls, CUCO_DESIRED_LOAD_FACTOR, stream, std::move(mr)}
 {
 }
 
 distinct_hash_join::distinct_hash_join(cudf::table_view const& right,
                                        cudf::null_equality compare_nulls,
                                        double load_factor,
-                                       rmm::cuda_stream_view stream)
+                                       rmm::cuda_stream_view stream,
+                                       cuda::mr::any_resource<cuda::mr::device_accessible> mr)
   : _has_nested_columns{cudf::has_nested_columns(right)},
     _nulls_equal{compare_nulls},
     _right{right},
@@ -172,7 +174,7 @@ distinct_hash_join::distinct_hash_join(cudf::table_view const& right,
                 {},
                 cuco::thread_scope_device,
                 cuco_storage_type{},
-                rmm::mr::polymorphic_allocator<char>{},
+                rmm::mr::polymorphic_allocator<char>{std::move(mr)},
                 stream.value()}
 {
   CUDF_FUNC_RANGE();
@@ -408,8 +410,9 @@ distinct_hash_join::~distinct_hash_join() = default;
 distinct_hash_join::distinct_hash_join(cudf::table_view const& right,
                                        null_equality compare_nulls,
                                        double load_factor,
-                                       rmm::cuda_stream_view stream)
-  : _impl{std::make_unique<impl_type>(right, compare_nulls, load_factor, stream)}
+                                       rmm::cuda_stream_view stream,
+                                       cuda::mr::any_resource<cuda::mr::device_accessible> mr)
+  : _impl{std::make_unique<impl_type>(right, compare_nulls, load_factor, stream, std::move(mr))}
 {
 }
 

cpp/tests/join/distinct_join_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -16,6 +16,8 @@
 #include <cudf/types.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
+#include <rmm/mr/statistics_resource_adaptor.hpp>
+
 #include <cuco/utility/error.hpp>
 
 #include <limits>
@@ -91,6 +93,54 @@ TEST_F(DistinctJoinTest, IntegerInnerJoin)
   this->compare_to_reference(right_table, left_table, result, cudf::table_view{{gold->view()}});
 }
 
+TEST_F(DistinctJoinTest, MemoryResource)
+{
+  column_wrapper<int32_t> col0_0{{1, 2, 3, 4, 5}};
+  strcol_wrapper col0_1({"s0", "s0", "s3", "s4", "s5"});
+  column_wrapper<int32_t> col0_2{{9, 9, 9, 9, 9}};
+
+  column_wrapper<int32_t> col1_0{{1, 2, 3, 4, 9}};
+  strcol_wrapper col1_1({"s0", "s0", "s0", "s4", "s4"});
+  column_wrapper<int32_t> col1_2{{9, 9, 9, 0, 9}};
+
+  CVector cols0, cols1;
+  cols0.push_back(col0_0.release());
+  cols0.push_back(col0_1.release());
+  cols0.push_back(col0_2.release());
+  cols1.push_back(col1_0.release());
+  cols1.push_back(col1_1.release());
+  cols1.push_back(col1_2.release());
+
+  Table right(std::move(cols0));
+  Table left(std::move(cols1));
+
+  auto mr = rmm::mr::statistics_resource_adaptor(cudf::get_current_device_resource_ref());
+
+  auto distinct_join = cudf::distinct_hash_join{
+    right.view(), cudf::null_equality::EQUAL, 0.5, cudf::get_default_stream(), mr};
+
+  EXPECT_GT(mr.get_bytes_counter().peak, 0);
+
+  auto result = distinct_join.inner_join(left.view());
+
+  column_wrapper<int32_t> col_gold_0{{1, 2}};
+  strcol_wrapper col_gold_1({"s0", "s0"});
+  column_wrapper<int32_t> col_gold_2{{9, 9}};
+  column_wrapper<int32_t> col_gold_3{{1, 2}};
+  strcol_wrapper col_gold_4({"s0", "s0"});
+  column_wrapper<int32_t> col_gold_5{{9, 9}};
+  CVector cols_gold;
+  cols_gold.push_back(col_gold_0.release());
+  cols_gold.push_back(col_gold_1.release());
+  cols_gold.push_back(col_gold_2.release());
+  cols_gold.push_back(col_gold_3.release());
+  cols_gold.push_back(col_gold_4.release());
+  cols_gold.push_back(col_gold_5.release());
+  Table gold(std::move(cols_gold));
+
+  this->compare_to_reference(right.view(), left.view(), result, gold.view());
+}
+
 TEST_F(DistinctJoinTest, InnerJoinNoNulls)
 {
   column_wrapper<int32_t> col0_0{{1, 2, 3, 4, 5}};