Add memory resource parameter to cudf::hash_join constructors

cpp/include/cudf/detail/join/hash_join.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -59,22 +59,27 @@ class hash_join {
    *        any `left` table that will be used later for join.
    * @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, which must remain
+   *           valid for the lifetime of this object.
    */
   hash_join(cudf::table_view const& right,
             bool has_nulls,
             cudf::null_equality compare_nulls,
-            rmm::cuda_stream_view stream);
+            rmm::cuda_stream_view stream,
+            rmm::device_async_resource_ref mr);
 
   /**
-   * @copydoc hash_join(cudf::table_view const&, bool, null_equality, rmm::cuda_stream_view)
+   * @copydoc hash_join(cudf::table_view const&, bool, null_equality, rmm::cuda_stream_view,
+   * rmm::device_async_resource_ref)
    *
    * @param load_factor The hash table occupancy ratio in (0,1]. A value of 0.5 means 50% occupancy.
    */
   hash_join(cudf::table_view const& right,
             bool has_nulls,
             cudf::null_equality compare_nulls,
             double load_factor,
-            rmm::cuda_stream_view stream);
+            rmm::cuda_stream_view stream,
+            rmm::device_async_resource_ref mr);
 
   /**
    * @copydoc cudf::hash_join::inner_join

cpp/include/cudf/join/hash_join.hpp

@@ -87,26 +87,39 @@ class 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, which must remain
+   *           valid for the lifetime of this object
    */
   hash_join(cudf::table_view const& right,
             null_equality compare_nulls,
-            rmm::cuda_stream_view stream = cudf::get_default_stream());
+            rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+            rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
   /**
-   * @copydoc hash_join(cudf::table_view const&, null_equality, rmm::cuda_stream_view)
+   * @brief Construct a hash join object for subsequent probe calls.
+   *
+   * @note The `hash_join` object must not outlive the table viewed by `right`, else behavior is
+   * undefined.
    *
+   * @throws std::invalid_argument if the right table has no columns
    * @throws std::invalid_argument if load_factor is not greater than 0 and less than or equal to 1
    *
+   * @param right The right table, from which the hash table is built
    * @param has_nulls Flag to indicate if there exists any nulls in the `right` table or
    *                  any `left` table that will be used later for join
+   * @param compare_nulls Controls whether null join-key values should match or not
    * @param load_factor The hash table occupancy ratio in (0,1]. A value of 0.5 means 50% desired
    * occupancy.
+   * @param stream CUDA stream used for device memory operations and kernel launches
+   * @param mr Device memory resource used to allocate the internal hash table, which must remain
+   *           valid for the lifetime of this object
    */
   hash_join(cudf::table_view const& right,
             nullable_join has_nulls,
             null_equality compare_nulls,
             double load_factor,
-            rmm::cuda_stream_view stream = cudf::get_default_stream());
+            rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+            rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
   /**
    * Returns the row indices that can be used to construct the result of performing

cpp/src/join/hash_join/hash_join.cu

@@ -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
  */
 
@@ -102,8 +102,9 @@ template <typename Hasher>
 hash_join<Hasher>::hash_join(cudf::table_view const& right,
                              bool has_nulls,
                              cudf::null_equality compare_nulls,
-                             rmm::cuda_stream_view stream)
-  : hash_join{right, has_nulls, compare_nulls, CUCO_DESIRED_LOAD_FACTOR, stream}
+                             rmm::cuda_stream_view stream,
+                             rmm::device_async_resource_ref mr)
+  : hash_join{right, has_nulls, compare_nulls, CUCO_DESIRED_LOAD_FACTOR, stream, mr}
 {
 }
 
@@ -112,7 +113,8 @@ hash_join<Hasher>::hash_join(cudf::table_view const& right,
                              bool has_nulls,
                              cudf::null_equality compare_nulls,
                              double load_factor,
-                             rmm::cuda_stream_view stream)
+                             rmm::cuda_stream_view stream,
+                             rmm::device_async_resource_ref mr)
   : _has_nulls(has_nulls),
     _is_empty{right.num_rows() == 0},
     _nulls_equal{compare_nulls},
