diff --git a/cpp/include/cudf/detail/join/distinct_hash_join.cuh b/cpp/include/cudf/detail/join/distinct_hash_join.cuh index 3d0b1852e896..9aedb49749f4 100644 --- a/cpp/include/cudf/detail/join/distinct_hash_join.cuh +++ b/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 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) * * @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 mr); /** * @copydoc cudf::distinct_hash_join::inner_join diff --git a/cpp/include/cudf/join/distinct_hash_join.hpp b/cpp/include/cudf/join/distinct_hash_join.hpp index 6dc2589f41dc..326d9cbb3626 100644 --- a/cpp/include/cudf/join/distinct_hash_join.hpp +++ b/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 mr = + cudf::get_current_device_resource_ref()); /** * @brief Returns the row indices that can be used to construct the result of performing diff --git a/cpp/src/join/distinct_hash_join.cu b/cpp/src/join/distinct_hash_join.cu index 718f96bf29db..23697c611df1 100644 --- a/cpp/src/join/distinct_hash_join.cu +++ b/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 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 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{}, + rmm::mr::polymorphic_allocator{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(right, compare_nulls, load_factor, stream)} + rmm::cuda_stream_view stream, + cuda::mr::any_resource mr) + : _impl{std::make_unique(right, compare_nulls, load_factor, stream, std::move(mr))} { } diff --git a/cpp/tests/join/distinct_join_tests.cpp b/cpp/tests/join/distinct_join_tests.cpp index 744caddeb493..6e251c3d9413 100644 --- a/cpp/tests/join/distinct_join_tests.cpp +++ b/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 #include +#include + #include #include @@ -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 col0_0{{1, 2, 3, 4, 5}}; + strcol_wrapper col0_1({"s0", "s0", "s3", "s4", "s5"}); + column_wrapper col0_2{{9, 9, 9, 9, 9}}; + + column_wrapper col1_0{{1, 2, 3, 4, 9}}; + strcol_wrapper col1_1({"s0", "s0", "s0", "s4", "s4"}); + column_wrapper 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 col_gold_0{{1, 2}}; + strcol_wrapper col_gold_1({"s0", "s0"}); + column_wrapper col_gold_2{{9, 9}}; + column_wrapper col_gold_3{{1, 2}}; + strcol_wrapper col_gold_4({"s0", "s0"}); + column_wrapper 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 col0_0{{1, 2, 3, 4, 5}};