diff --git a/cpp/include/cudf/detail/join/hash_join.hpp b/cpp/include/cudf/detail/join/hash_join.hpp index 81ab53fd5f3f..b3f5f8f5a39a 100644 --- a/cpp/include/cudf/detail/join/hash_join.hpp +++ b/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,14 +59,17 @@ 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. */ 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, + cuda::mr::any_resource 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, + * cuda::mr::any_resource) * * @param load_factor The hash table occupancy ratio in (0,1]. A value of 0.5 means 50% occupancy. */ @@ -74,7 +77,8 @@ class hash_join { bool has_nulls, 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::hash_join::inner_join diff --git a/cpp/include/cudf/join/hash_join.hpp b/cpp/include/cudf/join/hash_join.hpp index c2577f41b6ff..73c1944b54fe 100644 --- a/cpp/include/cudf/join/hash_join.hpp +++ b/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 */ 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(), + cuda::mr::any_resource 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 */ 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(), + cuda::mr::any_resource mr = + cudf::get_current_device_resource_ref()); /** * Returns the row indices that can be used to construct the result of performing diff --git a/cpp/src/join/hash_join/hash_join.cu b/cpp/src/join/hash_join/hash_join.cu index 163b558eb185..e2478802fc89 100644 --- a/cpp/src/join/hash_join/hash_join.cu +++ b/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 */ @@ -25,6 +25,7 @@ #include #include +#include namespace cudf::detail { @@ -102,8 +103,9 @@ template hash_join::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, + cuda::mr::any_resource mr) + : hash_join{right, has_nulls, compare_nulls, CUCO_DESIRED_LOAD_FACTOR, stream, std::move(mr)} { } @@ -112,7 +114,8 @@ hash_join::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, + cuda::mr::any_resource mr) : _has_nulls(has_nulls), _is_empty{right.num_rows() == 0}, _nulls_equal{compare_nulls}, @@ -124,7 +127,7 @@ hash_join::hash_join(cudf::table_view const& right, {}, {}, {}, - rmm::mr::polymorphic_allocator{}, + rmm::mr::polymorphic_allocator{std::move(mr)}, stream.value()}})}, _right{right}, _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(_right, stream)} @@ -148,16 +151,20 @@ hash_join::hash_join(cudf::table_view const& right, stream); } -template hash_join::hash_join(cudf::table_view const& right, - bool has_nulls, - cudf::null_equality compare_nulls, - rmm::cuda_stream_view stream); +template hash_join::hash_join( + cudf::table_view const& right, + bool has_nulls, + cudf::null_equality compare_nulls, + rmm::cuda_stream_view stream, + cuda::mr::any_resource mr); -template hash_join::hash_join(cudf::table_view const& right, - bool has_nulls, - cudf::null_equality compare_nulls, - double load_factor, - rmm::cuda_stream_view stream); +template hash_join::hash_join( + cudf::table_view const& right, + bool has_nulls, + cudf::null_equality compare_nulls, + double load_factor, + rmm::cuda_stream_view stream, + cuda::mr::any_resource mr); template hash_join::~hash_join() = default; @@ -172,9 +179,14 @@ hash_join::~hash_join() = default; hash_join::hash_join(cudf::table_view const& right, null_equality compare_nulls, - rmm::cuda_stream_view stream) - : hash_join( - right, nullable_join::YES, compare_nulls, cudf::detail::CUCO_DESIRED_LOAD_FACTOR, stream) + rmm::cuda_stream_view stream, + cuda::mr::any_resource mr) + : hash_join(right, + nullable_join::YES, + compare_nulls, + cudf::detail::CUCO_DESIRED_LOAD_FACTOR, + stream, + std::move(mr)) { } @@ -182,9 +194,10 @@ 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, + cuda::mr::any_resource mr) : _impl{std::make_unique( - right, has_nulls == nullable_join::YES, compare_nulls, load_factor, stream)} + right, has_nulls == nullable_join::YES, compare_nulls, load_factor, stream, std::move(mr))} { } diff --git a/cpp/tests/join/join_tests.cpp b/cpp/tests/join/join_tests.cpp index f6a6633f09cb..00dd1ce1431d 100644 --- a/cpp/tests/join/join_tests.cpp +++ b/cpp/tests/join/join_tests.cpp @@ -28,6 +28,7 @@ #include #include +#include #include @@ -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{{3, 1, 2, 0, 2}}.release()); + Table t0(std::move(cols0)); + + CVector cols1; + cols1.emplace_back(column_wrapper{{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 col_gold_0{{0, 2, 2, 3, 4, 4}}; + column_wrapper 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