Fix dictionary concatenate for INT8/INT16 indices

cpp/benchmarks/dictionary/concatenate.cpp

@@ -9,13 +9,15 @@
 #include <cudf/column/column.hpp>
 #include <cudf/column/column_view.hpp>
 #include <cudf/concatenate.hpp>
+#include <cudf/dictionary/dictionary_column_view.hpp>
 #include <cudf/dictionary/encode.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/span.hpp>
 
 #include <nvbench/nvbench.cuh>
 
+#include <string>
 #include <vector>
 
 static void bench_dictionary_concatenate(nvbench::state& state)
@@ -56,3 +58,72 @@ NVBENCH_BENCH(bench_dictionary_concatenate)
   .add_int64_axis("num_rows", {262144, 2097152, 16777216, 67108864})
   .add_int64_axis("cardinality", {10})
   .add_int64_axis("num_cols", {2, 10, 20});
+
+/**
+ * @brief Concatenate dictionary columns with narrow, mixed, or widening indices
+ *
+ * The `indices` axis selects the input configuration:
+ * - `int8`, `int16`, `int32`: every column has the same key set and the given indices type,
+ *   so the output keeps that indices type;
+ * - `mixed`: alternating INT8 and INT16 indices over the same key set, so the narrower inputs are
+ *   cast to INT16 before concatenation;
+ * - `widen`: each column has its own 100 keys with INT8 indices, so the concatenated key set no
+ *   longer fits INT8 and the output indices are widened to INT16.
+ */
+static void bench_dictionary_concatenate_indices(nvbench::state& state)
+{
+  auto const num_rows = static_cast<cudf::size_type>(state.get_int64("num_rows"));
+  auto const num_cols = static_cast<cudf::size_type>(state.get_int64("num_cols"));
+  auto const indices  = state.get_string("indices");
+
+  auto stream = cudf::get_default_stream();
+
+  auto const indices_type = [&](cudf::size_type col) {
+    if (indices == "int16") { return cudf::type_id::INT16; }
+    if (indices == "int32") { return cudf::type_id::INT32; }
+    if (indices == "mixed") { return col % 2 == 0 ? cudf::type_id::INT8 : cudf::type_id::INT16; }
+    return cudf::type_id::INT8;  // int8, widen
+  };
+  // 100 distinct keys per column fit INT8 indices; "widen" gives every column its own key range
+  auto constexpr keys_per_column = 100;
+  auto const key_range_start     = [&](cudf::size_type col) {
+    return indices == "widen" ? col * keys_per_column : 0;
+  };
+
+  auto columns = std::vector<std::unique_ptr<cudf::column>>{};
+  auto views   = std::vector<cudf::column_view>{};
+  for (cudf::size_type i = 0; i < num_cols; ++i) {
+    auto const lo              = key_range_start(i);
+    data_profile const profile = data_profile_builder().distribution(
+      cudf::type_id::INT32, distribution_id::UNIFORM, lo, lo + keys_per_column - 1);
+    auto input = create_random_column(cudf::type_id::INT32, row_count{num_rows}, profile);
+    columns.emplace_back(
+      cudf::dictionary::encode(input->view(), cudf::data_type{indices_type(i)}, stream));
+    views.push_back(columns.back()->view());
+  }
+
+  auto input_table = cudf::table(std::move(columns));
+
+  state.add_global_memory_reads<uint8_t>(input_table.alloc_size());
+  auto result = cudf::concatenate(views, stream);
+  state.add_global_memory_writes<uint8_t>(result->alloc_size());
+  // throughput is per processed index; the resulting key count is a plain summary
+  state.add_element_count(static_cast<double>(num_rows) * num_cols);
+  auto& keys_summary = state.add_summary("output_keys");
+  keys_summary.set_string("description", "Number of keys in the concatenated dictionary");
+  keys_summary.set_int64("value", cudf::dictionary_column_view(result->view()).keys_size());
+
+  state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value()));
+  auto const mem_stats_logger = cudf::memory_stats_logger();
+
+  state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) { cudf::concatenate(views, stream); });
+
+  state.add_buffer_size(
+    mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage");
+}
+
+NVBENCH_BENCH(bench_dictionary_concatenate_indices)
+  .set_name("concatenate_indices")
+  .add_int64_axis("num_rows", {262144, 2097152, 16777216})
+  .add_int64_axis("num_cols", {2, 10})
+  .add_string_axis("indices", {"int8", "int16", "int32", "mixed", "widen"});

