Use dictionary indices in join functions

cpp/include/cudf/dictionary/detail/update_keys.hpp

@@ -61,10 +61,30 @@ std::unique_ptr<column> set_keys(dictionary_column_view const& dictionary_column
                                  rmm::device_async_resource_ref mr);
 
 /**
- * @copydoc
- * cudf::dictionary::match_dictionaries(std::vector<cudf::dictionary_column_view>,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<column> 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<cudf::dictionary_column_view>,rmm::device_async_resource_ref)
  */
 std::vector<std::unique_ptr<column>> match_dictionaries(
   std::span<dictionary_column_view const> input,
@@ -93,5 +113,49 @@ std::vector<std::unique_ptr<column>> match_dictionaries(
 std::pair<std::vector<std::unique_ptr<column>>, std::vector<table_view>> match_dictionaries(
   std::vector<table_view> 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<std::unique_ptr<column>> match_dictionaries_to_indices(
+  std::span<dictionary_column_view const> 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<std::unique_ptr<column>>, std::vector<table_view>>
+match_dictionaries_to_indices(std::vector<table_view> tables,
+                              cuda::stream_ref stream,
+                              rmm::device_async_resource_ref mr);
+
 }  // namespace dictionary::detail
 }  // namespace cudf

cpp/src/dictionary/match_keys.cu

@@ -158,6 +158,78 @@ std::pair<std::vector<std::unique_ptr<column>>, std::vector<table_view>> match_d
   return {std::move(dictionary_columns), std::move(updated_tables)};
 }
 
+std::vector<std::unique_ptr<column>> match_dictionaries_to_indices(
+  std::span<dictionary_column_view const> 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<cudf::column_view> 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<std::unique_ptr<column>> 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<std::unique_ptr<column>>, std::vector<table_view>>
+match_dictionaries_to_indices(std::vector<table_view> 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<std::vector<column_view>> updated_columns;
+  std::transform(tables.begin(), tables.end(), std::back_inserter(updated_columns), [](auto& t) {
+    return std::vector<column_view>(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<std::unique_ptr<column>> 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<dictionary_column_view> 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<table_view> 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

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<cudf::column> indices;
+  rmm::device_buffer null_mask;
+  cudf::size_type null_count;
+};
+
+struct remap_indices_dispatch_fn {
   template <typename T>
-  std::unique_ptr<cudf::column> 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<T>())
   {
     // 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 <typename T>
+  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<T>())
+  {
+    CUDF_UNREACHABLE("not a valid dictionary key type");
+  }
+};
+
+struct set_keys_dispatch_fn {
+  template <typename T>
+  std::unique_ptr<cudf::column> 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<T>())
+  {
+    auto [indices, null_mask, null_count] = type_dispatcher<dispatch_storage_type>(
+      new_keys.type(), remap_indices_dispatch_fn{}, input, new_keys, stream, mr);
     auto keys_column = std::make_unique<cudf::column>(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 <typename T>
@@ -133,6 +163,22 @@ struct set_keys_dispatch_fn {
 };
 }  // namespace
 
+std::unique_ptr<column> 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<dispatch_storage_type>(
+    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<column> set_keys(dictionary_column_view const& input,
                                  column_view const& new_keys,
                                  cuda::stream_ref stream,

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();

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<int32_t> 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<int32_t> col0_2{{0, 1, 2, 4, 1}};
+
+  column_wrapper<int32_t> 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<int32_t> 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<cudf::column_view> 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<int32_t> col0_0{{3, 1, 2, 0, 3}};