@@ -124,7 +126,7 @@ hash_join<Hasher>::hash_join(cudf::table_view const& right,
       {},
       {},
       {},
-      rmm::mr::polymorphic_allocator<char>{},
+      rmm::mr::polymorphic_allocator<char>{mr},
       stream.value()}})},
     _right{right},
     _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(_right, stream)}
@@ -151,13 +153,15 @@ hash_join<Hasher>::hash_join(cudf::table_view const& right,
 template hash_join<hash_join_hasher>::hash_join(cudf::table_view const& right,
                                                 bool has_nulls,
                                                 cudf::null_equality compare_nulls,
-                                                rmm::cuda_stream_view stream);
+                                                rmm::cuda_stream_view stream,
+                                                rmm::device_async_resource_ref mr);
 
 template hash_join<hash_join_hasher>::hash_join(cudf::table_view const& right,
                                                 bool has_nulls,
                                                 cudf::null_equality compare_nulls,
                                                 double load_factor,
-                                                rmm::cuda_stream_view stream);
+                                                rmm::cuda_stream_view stream,
+                                                rmm::device_async_resource_ref mr);
 
 template <typename Hasher>
 hash_join<Hasher>::~hash_join() = default;
@@ -172,19 +176,21 @@ hash_join::~hash_join() = default;
 
 hash_join::hash_join(cudf::table_view const& right,
                      null_equality compare_nulls,
-                     rmm::cuda_stream_view stream)
+                     rmm::cuda_stream_view stream,
+                     rmm::device_async_resource_ref mr)
   : hash_join(
-      right, nullable_join::YES, compare_nulls, cudf::detail::CUCO_DESIRED_LOAD_FACTOR, stream)
+      right, nullable_join::YES, compare_nulls, cudf::detail::CUCO_DESIRED_LOAD_FACTOR, stream, mr)
 {
 }
 
 hash_join::hash_join(cudf::table_view const& right,
                      nullable_join has_nulls,
                      null_equality compare_nulls,
                      double load_factor,
-                     rmm::cuda_stream_view stream)
+                     rmm::cuda_stream_view stream,
+                     rmm::device_async_resource_ref mr)
   : _impl{std::make_unique<impl_type const>(
-      right, has_nulls == nullable_join::YES, compare_nulls, load_factor, stream)}
+      right, has_nulls == nullable_join::YES, compare_nulls, load_factor, stream, mr)}
 {
 }
 

cpp/tests/join/join_tests.cpp

@@ -28,6 +28,7 @@
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/mr/statistics_resource_adaptor.hpp>
 
 #include <cuco/utility/error.hpp>
 
@@ -2313,6 +2314,29 @@ TEST_F(JoinTest, HashJoinLargeOutputSize)
   EXPECT_EQ(col_size * col_size, output_size);
 }
 
+TEST_F(JoinTest, HashJoinMemoryResource)
+{
+  CVector cols0;
+  cols0.emplace_back(column_wrapper<int32_t>{{3, 1, 2, 0, 2}}.release());
+  Table t0(std::move(cols0));
+
+  CVector cols1;
+  cols1.emplace_back(column_wrapper<int32_t>{{2, 2, 0, 4, 3}}.release());
+  Table t1(std::move(cols1));
+
+  auto mr = rmm::mr::statistics_resource_adaptor(cudf::get_current_device_resource_ref());
+
+  cudf::hash_join hash_join(t1, cudf::null_equality::EQUAL, cudf::get_default_stream(), mr);
+
+  EXPECT_GT(mr.get_bytes_counter().peak, 0);
+
+  auto result = hash_join.inner_join(t0);
+  column_wrapper<int32_t> col_gold_0{{0, 2, 2, 3, 4, 4}};
+  column_wrapper<int32_t> col_gold_1{{4, 0, 1, 2, 0, 1}};
+  auto const [sorted_gold, sorted_result] = gather_maps_as_tables(col_gold_0, col_gold_1, result);
+  CUDF_TEST_EXPECT_TABLES_EQUIVALENT(*sorted_gold, *sorted_result);
+}
+
 TEST_F(JoinTest, HashJoinInnerMatchContext)
 {
   // Test inner join match context functionality with multiple matches and nulls