diff --git a/cpp/include/cudf/dictionary/detail/update_keys.hpp b/cpp/include/cudf/dictionary/detail/update_keys.hpp index 0fb608acfa0..a52ebaf27f1 100644 --- a/cpp/include/cudf/dictionary/detail/update_keys.hpp +++ b/cpp/include/cudf/dictionary/detail/update_keys.hpp @@ -61,10 +61,30 @@ std::unique_ptr set_keys(dictionary_column_view const& dictionary_column rmm::device_async_resource_ref mr); /** - * @copydoc - * cudf::dictionary::match_dictionaries(std::vector,rmm::device_async_resource_ref) + * @brief Remap the indices of a dictionary column to a new key set, returning + * only the remapped index column with its null mask. + * + * Like set_keys() but does not copy the key set or build a dictionary column. + * Rows whose key value is not found in new_keys are mapped to null. + * + * @throws std::invalid_argument if new_keys is empty or contains nulls + * @throws cudf::data_type_error if new_keys and input.keys() data_types do not match * + * @param input Dictionary column whose indices are to be remapped. + * @param new_keys Key column to remap indices into. Must be non-empty, null-free, + * and the same type as input's keys. * @param stream CUDA stream used for device memory operations and kernel launches. + * @param mr Device memory resource used to allocate the returned column's device memory. + * @return Integer column of remapped indices with the null mask from the remap applied. + */ +std::unique_ptr remap_indices(dictionary_column_view const& input, + column_view const& new_keys, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr); + +/** + * @copydoc + * cudf::dictionary::match_dictionaries(std::vector,rmm::device_async_resource_ref) */ std::vector> match_dictionaries( std::span input, @@ -93,5 +113,49 @@ std::vector> match_dictionaries( std::pair>, std::vector> match_dictionaries( std::vector tables, cuda::stream_ref stream, rmm::device_async_resource_ref mr); +/** + * @brief Like match_dictionaries() but returns index columns in place of the + * dictionary columns + * + * Computes the merged unique key set across all input dictionaries, remaps each + * dictionary's indices to that key set, and returns the resulting index columns + * (with null masks). The merged keys are not returned. + * + * This is more efficient than match_dictionaries() when the caller only needs + * to compare values by index and does not need the actual keys. + * + * @throws std::invalid_argument if input is empty + * + * @param input Span of dictionary column views to match + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the returned columns' device memory + * @return One index column per input dictionary, in the same order as input + */ +std::vector> match_dictionaries_to_indices( + std::span input, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr); + +/** + * @brief Like match_dictionaries() but substitutes index columns in place + * of dictionary columns in the returned table_views + * + * For each dictionary column found across the input `tables`, computes the merged + * unique key set and remaps each column's indices to it. The returned table_views + * reference these index columns (and the originals for non-dictionary columns). + * The merged keys are not returned. + * + * @throws std::invalid_argument if tables is empty + * + * @param tables Vector of table_views containing dictionary columns to be matched + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the returned columns' device memory + * @return Index column owners and updated table_views with index columns substituted + */ +std::pair>, std::vector> +match_dictionaries_to_indices(std::vector tables, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr); + } // namespace dictionary::detail } // namespace cudf diff --git a/cpp/src/dictionary/match_keys.cu b/cpp/src/dictionary/match_keys.cu index 269cb403998..5cf2d70cfac 100644 --- a/cpp/src/dictionary/match_keys.cu +++ b/cpp/src/dictionary/match_keys.cu @@ -158,6 +158,78 @@ std::pair>, std::vector> match_d return {std::move(dictionary_columns), std::move(updated_tables)}; } +std::vector> match_dictionaries_to_indices( + std::span input, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + CUDF_EXPECTS(not input.empty(), "expect at least one dictionary", std::invalid_argument); + + auto temp_mr = cudf::get_current_device_resource_ref(); + std::vector keys(input.size()); + std::transform(input.begin(), input.end(), keys.begin(), [](auto& col) { return col.keys(); }); + auto all_keys = cudf::detail::concatenate(keys, stream, temp_mr); + + auto new_keys = cudf::type_dispatcher( + keys.front().type(), unique_keys_dispatch_fn{}, all_keys->view(), stream, temp_mr); + auto keys_view = new_keys->view(); + + std::vector> result(input.size()); + std::transform(input.begin(), input.end(), result.begin(), [keys_view, mr, stream](auto& col) { + return remap_indices(col, keys_view, stream, mr); + }); + return result; +} + +std::pair>, std::vector> +match_dictionaries_to_indices(std::vector tables, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + CUDF_EXPECTS(not tables.empty(), "expect at least one table", std::invalid_argument); + + // Make a copy of all the column views from each table_view + std::vector> updated_columns; + std::transform(tables.begin(), tables.end(), std::back_inserter(updated_columns), [](auto& t) { + return std::vector(t.begin(), t.end()); + }); + + // Each column in a table must match in type. + // Once a dictionary column is found, all the corresponding column_views in the + // other table_views are matched. The matched column_views then replace the originals. + std::vector> index_columns; + auto first_table = tables.front(); + for (size_type col_idx = 0; col_idx < first_table.num_columns(); ++col_idx) { + auto col = first_table.column(col_idx); + if (col.type().id() == type_id::DICTIONARY32) { + std::vector dict_views; // hold all column_views at col_idx + std::transform( + tables.begin(), tables.end(), std::back_inserter(dict_views), [col_idx](auto& t) { + return dictionary_column_view(t.column(col_idx)); + }); + // now match the keys and return index columns + auto idx_cols = dictionary::detail::match_dictionaries_to_indices(dict_views, stream, mr); + // replace the updated_columns vector entries for the set of columns at col_idx + auto dict_col_idx = 0; + for (auto& v : updated_columns) + v[col_idx] = idx_cols[dict_col_idx++]->view(); + // move the updated index columns into the main output vector + std::move(idx_cols.begin(), idx_cols.end(), std::back_inserter(index_columns)); + } + } + // All the new column_views are now included in updated_columns + + // Rebuild the table_views from the column_views + std::vector updated_tables; + std::transform(updated_columns.begin(), + updated_columns.end(), + std::back_inserter(updated_tables), + [](auto& v) { return table_view{v}; }); + + // Return the new index columns and table_views + return {std::move(index_columns), std::move(updated_tables)}; +} + } // namespace detail // external API diff --git a/cpp/src/dictionary/set_keys.cu b/cpp/src/dictionary/set_keys.cu index 04475f2bbcd..fd380c85b0c 100644 --- a/cpp/src/dictionary/set_keys.cu +++ b/cpp/src/dictionary/set_keys.cu @@ -65,12 +65,18 @@ struct apply_indices_map_fn { } }; -struct set_keys_dispatch_fn { +struct remap_result { + std::unique_ptr indices; + rmm::device_buffer null_mask; + cudf::size_type null_count; +}; + +struct remap_indices_dispatch_fn { template - std::unique_ptr operator()(cudf::dictionary_column_view const& input, - cudf::column_view const& new_keys, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + remap_result operator()(cudf::dictionary_column_view const& input, + cudf::column_view const& new_keys, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) requires(cudf::is_dictionary_key()) { // compute sorted-order so the new_keys can be searched more quickly @@ -116,9 +122,33 @@ struct set_keys_dispatch_fn { stream, mr); + return {std::move(indices_column), std::move(null_mask), null_count}; + } + + template + remap_result operator()(cudf::dictionary_column_view const&, + cudf::column_view const&, + cuda::stream_ref, + rmm::device_async_resource_ref) + requires(not cudf::is_dictionary_key()) + { + CUDF_UNREACHABLE("not a valid dictionary key type"); + } +}; + +struct set_keys_dispatch_fn { + template + std::unique_ptr operator()(cudf::dictionary_column_view const& input, + cudf::column_view const& new_keys, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) + requires(cudf::is_dictionary_key()) + { + auto [indices, null_mask, null_count] = type_dispatcher( + new_keys.type(), remap_indices_dispatch_fn{}, input, new_keys, stream, mr); auto keys_column = std::make_unique(new_keys, stream, mr); return make_dictionary_column( - std::move(keys_column), std::move(indices_column), std::move(null_mask), null_count); + std::move(keys_column), std::move(indices), std::move(null_mask), null_count); } template @@ -133,6 +163,22 @@ struct set_keys_dispatch_fn { }; } // namespace +std::unique_ptr remap_indices(dictionary_column_view const& input, + column_view const& new_keys, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + CUDF_EXPECTS(!new_keys.has_nulls(), "keys parameter must not have nulls", std::invalid_argument); + CUDF_EXPECTS(!new_keys.is_empty(), "keys cannot be empty", std::invalid_argument); + CUDF_EXPECTS( + cudf::have_same_types(input.keys(), new_keys), "keys types must match", cudf::data_type_error); + + auto [indices, null_mask, null_count] = type_dispatcher( + new_keys.type(), remap_indices_dispatch_fn{}, input, new_keys, stream, mr); + indices->set_null_mask(std::move(null_mask), null_count); + return std::move(indices); +} + std::unique_ptr set_keys(dictionary_column_view const& input, column_view const& new_keys, cuda::stream_ref stream, diff --git a/cpp/src/join/join.cu b/cpp/src/join/join.cu index 9334125e05d..2a8a2352428 100644 --- a/cpp/src/join/join.cu +++ b/cpp/src/join/join.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "join_common_utils.hpp" @@ -32,12 +32,9 @@ inner_join(table_view const& left_input, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - // Make sure any dictionary columns have matched key sets. - // This will return any new dictionary columns created as well as updated table_views. - auto matched = cudf::dictionary::detail::match_dictionaries( - {left_input, right_input}, - stream, - cudf::get_current_device_resource_ref()); // temporary objects returned + // match dictionary key sets so indices are comparable across tables + auto matched = cudf::dictionary::detail::match_dictionaries_to_indices( + {left_input, right_input}, stream, cudf::get_current_device_resource_ref()); // now rebuild the table views with the updated ones auto const left = matched.second.front(); @@ -67,12 +64,10 @@ left_join(table_view const& left_input, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - // Make sure any dictionary columns have matched key sets. - // This will return any new dictionary columns created as well as updated table_views. - auto matched = cudf::dictionary::detail::match_dictionaries( - {left_input, right_input}, // these should match - stream, - cudf::get_current_device_resource_ref()); // temporary objects returned + // match dictionary keys so indices are comparable across tables + auto matched = cudf::dictionary::detail::match_dictionaries_to_indices( + {left_input, right_input}, stream, cudf::get_current_device_resource_ref()); + // now rebuild the table views with the updated ones table_view const left = matched.second.front(); table_view const right = matched.second.back(); @@ -92,12 +87,10 @@ full_join(table_view const& left_input, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - // Make sure any dictionary columns have matched key sets. - // This will return any new dictionary columns created as well as updated table_views. - auto matched = cudf::dictionary::detail::match_dictionaries( - {left_input, right_input}, // these should match - stream, - cudf::get_current_device_resource_ref()); // temporary objects returned + // match dictionary key sets so indices are comparable across tables + auto matched = cudf::dictionary::detail::match_dictionaries_to_indices( + {left_input, right_input}, stream, cudf::get_current_device_resource_ref()); + // now rebuild the table views with the updated ones table_view const left = matched.second.front(); table_view const right = matched.second.back(); diff --git a/cpp/tests/join/join_tests.cpp b/cpp/tests/join/join_tests.cpp index 60717d35b89..6bda4c4274d 100644 --- a/cpp/tests/join/join_tests.cpp +++ b/cpp/tests/join/join_tests.cpp @@ -2783,6 +2783,45 @@ TEST_P(JoinParameterizedTest, DictionaryInnerJoinWithNulls) CUDF_TEST_EXPECT_TABLES_EQUIVALENT(*sorted_gold, *sorted_result); } +TEST_P(JoinParameterizedTest, DictionaryInnerJoinKeyWithNulls) +{ + auto algo = GetParam(); + column_wrapper col0_0{{3, 1, 2, 0, 2}}; + strcol_wrapper col0_1_w({"s1", "s1", "", "s4", "s0"}, {true, true, false, true, true}); + auto col0_1 = cudf::dictionary::encode(col0_1_w); + column_wrapper col0_2{{0, 1, 2, 4, 1}}; + + column_wrapper col1_0{{2, 2, 0, 4, 3}}; + strcol_wrapper col1_1_w({"s1", "", "", "s2", "s1"}, {true, false, false, true, true}); + auto col1_1 = cudf::dictionary::encode(col1_1_w); + column_wrapper col1_2{{1, 0, 1, 2, 1}}; + + auto t0 = cudf::table_view({col0_0, col0_1->view(), col0_2}); + auto t1 = cudf::table_view({col1_0, col1_1->view(), col1_2}); + + // left[2](2, null) matches right[1](2, null); left[0](3, "s1") matches right[4](3, "s1") + auto result = inner_join(t0, t1, {0, 1}, {0, 1}, cudf::null_equality::EQUAL, algo); + auto result_view = result->view(); + auto decoded1 = cudf::dictionary::decode(result_view.column(1)); + auto decoded4 = cudf::dictionary::decode(result_view.column(4)); + std::vector result_decoded({result_view.column(0), + decoded1->view(), + result_view.column(2), + result_view.column(3), + decoded4->view(), + result_view.column(5)}); + auto result_sort_order = cudf::sorted_order(cudf::table_view(result_decoded)); + auto sorted_result = cudf::gather(cudf::table_view(result_decoded), *result_sort_order); + + auto g0 = cudf::table_view({col0_0, col0_1_w, col0_2}); + auto g1 = cudf::table_view({col1_0, col1_1_w, col1_2}); + auto gold = inner_join(g0, g1, {0, 1}, {0, 1}); + auto gold_sort_order = cudf::sorted_order(gold->view()); + auto sorted_gold = cudf::gather(gold->view(), *gold_sort_order); + + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(*sorted_gold, *sorted_result); +} + TEST_F(JoinDictionaryTest, FullJoinNoNulls) { column_wrapper col0_0{{3, 1, 2, 0, 3}};