cpp/src/dictionary/detail/concatenate.cu

@@ -11,14 +11,17 @@
 #include <cudf/detail/iterator.cuh>
 #include <cudf/detail/row_operator/equality.cuh>
 #include <cudf/detail/row_operator/hashing.cuh>
+#include <cudf/detail/unary.hpp>
 #include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/dictionary/detail/concatenate.hpp>
+#include <cudf/dictionary/detail/encode.hpp>
 #include <cudf/dictionary/dictionary_column_view.hpp>
 #include <cudf/dictionary/dictionary_factories.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/memory_resource.hpp>
+#include <cudf/utilities/traits.hpp>
 #include <cudf/utilities/type_checks.hpp>
 
 #include <rmm/device_uvector.hpp>
@@ -37,6 +40,7 @@
 #include <thrust/transform_scan.h>
 
 #include <algorithm>
+#include <numeric>
 #include <vector>
 
 namespace cudf {
@@ -133,11 +137,15 @@ struct compute_children_offsets_fn {
 /**
  * @brief Functor for mapping the old indices values to the new indices values
  *        based on the new keys arrangement after concatenation
+ *
+ * Indices are read through an indexalator, so any integral indices type is
+ * remapped without type dispatch.
  */
 struct map_indices_fn {
   cuda::std::span<offsets_pair const> d_offsets;
   cuda::std::span<size_type const> d_keys_remap;
-  column_device_view d_indices;
+  column_device_view d_indices;  // for null checks only
+  cudf::detail::input_indexalator d_input;
 
   __device__ size_type operator()(size_type idx) const
   {
@@ -148,7 +156,7 @@ struct map_indices_fn {
                     1;
     auto col_idx    = cuda::std::distance(d_offsets.begin(), col_iter);
     auto key_offset = d_offsets[col_idx].first;
-    return d_keys_remap[key_offset + d_indices.element<size_type>(idx)];
+    return d_keys_remap[key_offset + d_input[idx]];
   }
 };
 
@@ -231,26 +239,43 @@ std::unique_ptr<column> concatenate(host_span<column_view const> columns,
   thrust::gather(
     policy, d_indices.begin(), d_indices.end(), all_keys_remap.begin(), final_remap.begin());
 
-  // next, concatenate the indices
+  // next, concatenate the indices.
+  // The output indices type is the widest of the input indices types, widened further if the
+  // concatenated keys no longer fit in it (e.g. two INT8 dictionaries with 200 distinct keys).
+  auto indices_type = std::accumulate(
+    columns.begin(), columns.end(), data_type{type_id::INT8}, [](data_type widest, auto cv) {
+      // an empty dictionary column may carry no children at all, but a sliced-to-empty
+      // view still has an indices child whose type participates in the selection
+      if (cv.num_children() == 0) { return widest; }
+      auto const t = dictionary_column_view(cv).indices().type();
+      return cudf::size_of(t) > cudf::size_of(widest) ? t : widest;
+    });
+  auto const needed_type = get_indices_type_for_size(keys_column->size());
+  if (cudf::size_of(needed_type) > cudf::size_of(indices_type)) { indices_type = needed_type; }
+
+  std::vector<std::unique_ptr<column>> widened_indices;  // keeps casted indices alive
   std::vector<column_view> indices_views(columns.size());
-  std::transform(columns.begin(), columns.end(), indices_views.begin(), [](auto cv) {
-    auto dict_view = dictionary_column_view(cv);
-    if (dict_view.is_empty()) {
-      return column_view{data_type{type_id::INT32}, 0, nullptr, nullptr, 0};
-    }
-    return dict_view.get_indices_annotated();  // nicely includes validity mask and view offset
-  });
+  std::transform(
+    columns.begin(), columns.end(), indices_views.begin(), [&](auto cv) -> column_view {
+      auto dict_view = dictionary_column_view(cv);
+      if (dict_view.is_empty()) { return column_view{indices_type, 0, nullptr, nullptr, 0}; }
+      auto indices = dict_view.get_indices_annotated();  // includes validity mask and view offset
+      if (indices.type() == indices_type) { return indices; }
+      widened_indices.emplace_back(cudf::detail::cast(indices, indices_type, stream, temp_mr));
+      return widened_indices.back()->view();
+    });
   auto all_indices = cudf::detail::concatenate(indices_views, stream, mr);
 
   // remap the input indices values to the new indices for the new keys order
   auto indices_column = make_numeric_column(
     all_indices->type(), all_indices->size(), mask_state::UNALLOCATED, stream, mr);
-  auto output_view      = indices_column->mutable_view();
-  auto input_view       = column_device_view::create(all_indices->view(), stream, temp_mr);
+  auto input_view = column_device_view::create(all_indices->view(), stream, temp_mr);
+  auto d_input    = cudf::detail::indexalator_factory::make_input_iterator(all_indices->view());
+  auto d_output =
+    cudf::detail::indexalator_factory::make_output_iterator(indices_column->mutable_view());
   auto children_offsets = child_offsets_fn.create_children_offsets(stream, temp_mr);
-  auto map_fn           = map_indices_fn{children_offsets, final_remap, *input_view};
-  thrust::transform(
-    policy, iota, iota + all_indices->size(), output_view.begin<size_type>(), map_fn);
+  auto map_fn           = map_indices_fn{children_offsets, final_remap, *input_view, d_input};
+  thrust::transform(policy, iota, iota + all_indices->size(), d_output, map_fn);
 
   // remove the bitmask from the all_indices
   auto null_count = all_indices->null_count();  // get before release()

cpp/tests/copying/concatenate_tests.cpp

@@ -16,6 +16,7 @@
 #include <cudf/concatenate.hpp>
 #include <cudf/copying.hpp>
 #include <cudf/detail/iterator.cuh>
+#include <cudf/dictionary/dictionary_column_view.hpp>
 #include <cudf/dictionary/encode.hpp>
 #include <cudf/filling.hpp>
 #include <cudf/null_mask.hpp>
@@ -1651,6 +1652,87 @@ TYPED_TEST(DictionaryConcatTestFW, FixedWidthKeys)
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(*decoded, original);
 }
 
+// INT8/INT16 indices survive slice + concatenate and decode back to the original
+TEST_F(DictionaryConcatTest, NarrowIndices)
+{
+  cudf::test::fixed_width_column_wrapper<int32_t> original({20, 10, 0, 5, 15, 15, 10, 5, 20},
+                                                           {1, 1, 0, 1, 1, 1, 1, 1, 1});
+  for (auto const indices_type : {cudf::type_id::INT8, cudf::type_id::INT16}) {
+    auto dictionary = cudf::dictionary::encode(original, cudf::data_type{indices_type});
+    std::vector<cudf::size_type> splits{0, 3, 3, 5, 5, 9};
+    std::vector<cudf::column_view> views = cudf::slice(dictionary->view(), splits);
+    auto result                          = cudf::concatenate(views);
+    // the indices keep their width when the keys still fit
+    EXPECT_EQ(cudf::dictionary_column_view(result->view()).indices().type().id(), indices_type);
+    auto decoded = cudf::dictionary::decode(result->view());
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*decoded, original);
+  }
+}
+
+// inputs with different indices types concatenate to the widest of them
+TEST_F(DictionaryConcatTest, MixedIndicesTypes)
+{
+  cudf::test::fixed_width_column_wrapper<int32_t> first({1, 2, 3, 2});
+  cudf::test::fixed_width_column_wrapper<int32_t> second({4, 5, 6, 1});
+  auto dictionary1 = cudf::dictionary::encode(first, cudf::data_type{cudf::type_id::INT8});
+  auto dictionary2 = cudf::dictionary::encode(second, cudf::data_type{cudf::type_id::INT16});
+  auto result      = cudf::concatenate(std::vector<cudf::column_view>{*dictionary1, *dictionary2});
+  // widest input indices type wins
+  EXPECT_EQ(cudf::dictionary_column_view(result->view()).indices().type().id(),
+            cudf::type_id::INT16);
+  auto decoded = cudf::dictionary::decode(result->view());
+  cudf::test::fixed_width_column_wrapper<int32_t> expected({1, 2, 3, 2, 4, 5, 6, 1});
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*decoded, expected);
+}
+
+// the indices are widened when the concatenated keys no longer fit the input indices type
+TEST_F(DictionaryConcatTest, WidenIndicesWhenKeysOverflow)
+{
+  // two INT8 dictionaries with 100 distinct keys each and no overlap: 200 keys need INT16
+  auto first_begin  = cuda::counting_iterator<int32_t>{0};
+  auto second_begin = cuda::counting_iterator<int32_t>{100};
+  cudf::test::fixed_width_column_wrapper<int32_t> first(first_begin, first_begin + 100);
+  cudf::test::fixed_width_column_wrapper<int32_t> second(second_begin, second_begin + 100);
+  auto dictionary1 = cudf::dictionary::encode(first, cudf::data_type{cudf::type_id::INT8});
+  auto dictionary2 = cudf::dictionary::encode(second, cudf::data_type{cudf::type_id::INT8});
+  auto result      = cudf::concatenate(std::vector<cudf::column_view>{*dictionary1, *dictionary2});
+  EXPECT_EQ(cudf::dictionary_column_view(result->view()).keys_size(), 200);
+  EXPECT_EQ(cudf::dictionary_column_view(result->view()).indices().type().id(),
+            cudf::type_id::INT16);
+  auto decoded = cudf::dictionary::decode(result->view());
+  cudf::test::fixed_width_column_wrapper<int32_t> expected(first_begin, first_begin + 200);
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*decoded, expected);
+}
+
+// an empty (sliced) input still contributes its indices type to the output selection
+TEST_F(DictionaryConcatTest, EmptyViewKeepsIndicesType)
+{
+  cudf::test::fixed_width_column_wrapper<int32_t> narrow({1, 2, 3, 2});
+  auto dictionary1 = cudf::dictionary::encode(narrow, cudf::data_type{cudf::type_id::INT8});
+  cudf::test::fixed_width_column_wrapper<int32_t> wide({4, 5, 6});
+  auto dictionary2 = cudf::dictionary::encode(wide, cudf::data_type{cudf::type_id::INT16});
+  auto empty_wide  = cudf::slice(dictionary2->view(), {0, 0}).front();
+  auto result      = cudf::concatenate(std::vector<cudf::column_view>{*dictionary1, empty_wide});
+  EXPECT_EQ(cudf::dictionary_column_view(result->view()).indices().type().id(),
+            cudf::type_id::INT16);
+  auto decoded = cudf::dictionary::decode(result->view());
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(*decoded, narrow);
+}
+
+TEST_F(DictionaryConcatTest, AllEmptyViews)
+{
+  cudf::test::fixed_width_column_wrapper<int32_t> first({1, 2, 3});
+  auto dictionary1 = cudf::dictionary::encode(first, cudf::data_type{cudf::type_id::INT8});
+  cudf::test::fixed_width_column_wrapper<int32_t> second({4, 5});
+  auto dictionary2 = cudf::dictionary::encode(second, cudf::data_type{cudf::type_id::INT16});
+  auto empty1      = cudf::slice(dictionary1->view(), {0, 0}).front();
+  auto empty2      = cudf::slice(dictionary2->view(), {1, 1}).front();
+  auto result      = cudf::concatenate(std::vector<cudf::column_view>{empty1, empty2});
+  // all-empty inputs short-circuit to an empty (childless) dictionary column
+  EXPECT_EQ(result->size(), 0);
+  EXPECT_EQ(result->type().id(), cudf::type_id::DICTIONARY32);
+}
+
 TEST_F(DictionaryConcatTest, ErrorsTest)
 {
   cudf::test::strings_column_wrapper strings({"aaa", "ddd", "bbb"});