diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt index 939484f91146..a968660e73bc 100644 --- a/cpp/benchmarks/CMakeLists.txt +++ b/cpp/benchmarks/CMakeLists.txt @@ -297,7 +297,8 @@ ConfigureNVBench( # * parquet reader benchmark ---------------------------------------------------------------------- ConfigureNVBench( PARQUET_READER_NVBENCH io/parquet/parquet_reader_input.cpp io/parquet/parquet_reader_encoding.cpp - io/parquet/parquet_reader_options.cpp io/parquet/reader_common.cpp + io/parquet/parquet_reader_options.cpp io/parquet/parquet_reader_dict_transcode.cpp + io/parquet/reader_common.cpp ) # ################################################################################################## diff --git a/cpp/benchmarks/io/parquet/parquet_reader_dict_transcode.cpp b/cpp/benchmarks/io/parquet/parquet_reader_dict_transcode.cpp new file mode 100644 index 000000000000..f220f3cca2a4 --- /dev/null +++ b/cpp/benchmarks/io/parquet/parquet_reader_dict_transcode.cpp @@ -0,0 +1,79 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include + +#include + +// Size of the data in the benchmark dataframe; chosen to be low enough to allow benchmarks to +// run on most GPUs, but large enough to allow highest throughput +constexpr std::size_t data_size = 512 << 20; + +// Measures reads of dictionary-encoded low-cardinality columns with and without the direct +// Parquet-dictionary -> DICTIONARY32 transcode (`output_dict_columns`), over the column types the +// fast path accepts: flat strings and flat fixed-width INT32/INT64/TIMESTAMP_DAYS columns. +template +void BM_parquet_read_dict_transcode(nvbench::state& state, + nvbench::type_list>) +{ + auto constexpr output_dict_columns = OutputDict == output_dict::YES; + + auto const cardinality = static_cast(state.get_int64("cardinality")); + auto const run_length = static_cast(state.get_int64("run_length")); + + auto const data_types = std::vector{cudf::type_id::INT32, + cudf::type_id::INT64, + cudf::type_id::TIMESTAMP_DAYS, + cudf::type_id::STRING}; + data_profile const profile = + data_profile_builder().cardinality(cardinality).avg_run_length(run_length); + auto const tbl = create_random_table(data_types, table_size_bytes{data_size}, profile); + auto const view = tbl->view(); + + cuio_source_sink_pair source_sink(io_type::HOST_BUFFER); + auto const write_options = + cudf::io::parquet_writer_options::builder(source_sink.make_sink_info(), view) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .build(); + cudf::io::write_parquet(write_options); + + cudf::io::parquet_reader_options read_options = + cudf::io::parquet_reader_options::builder(source_sink.make_source_info()) + .output_dict_columns(output_dict_columns) + .build(); + + auto mem_stats_logger = cudf::memory_stats_logger(); + state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().value())); + state.exec( + nvbench::exec_tag::sync | nvbench::exec_tag::timer, [&](nvbench::launch& launch, auto& timer) { + drop_page_cache_if_enabled(read_options.get_source().filepaths()); + timer.start(); + auto const result = cudf::io::read_parquet(read_options); + timer.stop(); + CUDF_EXPECTS(result.tbl->num_rows() == view.num_rows(), + "Benchmark did not read the entire table"); + CUDF_EXPECTS(result.tbl->num_columns() == view.num_columns(), "Unexpected number of columns"); + }); + + auto const elapsed_time = state.get_summary("nv/cold/time/gpu/mean").get_float64("value"); + state.add_element_count(static_cast(data_size) / elapsed_time, "bytes_per_second"); + state.add_buffer_size( + mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); + state.add_buffer_size(source_sink.size(), "encoded_file_size", "encoded_file_size"); +} + +NVBENCH_BENCH_TYPES(BM_parquet_read_dict_transcode, + NVBENCH_TYPE_AXES(nvbench::enum_type_list)) + .set_name("parquet_read_dict_transcode") + .set_type_axes_names({"output_dict_columns"}) + .set_min_samples(4) + .add_int64_axis("cardinality", {1000}) + .add_int64_axis("run_length", {4}); diff --git a/cpp/include/cudf/io/parquet.hpp b/cpp/include/cudf/io/parquet.hpp index 4cad0aba032f..fdce8383e2c1 100644 --- a/cpp/include/cudf/io/parquet.hpp +++ b/cpp/include/cudf/io/parquet.hpp @@ -343,16 +343,23 @@ class parquet_reader_options { } /** - * @brief Returns whether the reader returns flat string columns as DICTIONARY32 encoded columns + * @brief Returns whether the reader returns flat string and fixed-width columns as DICTIONARY32 + * encoded columns * - * When true, the reader outputs STRING columns as DICTIONARY32 encoded columns. A DICTIONARY32 - * column consists of an INT32 indices child and a STRING keys child. + * When true, the reader outputs eligible flat columns as DICTIONARY32 encoded columns: a signed + * integer indices child (INT8/INT16/INT32, sized to the largest row-group dictionary with the + * same width rule as cudf::dictionary::encode) plus a keys child of the column's logical type. + * Eligible are flat STRING columns and flat fixed-width columns whose physical storage is INT32 + * or INT64 and whose decode is a plain copy (no decimal, timestamp-unit, or width conversion), + * when every data page of the column is dictionary encoded. * * When AST/JIT filters are set, the direct transcode fast path is disabled. - * String columns are materialized, then operated on by the filter. The filtered results are then - * encoded as DICTIONARY32 columns. + * String columns are materialized, then operated on by the filter, and the filtered results are + * encoded as DICTIONARY32 columns. Fixed-width columns participate only in the direct fast path + * and are returned as plain columns whenever it does not apply (filters, chunked or bounded + * reads, or pages that are not dictionary encoded). * - * @return `true` if the reader returns flat string columns as DICTIONARY32 encoded columns + * @return `true` if the reader returns eligible flat columns as DICTIONARY32 encoded columns */ [[nodiscard]] bool is_enabled_output_dict_columns() const { return _output_dict_columns; } @@ -651,9 +658,11 @@ class parquet_reader_options { void enable_prepend_row_index_column(bool val) { _prepend_row_index_column = val; } /** - * @brief Sets to enable/disable trying to output DICTIONARY32 columns for flat string columns. + * @brief Sets to enable/disable trying to output DICTIONARY32 columns for eligible flat string + * and fixed-width columns (see is_enabled_output_dict_columns for eligibility and fallback + * behavior). * - * @param val Boolean indicating whether to output DICTIONARY32 columns for flat string columns + * @param val Boolean indicating whether to output eligible flat columns as DICTIONARY32 */ void enable_output_dict_columns(bool val) { _output_dict_columns = val; } }; @@ -955,13 +964,17 @@ class parquet_reader_options_builder { } /** - * @brief Sets options for enabling/disabling output of DICTIONARY32 columns for flat string - * columns. + * @brief Sets options for enabling/disabling output of DICTIONARY32 columns for eligible flat + * string and fixed-width columns. * - * @param val Boolean value whether to output flat string columns as DICTIONARY32 encoded columns + * @param val Boolean value whether to output eligible flat columns as DICTIONARY32 encoded + * columns * - * @note When enabled, the output columns will be of type DICTIONARY32. When disabled, the output - * columns will be of type STRING. + * @note When enabled, eligible columns are returned as DICTIONARY32 with a signed integer + * indices child sized to the dictionary and a keys child of the logical type; string columns + * are always delivered as DICTIONARY32 (post-read encoded when the direct path does not apply), + * while fixed-width columns fall back to their plain type + * (see parquet_reader_options::is_enabled_output_dict_columns). * * @return this for chaining */ diff --git a/cpp/src/dictionary/detail/concatenate.cu b/cpp/src/dictionary/detail/concatenate.cu index 152527893e4d..7f9edab7adea 100644 --- a/cpp/src/dictionary/detail/concatenate.cu +++ b/cpp/src/dictionary/detail/concatenate.cu @@ -11,15 +11,19 @@ #include #include #include +#include #include #include +#include #include #include #include #include #include #include +#include #include +#include #include #include @@ -37,6 +41,7 @@ #include #include +#include #include namespace cudf { @@ -133,22 +138,57 @@ 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 + * + * @tparam IndexType Integral type of the (already widened) indices column */ +template struct map_indices_fn { cuda::std::span d_offsets; cuda::std::span d_keys_remap; column_device_view d_indices; - __device__ size_type operator()(size_type idx) const + __device__ IndexType operator()(size_type idx) const { - if (d_indices.is_null(idx)) { return 0; } + if (d_indices.is_null(idx)) { return IndexType{0}; } auto cmp = [] __device__(auto const& lhs, auto const& rhs) { return lhs.second < rhs.second; }; auto col_iter = thrust::upper_bound( thrust::seq, d_offsets.begin(), d_offsets.end(), offsets_pair{0, idx}, cmp) - 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)]; + auto col_idx = cuda::std::distance(d_offsets.begin(), col_iter); + auto key_offset = d_offsets[col_idx].first; + auto const old_index = static_cast(d_indices.element(idx)); + return static_cast(d_keys_remap[key_offset + old_index]); + } +}; + +/** + * @brief Dispatches map_indices_fn on the indices type so that the indices are + * read and written with their actual width (INT8/INT16/INT32/...). + */ +struct remap_indices_fn { + template + void operator()(column_device_view const& d_indices, + mutable_column_view const& output, + cuda::std::span d_offsets, + cuda::std::span d_keys_remap, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const + requires(cudf::is_index_type()) + { + auto policy = rmm::exec_policy_nosync(stream, mr); + auto iota = cuda::counting_iterator{0}; + thrust::transform(policy, + iota, + iota + output.size(), + output.begin(), + map_indices_fn{d_offsets, d_keys_remap, d_indices}); + } + + template + void operator()(Args&&...) const + requires(!cudf::is_index_type()) + { + CUDF_FAIL("dictionary indices must be an integral index type"); } }; @@ -231,15 +271,31 @@ 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 @@ -248,9 +304,14 @@ std::unique_ptr concatenate(host_span columns, auto output_view = indices_column->mutable_view(); auto input_view = column_device_view::create(all_indices->view(), stream, temp_mr); 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); + cudf::type_dispatcher(all_indices->type(), + remap_indices_fn{}, + *input_view, + output_view, + cuda::std::span{children_offsets}, + cuda::std::span{final_remap}, + stream, + temp_mr); // remove the bitmask from the all_indices auto null_count = all_indices->null_count(); // get before release() diff --git a/cpp/src/io/parquet/decode_fixed.cu b/cpp/src/io/parquet/decode_fixed.cu index 2d262342ea9f..b735dc3eb9c5 100644 --- a/cpp/src/io/parquet/decode_fixed.cu +++ b/cpp/src/io/parquet/decode_fixed.cu @@ -83,10 +83,11 @@ __device__ static void scan_block_exclusive_sum( } /** - * @brief Write a batch of decoded dictionary indices directly as INT32 output values. + * @brief Write a batch of decoded dictionary indices directly as signed integer output values. * * Used by the Parquet-dict → DICTIONARY32 transcode path: instead of materializing the dictionary - * keys, the per-row dictionary indices are emitted verbatim as the INT32 indices child of the + * keys, the per-row dictionary indices are emitted verbatim as the signed integer indices child + * (INT8/INT16/INT32, per the chunk's dict_index_bytes) of the * output DICTIONARY32 column. * * @tparam block_size Number of threads per block @@ -133,13 +134,22 @@ __device__ void decode_dict_indices_as_int32( return thread_pos; }(); - auto* dst = reinterpret_cast(data_out) + dst_pos; - auto const num_keys = static_cast(s->stream.dict_size / sizeof(string_index_pair)); - auto const idx = sb->dict_idx[rolling_index(src_pos)]; + // string chunks index a string_index_pair table; fixed-width chunks index the raw + // dictionary page, whose entries are the physical value width + auto const entry_size = s->setup.col.physical_type == Type::BYTE_ARRAY + ? sizeof(string_index_pair) + : static_cast(s->output_cvt.dtype_len_in); + auto const num_keys = static_cast(s->stream.dict_size / entry_size); + auto const idx = sb->dict_idx[rolling_index(src_pos)]; if (idx >= num_keys) { s->set_error_code(decode_error::DATA_STREAM_OVERRUN); } else { - *dst = idx; + // the index width follows dictionary::get_indices_type_for_size for the key count + switch (s->output_cvt.dtype_len) { + case 1: reinterpret_cast(data_out)[dst_pos] = static_cast(idx); break; + case 2: reinterpret_cast(data_out)[dst_pos] = static_cast(idx); break; + default: reinterpret_cast(data_out)[dst_pos] = idx; break; + } } } @@ -1306,7 +1316,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8) if constexpr (has_strings_t || has_lists_t || is_dict_int32_t) { if (process_nulls) { uint32_t const dtype_len = [&]() -> uint32_t { - if constexpr (is_dict_int32_t) { return sizeof(int32_t); } + if constexpr (is_dict_int32_t) { return s->setup.col.dict_index_bytes; } if constexpr (has_strings_t) { return sizeof(cudf::size_type); } return s->output_cvt.dtype_len; }(); diff --git a/cpp/src/io/parquet/page_decode.cuh b/cpp/src/io/parquet/page_decode.cuh index 3c24912095aa..0d2e8faff77b 100644 --- a/cpp/src/io/parquet/page_decode.cuh +++ b/cpp/src/io/parquet/page_decode.cuh @@ -1257,6 +1257,12 @@ inline __device__ bool setup_local_page_info(auto* const s, } else if (data_type == Type::INT96) { s->output_cvt.dtype_len = 8; // Convert to 64-bit timestamp } + // Parquet-dict -> DICTIONARY32 transcode: the output holds dictionary indices, so both + // the per-page output offset and the value width follow the index width, not the + // logical type width. (String pages coincidentally match via the string special case.) + if (s->setup.page.kernel_mask == decode_kernel_mask::DICT_INT32) { + s->output_cvt.dtype_len = s->setup.col.dict_index_bytes; + } } // during the decoding step we need to offset the global output buffers @@ -1301,7 +1307,10 @@ inline __device__ bool setup_local_page_info(auto* const s, idx < max_depth - 1 ? sizeof(cudf::size_type) : s->output_cvt.dtype_len; // if this is a string column, then dtype_len is a lie. data will be offsets rather // than (ptr,len) tuples. - if (is_string_col(s->setup.col)) { len = sizeof(cudf::size_type); } + if (is_string_col(s->setup.col) && + s->setup.page.kernel_mask != decode_kernel_mask::DICT_INT32) { + len = sizeof(cudf::size_type); + } nesting_info->data_out += (output_offset * len); } if (nesting_info->string_out != nullptr) { diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index f0399fcbc646..f595da2bdd67 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -257,7 +257,8 @@ enum class decode_kernel_mask { STRING_STREAM_SPLIT_NESTED = (1 << 24), // Run decode kernel for nested BYTE_STREAM_SPLIT string data STRING_STREAM_SPLIT_LIST = (1 << 25), // Run decode kernel for list BYTE_STREAM_SPLIT string data - DICT_INT32 = (1 << 26), // Run decode kernel for dict string → INT32 indices + DICT_INT32 = (1 << 26), // Emit dictionary indices (string or fixed-width chunks) + // as a signed integer column; name predates sized indices }; constexpr uint32_t STRINGS_MASK_NON_DELTA = BitOr(decode_kernel_mask::STRING, @@ -510,9 +511,10 @@ struct ColumnChunkDesc { float list_bytes_per_row_est{}; // for LIST columns, an estimate on number of bytes per row - bool is_strings_to_cat{}; // convert strings to hashes - bool is_large_string_col{}; // `true` if string data uses 64-bit offsets - int32_t src_file_idx{}; // source file index + bool is_strings_to_cat{}; // convert strings to hashes + uint8_t dict_index_bytes{4}; // width of the emitted dictionary index for DICT_INT32 (1/2/4) + bool is_large_string_col{}; // `true` if string data uses 64-bit offsets + int32_t src_file_idx{}; // source file index }; /** diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index b5cf2005f0cf..204c2991b63c 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -768,7 +768,8 @@ table_with_metadata reader_impl::read_chunk_internal(read_mode mode) } // For any columns that were selected for direct parquet-dict → DICTIONARY32 transcode in - // `prepare_dict_transcode`, the entries in `out_columns` are currently INT32 indices columns. + // `prepare_dict_transcode`, the entries in `out_columns` are currently signed integer indices + // columns (INT8/INT16/INT32, sized to the dictionary). // Assemble them into DICTIONARY32 columns here by attaching per-chunk keys; concatenate // remaps indices to the unified keys child. assemble_dict_transcoded_columns(out_columns); diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 6fbd4789c6dd..1fd2379f2a66 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -602,6 +602,9 @@ class reader_impl { // Per-input-column flag indicating whether that column was selected for direct // Parquet-dict → DICTIONARY32 transcode. std::vector _dict_transcode_eligible; + // Per input column: the logical output type of a column selected for direct transcode + // (the DICTIONARY32 keys type); EMPTY when the column is not transcoded. + std::vector _dict_transcode_key_types; }; } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/reader_impl_dict_transcode.cu b/cpp/src/io/parquet/reader_impl_dict_transcode.cu index ef6425ed5e84..424ba79119cf 100644 --- a/cpp/src/io/parquet/reader_impl_dict_transcode.cu +++ b/cpp/src/io/parquet/reader_impl_dict_transcode.cu @@ -13,11 +13,16 @@ #include #include #include +#include #include #include #include #include #include +#include +#include + +#include #include @@ -46,16 +51,39 @@ namespace { return is_string_col(chunk) and chunk.physical_type == Type::BYTE_ARRAY; } +/** + * @brief Whether a column chunk is a flat fixed-width chunk whose decode is a plain copy of the + * physical values (no decimal, timestamp-unit, or width conversion), so its dictionary page holds + * the output values verbatim. + * + * @param chunk The column chunk descriptor to classify + * @param out_type The current output buffer type of the chunk's column + * @return True if the chunk's dictionary page entries are `out_type` values as stored + */ +[[nodiscard]] bool is_fixed_width_dict_chunk(ColumnChunkDesc const& chunk, data_type out_type) +{ + if (chunk.physical_type != Type::INT32 and chunk.physical_type != Type::INT64) { return false; } + if (chunk.logical_type.has_value() and chunk.logical_type->type == LogicalType::DECIMAL) { + return false; + } + // a non-zero clock rate means the decode rescales timestamp values + if (chunk.ts_clock_rate != 0) { return false; } + if (not cudf::is_fixed_width(out_type)) { return false; } + auto const physical_size = chunk.physical_type == Type::INT32 ? std::size_t{4} : std::size_t{8}; + return cudf::size_of(out_type) == physical_size; +} + /** * @brief Per-input-column eligibility flags for Parquet-dict → DICTIONARY32 transcode. * * Each column must satisfy all of these conditions to be eligible for direct transcode. */ struct column_eligibility { - bool has_string_buffer = false; ///< Output buffer is currently typed as STRING - bool has_any_chunk = false; ///< At least one chunk was seen for this column - bool all_chunks_string = true; ///< Every chunk is a flat BYTE_ARRAY string chunk with a dict - bool all_pages_dict = true; ///< Every data page uses a dictionary encoding + data_type key_type{type_id::EMPTY}; ///< Logical output type; the DICTIONARY32 keys type + bool has_transcodable_buffer = false; ///< Output buffer is a flat STRING or fixed-width column + bool has_any_chunk = false; ///< At least one chunk was seen for this column + bool all_chunks_eligible = true; ///< Every chunk is a flat dictionary chunk of the buffer type + bool all_pages_dict = true; ///< Every data page uses a dictionary encoding /** * @brief Whether the column satisfies every transcode-eligibility condition. @@ -64,7 +92,7 @@ struct column_eligibility { */ [[nodiscard]] bool is_eligible() const { - return has_string_buffer and has_any_chunk and all_chunks_string and all_pages_dict; + return has_transcodable_buffer and has_any_chunk and all_chunks_eligible and all_pages_dict; } }; @@ -76,10 +104,13 @@ struct column_eligibility { */ void update_from_chunk(column_eligibility& e, ColumnChunkDesc const& chunk) { - e.has_any_chunk = true; + e.has_any_chunk = true; + bool const chunk_matches_buffer = e.key_type.id() == type_id::STRING + ? is_byte_array_string_chunk(chunk) + : is_fixed_width_dict_chunk(chunk, e.key_type); if (chunk.max_nesting_depth != 1 or chunk.max_level[level_type::REPETITION] != 0 or - not is_byte_array_string_chunk(chunk) or chunk.num_dict_pages < 1) { - e.all_chunks_string = false; + not chunk_matches_buffer or chunk.num_dict_pages < 1) { + e.all_chunks_eligible = false; } } @@ -87,14 +118,17 @@ void update_from_chunk(column_eligibility& e, ColumnChunkDesc const& chunk) * @brief Compute per-input-column eligibility for Parquet-dict → DICTIONARY32 transcode. * * A column is eligible iff - * - the corresponding output buffer is currently typed as STRING (i.e. a flat string column), - * - every chunk of that column is a BYTE_ARRAY string chunk with a dictionary page, + * - the corresponding output buffer is a flat STRING column, or a flat fixed-width column whose + * decode is a plain copy of INT32/INT64 physical values (no decimal, timestamp-unit, or width + * conversion) -- the buffer's logical type becomes the DICTIONARY32 keys type, + * - every chunk of that column is a matching chunk of that type with a dictionary page, * - every data page of every chunk of that column uses DICTIONARY encoding, * - the chunk has a flat (non-list, non-nested) schema. * * @param pass The pass intermediate data holding host-side chunks and pages * @param input_columns The reader's input column descriptors - * @param output_buffers The output column buffers (used to detect flat STRING columns) + * @param output_buffers The output column buffers (used to detect eligible flat columns and + * their logical key types) * @return A vector of per-input-column eligibility records, indexed by input column */ [[nodiscard]] std::vector compute_dict_transcode_eligibility( @@ -105,12 +139,17 @@ void update_from_chunk(column_eligibility& e, ColumnChunkDesc const& chunk) auto const num_input_cols = input_columns.size(); std::vector elig(num_input_cols); - // Mark columns whose output buffer is a flat string column. + // Mark columns whose output buffer is a flat string or fixed-width column, and record the + // buffer's logical type: it becomes the DICTIONARY32 keys type. std::transform( input_columns.begin(), input_columns.end(), elig.begin(), [&](input_column_info const& col) { column_eligibility e{}; - e.has_string_buffer = - col.nesting_depth() == 1 and output_buffers[col.nesting[0]].type.id() == type_id::STRING; + if (col.nesting_depth() != 1) { return e; } + auto const out_type = output_buffers[col.nesting[0]].type; + e.has_transcodable_buffer = + out_type.id() == type_id::STRING or + (cudf::is_fixed_width(out_type) and out_type.id() != type_id::BOOL8); + if (e.has_transcodable_buffer) { e.key_type = out_type; } return e; }); @@ -151,6 +190,40 @@ void update_from_chunk(column_eligibility& e, ColumnChunkDesc const& chunk) return cudf::strings::detail::make_strings_column(begin, begin + entry_count, stream, mr); } +/** + * @brief Build a keys column of `key_type` from a chunk's dictionary entries. + * + * String chunks use the reader's `str_dict_index` (pointer/length pairs). Fixed-width chunks copy + * the PLAIN-encoded dictionary page, whose entries are the output values verbatim (guaranteed by + * `is_fixed_width_dict_chunk`). + * + * @param chunk The column chunk whose dictionary becomes the keys + * @param dict_page_data Device pointer to the chunk's (decompressed) dictionary page payload + * @param key_type The logical output type of the column + * @param entry_count Number of dictionary entries (keys) for this chunk + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the returned column's memory + * @return A `key_type` column holding this chunk's dictionary keys + */ +[[nodiscard]] std::unique_ptr make_keys_column(ColumnChunkDesc const& chunk, + uint8_t const* dict_page_data, + data_type key_type, + size_type entry_count, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + if (key_type.id() == type_id::STRING) { + return make_keys_column_from_index_pairs(chunk.str_dict_index, entry_count, stream, mr); + } + if (entry_count <= 0) { return cudf::make_empty_column(key_type); } + auto const keys_bytes = static_cast(entry_count) * cudf::size_of(key_type); + return std::make_unique(key_type, + entry_count, + rmm::device_buffer{dict_page_data, keys_bytes, stream, mr}, + rmm::device_buffer{}, + 0); +} + } // namespace void reader_impl::prepare_dict_transcode(read_mode mode) @@ -158,6 +231,7 @@ void reader_impl::prepare_dict_transcode(read_mode mode) CUDF_FUNC_RANGE(); _dict_transcode_eligible.assign(_input_columns.size(), false); + _dict_transcode_key_types.assign(_input_columns.size(), data_type{type_id::EMPTY}); if (not _options.output_dict_columns) { return; } @@ -184,6 +258,10 @@ void reader_impl::prepare_dict_transcode(read_mode mode) elig.begin(), elig.end(), _dict_transcode_eligible.begin(), [](column_eligibility const& e) { return e.is_eligible(); }); + std::transform( + elig.begin(), elig.end(), _dict_transcode_key_types.begin(), [](column_eligibility const& e) { + return e.is_eligible() ? e.key_type : data_type{type_id::EMPTY}; + }); auto const num_eligible = std::count(_dict_transcode_eligible.begin(), _dict_transcode_eligible.end(), true); @@ -191,12 +269,29 @@ void reader_impl::prepare_dict_transcode(read_mode mode) auto const num_input_cols = _input_columns.size(); - // Change the output buffer type for eligible columns from STRING → INT32. + // Per-column index type from the largest chunk dictionary, following the same width rule as + // `dictionary::encode` (`get_indices_type_for_size`), so that concatenating batches keeps the + // width unless the merged keys overflow it. + std::vector index_types(num_input_cols, data_type{type_id::INT32}); + { + std::vector max_keys(num_input_cols, 0); + for (auto const& page : pass.pages) { + if ((page.flags & PAGEINFO_FLAGS_DICTIONARY) == 0) { continue; } + auto const col_idx = pass.chunks[page.chunk_idx].src_col_index; + max_keys[col_idx] = std::max(max_keys[col_idx], size_type{page.num_input_values}); + } + std::transform(max_keys.begin(), max_keys.end(), index_types.begin(), [](size_type keys) { + return cudf::dictionary::detail::get_indices_type_for_size(keys); + }); + } + + // Change the output buffer type for eligible columns to the index type: the decode writes the + // dictionary indices instead of the logical values. std::for_each( cuda::counting_iterator{0}, cuda::counting_iterator{num_input_cols}, [&](size_t i) { if (not _dict_transcode_eligible[i]) { return; } auto& out_buf = _output_buffers[_input_columns[i].nesting[0]]; - out_buf.type = data_type{type_id::INT32}; + out_buf.type = index_types[i]; }); // Rewrite per-page `kernel_mask` for eligible columns on the host subpass pages from @@ -207,9 +302,12 @@ void reader_impl::prepare_dict_transcode(read_mode mode) auto const chunk_idx = page.chunk_idx; auto const col_idx = pass.chunks[chunk_idx].src_col_index; if (not _dict_transcode_eligible[col_idx]) { return; } - if (page.kernel_mask == decode_kernel_mask::STRING_DICT) { + if (page.kernel_mask == decode_kernel_mask::STRING_DICT or + page.kernel_mask == decode_kernel_mask::FIXED_WIDTH_DICT) { page.kernel_mask = decode_kernel_mask::DICT_INT32; - any_rewritten = true; + pass.chunks[chunk_idx].dict_index_bytes = + static_cast(cudf::size_of(index_types[col_idx])); + any_rewritten = true; } }); @@ -220,9 +318,10 @@ void reader_impl::prepare_dict_transcode(read_mode mode) return; } - // Push the rewritten `kernel_mask`s back to device so subsequent decode kernels dispatch - // correctly. + // Push the rewritten `kernel_mask`s and chunk index widths back to device so subsequent decode + // kernels dispatch and write correctly. subpass.pages.host_to_device_async(_stream); + pass.chunks.host_to_device_async(_stream); subpass.kernel_mask = std::transform_reduce( subpass.pages.host_begin(), subpass.pages.host_end(), @@ -276,26 +375,32 @@ void reader_impl::assemble_dict_transcoded_columns( // column, so `nesting[0]` is the correct, and only, output-buffer index to use here. auto const out_idx = static_cast(_input_columns[i].nesting[0]); - // Per-chunk key counts from the dictionary page's `num_input_values`, mirrored back to - // host when `pass.pages` was copied by `decode_page_headers`. + auto const key_type = _dict_transcode_key_types[i]; + CUDF_EXPECTS(key_type.id() != type_id::EMPTY, + "Missing keys type for a dict-transcoded column"); + + // Per-chunk key counts and dictionary page payloads from the dictionary page, mirrored + // back to host when `pass.pages` was copied by `decode_page_headers`. std::vector chunk_key_counts(chunk_indices.size(), 0); - std::transform(chunk_indices.begin(), - chunk_indices.end(), - chunk_key_counts.begin(), - [&](size_t chunk_idx) -> size_type { - if (pass.chunks[chunk_idx].dict_page == nullptr) { return 0; } - for (auto const& page : pass.pages) { - if (page.chunk_idx == static_cast(chunk_idx) and - (page.flags & PAGEINFO_FLAGS_DICTIONARY) != 0) { - return static_cast(page.num_input_values); - } - } - return size_type{0}; - }); + std::vector chunk_dict_data(chunk_indices.size(), nullptr); + for (size_t k = 0; k < chunk_indices.size(); ++k) { + auto const chunk_idx = chunk_indices[k]; + if (pass.chunks[chunk_idx].dict_page == nullptr) { continue; } + for (auto const& page : pass.pages) { + if (page.chunk_idx == static_cast(chunk_idx) and + (page.flags & PAGEINFO_FLAGS_DICTIONARY) != 0) { + chunk_dict_data[k] = page.page_data; + chunk_key_counts[k] = static_cast(page.num_input_values); + break; + } + } + } auto& indices_col = out_columns[out_idx]; - CUDF_EXPECTS(indices_col != nullptr and indices_col->type().id() == type_id::INT32, - "Expected INT32 indices column for dict-transcoded flat string column"); + CUDF_EXPECTS(indices_col != nullptr and (indices_col->type().id() == type_id::INT8 or + indices_col->type().id() == type_id::INT16 or + indices_col->type().id() == type_id::INT32), + "Expected a signed integer indices column for a dict-transcoded flat column"); auto indices_owner = std::move(indices_col); // Single row group fast path: the Parquet dictionary page's entries become the keys as-is, @@ -305,8 +410,8 @@ void reader_impl::assemble_dict_transcoded_columns( // `indices_owner` intact for the path below. auto const emit_single_row_group_column = [&]() -> bool { auto const& chunk = pass.chunks[chunk_indices[0]]; - auto keys = make_keys_column_from_index_pairs( - chunk.str_dict_index, chunk_key_counts[0], _stream, _mr); + auto keys = + make_keys_column(chunk, chunk_dict_data[0], key_type, chunk_key_counts[0], _stream, _mr); auto const num_distinct_keys = cudf::detail::distinct_count( keys->view(), null_policy::INCLUDE, nan_policy::NAN_IS_VALID, _stream); if (num_distinct_keys != keys->size()) { return true; } // fall back: dedup below @@ -363,28 +468,31 @@ void reader_impl::assemble_dict_transcoded_columns( // `cudf::detail::concatenate` remaps the indices against the unified keys. std::vector> seg_keys_owners(chunk_indices.size()); std::vector dict_segment_views(chunk_indices.size()); - std::transform( - cuda::counting_iterator{0}, - cuda::counting_iterator{chunk_indices.size()}, - dict_segment_views.begin(), - [&](size_t k) { - auto const chunk_idx = chunk_indices[k]; - auto const& chunk = pass.chunks[chunk_idx]; - - seg_keys_owners[k] = make_keys_column_from_index_pairs( - chunk.str_dict_index, chunk_key_counts[k], _stream, get_current_device_resource_ref()); - - auto const seg_begin = chunk_row_offsets[k]; - auto const seg_end = chunk_row_offsets[k + 1]; - auto const seg_rows = seg_end - seg_begin; - return column_view{data_type{type_id::DICTIONARY32}, - seg_rows, - nullptr, // dictionary parent holds no data - indices_view.null_mask(), // shared with indices_view - seg_null_counts[k], - seg_begin, // reslices shared indices child + null mask - {indices_view, seg_keys_owners[k]->view()}}; - }); + std::transform(cuda::counting_iterator{0}, + cuda::counting_iterator{chunk_indices.size()}, + dict_segment_views.begin(), + [&](size_t k) { + auto const chunk_idx = chunk_indices[k]; + auto const& chunk = pass.chunks[chunk_idx]; + + seg_keys_owners[k] = make_keys_column(chunk, + chunk_dict_data[k], + key_type, + chunk_key_counts[k], + _stream, + get_current_device_resource_ref()); + + auto const seg_begin = chunk_row_offsets[k]; + auto const seg_end = chunk_row_offsets[k + 1]; + auto const seg_rows = seg_end - seg_begin; + return column_view{data_type{type_id::DICTIONARY32}, + seg_rows, + nullptr, // dictionary parent holds no data + indices_view.null_mask(), // shared with indices_view + seg_null_counts[k], + seg_begin, // reslices shared indices child + null mask + {indices_view, seg_keys_owners[k]->view()}}; + }); // `cudf::detail::concatenate` deduplicates + sorts keys and recomputes indices. out_columns[out_idx] = cudf::detail::concatenate(dict_segment_views, _stream, _mr); diff --git a/cpp/tests/copying/concatenate_tests.cpp b/cpp/tests/copying/concatenate_tests.cpp index e23bf57c3216..7d8110f0bfa9 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,84 @@ TYPED_TEST(DictionaryConcatTestFW, FixedWidthKeys) CUDF_TEST_EXPECT_COLUMNS_EQUAL(*decoded, 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); + } +} + +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); +} + +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"}); diff --git a/cpp/tests/io/parquet_reader_dict_test.cpp b/cpp/tests/io/parquet_reader_dict_test.cpp index 057e4ac91315..3f27a2b08211 100644 --- a/cpp/tests/io/parquet_reader_dict_test.cpp +++ b/cpp/tests/io/parquet_reader_dict_test.cpp @@ -480,8 +480,262 @@ TEST_F(ParquetReaderDictTest, MultiColumnMixedEligibility) << "List must remain LIST when output_dict_columns is on (transcode is flat-only)"; CUDF_TEST_EXPECT_COLUMNS_EQUAL(list_col->view(), read_list); - // Non-string column: unchanged INT32. + // All-unique INT32 column: the writer does not dictionary-encode it (data pages are PLAIN), + // so it is ineligible for transcode and stays a plain INT32 column. auto const read_key = read_table->view().column(2); ASSERT_EQ(read_key.type().id(), cudf::type_id::INT32); CUDF_TEST_EXPECT_COLUMNS_EQUAL(key_col, read_key); } + +namespace { + +/// A low-cardinality nullable INT32 column. +cudf::test::fixed_width_column_wrapper make_low_cardinality_ints() +{ + std::mt19937 engine(seed ^ 0xF17ED0UL); + std::uniform_int_distribution value_dist(0, cardinality - 1); + std::bernoulli_distribution null_dist(null_probability); + std::vector values(num_rows); + std::vector valids(num_rows); + for (cudf::size_type i = 0; i < num_rows; ++i) { + values[i] = 1'000'000 + value_dist(engine); + valids[i] = not null_dist(engine); + } + return cudf::test::fixed_width_column_wrapper( + values.begin(), values.end(), valids.begin()); +} + +/// A low-cardinality INT64 column. +cudf::test::fixed_width_column_wrapper make_low_cardinality_int64s() +{ + std::mt19937 engine(seed ^ 0x64B175UL); + std::uniform_int_distribution value_dist(0, cardinality - 1); + std::vector values(num_rows); + for (cudf::size_type i = 0; i < num_rows; ++i) { + values[i] = 3'000'000'000LL + value_dist(engine); + } + return cudf::test::fixed_width_column_wrapper(values.begin(), values.end()); +} + +/// A low-cardinality date (TIMESTAMP_DAYS, physical INT32) column. +cudf::test::fixed_width_column_wrapper make_low_cardinality_dates() +{ + std::mt19937 engine(seed ^ 0xDA7E5UL); + std::uniform_int_distribution value_dist(0, cardinality - 1); + std::vector values(num_rows); + for (cudf::size_type i = 0; i < num_rows; ++i) { + values[i] = 8000 + value_dist(engine); + } + return cudf::test::fixed_width_column_wrapper(values.begin(), + values.end()); +} + +} // namespace + +// Flat fixed-width columns (INT32 with nulls, INT64, DATE) written with dictionary encoding must +// transcode to DICTIONARY32 columns whose keys carry the logical type, across several row groups, +// and decode back to exactly the input. +TEST_F(ParquetReaderDictTest, FlatFixedWidthDictTranscode) +{ + auto const int32_col = make_low_cardinality_ints(); + auto const int64_col = make_low_cardinality_int64s(); + auto const date_col = make_low_cardinality_dates(); + auto const input = cudf::table_view{{int32_col, int64_col, date_col}}; + + auto const filepath = temp_env->get_temp_filepath("FlatFixedWidthDictTranscode.parquet"); + write_parquet(input, filepath); + + auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .output_dict_columns(true) + .build(); + auto const read_table = cudf::io::read_parquet(read_opts).tbl; + + std::array const key_types{ + cudf::type_id::INT32, cudf::type_id::INT64, cudf::type_id::TIMESTAMP_DAYS}; + for (cudf::size_type i = 0; i < read_table->num_columns(); ++i) { + auto const read_col = read_table->view().column(i); + ASSERT_EQ(read_col.type().id(), cudf::type_id::DICTIONARY32) << "column " << i; + cudf::dictionary_column_view const dict{read_col}; + EXPECT_EQ(dict.keys().type().id(), key_types[i]) << "column " << i; + // `cardinality` distinct keys fit INT16 indices (get_indices_type_for_size) + EXPECT_EQ(dict.indices().type().id(), cudf::type_id::INT16) << "column " << i; + auto const decoded = cudf::dictionary::decode(dict); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(input.column(i), decoded->view()); + } +} + +// Fixed-width columns stay plain by default (option off). +TEST_F(ParquetReaderDictTest, FlatFixedWidthNoTranscodeByDefault) +{ + auto const int32_col = make_low_cardinality_ints(); + auto const input = cudf::table_view{{int32_col}}; + auto const filepath = temp_env->get_temp_filepath("FlatFixedWidthNoTranscode.parquet"); + write_parquet(input, filepath); + + auto const read_table = + cudf::io::read_parquet( + cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}).build()) + .tbl; + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::INT32); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(int32_col, read_col); +} + +// Under an AST filter the direct fast path is off: string columns are still delivered as +// DICTIONARY32 via the post-hoc encode, while fixed-width columns fall back to plain. +TEST_F(ParquetReaderDictTest, FixedWidthFilterFallsBackToPlain) +{ + auto const int32_col = make_low_cardinality_ints(); + auto const string_col = make_low_cardinality_strings(); + auto const input = cudf::table_view{{int32_col, string_col}}; + auto const filepath = temp_env->get_temp_filepath("FixedWidthFilterFallback.parquet"); + write_parquet(input, filepath); + + auto const ref = cudf::ast::column_reference(0); + auto literal_value = cudf::numeric_scalar(1'000'000 + cardinality / 2); + auto const literal = cudf::ast::literal(literal_value); + auto const expr = cudf::ast::operation(cudf::ast::ast_operator::LESS, ref, literal); + + auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .filter(expr) + .output_dict_columns(true) + .build(); + auto const read_table = cudf::io::read_parquet(read_opts).tbl; + + auto const read_int = read_table->view().column(0); + ASSERT_EQ(read_int.type().id(), cudf::type_id::INT32); + auto const read_str = read_table->view().column(1); + ASSERT_EQ(read_str.type().id(), cudf::type_id::DICTIONARY32); + + // Cross-check the surviving rows against a plain filtered read. + auto const plain_table = cudf::io::read_parquet(cudf::io::parquet_reader_options::builder( + cudf::io::source_info{filepath}) + .filter(expr) + .build()) + .tbl; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(plain_table->view().column(0), read_int); + auto const decoded_str = cudf::dictionary::decode(cudf::dictionary_column_view(read_str)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(plain_table->view().column(1), decoded_str->view()); +} + +namespace { + +/// A nullable INT32 column with `distinct` distinct values over `rows` rows. +cudf::test::fixed_width_column_wrapper make_int_column_with_cardinality( + cudf::size_type rows, cudf::size_type distinct, unsigned int local_seed) +{ + std::mt19937 engine(local_seed); + std::uniform_int_distribution value_dist(0, distinct - 1); + std::vector values(rows); + for (cudf::size_type i = 0; i < rows; ++i) { + values[i] = 5'000'000 + value_dist(engine); + } + return cudf::test::fixed_width_column_wrapper(values.begin(), values.end()); +} + +} // namespace + +// The emitted index width follows the largest row-group dictionary: <= 127 keys use INT8 and +// > 32767 keys use INT32 (the INT16 middle case is covered by FlatFixedWidthDictTranscode). +TEST_F(ParquetReaderDictTest, FixedWidthIndexWidths) +{ + // INT8: 100 distinct keys across the default row groups. + { + auto const col = make_int_column_with_cardinality(num_rows, 100, seed ^ 0x1D8); + auto const filepath = temp_env->get_temp_filepath("FixedWidthIndexWidthInt8.parquet"); + write_parquet(cudf::table_view{{col}}, filepath); + auto const read_table = read_parquet_as_dict(filepath).tbl; + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::DICTIONARY32); + cudf::dictionary_column_view const dict{read_col}; + EXPECT_EQ(dict.indices().type().id(), cudf::type_id::INT8); + auto const decoded = cudf::dictionary::decode(dict); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col, decoded->view()); + } + // INT32: a single row group whose dictionary exceeds 32767 keys. + { + constexpr cudf::size_type wide_rows = 136'000; + constexpr cudf::size_type wide_distinct = 34'000; + // deterministic: each value appears four times in adjacent runs, so the dictionary holds + // `wide_distinct` keys and dictionary encoding beats PLAIN in size (the writer skips the + // dictionary, even with policy ALWAYS, when it would not be smaller) + std::vector wide_values(wide_rows); + for (cudf::size_type i = 0; i < wide_rows; ++i) { + wide_values[i] = 5'000'000 + (i / 4) % wide_distinct; + } + cudf::test::fixed_width_column_wrapper const col(wide_values.begin(), + wide_values.end()); + auto const filepath = temp_env->get_temp_filepath("FixedWidthIndexWidthInt32.parquet"); + auto const options = cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, + cudf::table_view{{col}}) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .row_group_size_rows(wide_rows) + .build(); + cudf::io::write_parquet(options); + auto const read_table = read_parquet_as_dict(filepath).tbl; + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::DICTIONARY32); + cudf::dictionary_column_view const dict{read_col}; + EXPECT_EQ(dict.indices().type().id(), cudf::type_id::INT32); + EXPECT_GT(dict.keys_size(), 32767); + auto const decoded = cudf::dictionary::decode(dict); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col, decoded->view()); + } +} + +// Bounded reads (skip_rows / num_rows) disable the direct fast path; fixed-width columns fall +// back to their plain type and match a plain bounded read. +TEST_F(ParquetReaderDictTest, FixedWidthBoundedReadFallsBackToPlain) +{ + auto const col = make_low_cardinality_ints(); + auto const filepath = temp_env->get_temp_filepath("FixedWidthBoundedRead.parquet"); + write_parquet(cudf::table_view{{col}}, filepath); + + auto const bounded = [&](bool output_dict) { + return cudf::io::read_parquet( + cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .skip_rows(row_group_size / 2) + .num_rows(row_group_size) + .output_dict_columns(output_dict) + .build()) + .tbl; + }; + auto const read_table = bounded(true); + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::INT32); + auto const plain_table = bounded(false); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(plain_table->view().column(0), read_col); +} + +// Chunked reads disable the direct fast path; fixed-width columns come back plain in every +// output chunk and reassemble to the input. +TEST_F(ParquetReaderDictTest, FixedWidthChunkedReadFallsBackToPlain) +{ + auto const col = make_low_cardinality_ints(); + auto const filepath = temp_env->get_temp_filepath("FixedWidthChunkedRead.parquet"); + write_parquet(cudf::table_view{{col}}, filepath); + + auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .output_dict_columns(true) + .build(); + auto reader = cudf::io::chunked_parquet_reader(/*chunk_read_limit=*/8 * 1024, read_opts); + + std::vector> chunks; + std::vector views; + cudf::size_type total_rows = 0; + int num_chunks = 0; + while (reader.has_next()) { + auto chunk = reader.read_chunk(); + auto const read_col = chunk.tbl->view().column(0); + if (read_col.size() == 0) { continue; } + ASSERT_EQ(read_col.type().id(), cudf::type_id::INT32); + total_rows += read_col.size(); + ++num_chunks; + chunks.push_back(std::move(chunk.tbl)); + views.push_back(chunks.back()->view().column(0)); + } + ASSERT_EQ(total_rows, num_rows); + EXPECT_GT(num_chunks, 1) << "byte limit should split the read into multiple chunks"; + auto const combined = cudf::concatenate(views); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col, combined->view()); +}