diff --git a/cpp/benchmarks/dictionary/concatenate.cpp b/cpp/benchmarks/dictionary/concatenate.cpp index 52540f371ea2..46095ec8f4c7 100644 --- a/cpp/benchmarks/dictionary/concatenate.cpp +++ b/cpp/benchmarks/dictionary/concatenate.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -16,6 +17,7 @@ #include +#include #include 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(state.get_int64("num_rows")); + auto const num_cols = static_cast(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>{}; + auto views = std::vector{}; + 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(input_table.alloc_size()); + auto result = cudf::concatenate(views, stream); + state.add_global_memory_writes(result->alloc_size()); + // throughput is per processed index; the resulting key count is a plain summary + state.add_element_count(static_cast(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"}); diff --git a/cpp/src/dictionary/detail/concatenate.cu b/cpp/src/dictionary/detail/concatenate.cu index 152527893e4d..3caa23407c52 100644 --- a/cpp/src/dictionary/detail/concatenate.cu +++ b/cpp/src/dictionary/detail/concatenate.cu @@ -11,14 +11,17 @@ #include #include #include +#include #include #include +#include #include #include #include #include #include #include +#include #include #include @@ -37,6 +40,7 @@ #include #include +#include #include 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 d_offsets; cuda::std::span 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(idx)]; + return d_keys_remap[key_offset + d_input[idx]]; } }; @@ -231,26 +239,43 @@ std::unique_ptr concatenate(host_span 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> widened_indices; // keeps casted indices alive std::vector 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(), 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() diff --git a/cpp/tests/copying/concatenate_tests.cpp b/cpp/tests/copying/concatenate_tests.cpp index e23bf57c3216..8c379b8feef5 100644 --- a/cpp/tests/copying/concatenate_tests.cpp +++ b/cpp/tests/copying/concatenate_tests.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -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 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 splits{0, 3, 3, 5, 5, 9}; + std::vector 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 first({1, 2, 3, 2}); + cudf::test::fixed_width_column_wrapper 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{*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 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{0}; + auto second_begin = cuda::counting_iterator{100}; + cudf::test::fixed_width_column_wrapper first(first_begin, first_begin + 100); + cudf::test::fixed_width_column_wrapper 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{*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 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 narrow({1, 2, 3, 2}); + auto dictionary1 = cudf::dictionary::encode(narrow, cudf::data_type{cudf::type_id::INT8}); + cudf::test::fixed_width_column_wrapper 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{*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 first({1, 2, 3}); + auto dictionary1 = cudf::dictionary::encode(first, cudf::data_type{cudf::type_id::INT8}); + cudf::test::fixed_width_column_wrapper 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{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"});