diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 81e447291a5b..605e78eb0c14 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -811,6 +811,7 @@ add_library( src/io/parquet/reader_impl.cpp src/io/parquet/reader_impl_chunking.cu src/io/parquet/reader_impl_chunking_utils.cu + src/io/parquet/reader_impl_dict_transcode.cu src/io/parquet/reader_impl_helpers.cpp src/io/parquet/reader_impl_preprocess.cu src/io/parquet/reader_impl_preprocess_utils.cu diff --git a/cpp/include/cudf/io/parquet.hpp b/cpp/include/cudf/io/parquet.hpp index 148dc15b5c69..bbb7766dff76 100644 --- a/cpp/include/cudf/io/parquet.hpp +++ b/cpp/include/cudf/io/parquet.hpp @@ -110,6 +110,8 @@ class parquet_reader_options { type_id _decimal_width{type_id::EMPTY}; // Whether to use JIT compilation for filtering bool _use_jit_filter = false; + // Whether to output flat string columns as DICT32 encoded columns + bool _output_dict_columns = false; // Whether column name matching is case sensitive. In case of multiple // case-insensitive matches, the first matched column is selected bool _case_sensitive_names = true; @@ -340,6 +342,20 @@ class parquet_reader_options { return _prepend_row_index_column; } + /** + * @brief Returns whether the reader returns flat string 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 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. + * + * @return `true` if the reader returns flat string columns as DICTIONARY32 encoded columns + */ + [[nodiscard]] bool is_enabled_output_dict_columns() const { return _output_dict_columns; } + /** * @brief Set a new source location * @@ -633,6 +649,13 @@ class parquet_reader_options { * @param val Boolean indicating whether to prepend the row index column. */ 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. + * + * @param val Boolean indicating whether to output DICTIONARY32 columns for flat string columns + */ + void enable_output_dict_columns(bool val) { _output_dict_columns = val; } }; /** @@ -931,6 +954,23 @@ class parquet_reader_options_builder { return *this; } + /** + * @brief Sets options for enabling/disabling output of DICTIONARY32 columns for flat string + * columns. + * + * @param val Boolean value whether to output flat string 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. + * + * @return this for chaining + */ + parquet_reader_options_builder& output_dict_columns(bool val) + { + options.enable_output_dict_columns(val); + return *this; + } + /** * @brief move parquet_reader_options member once it's built. */ diff --git a/cpp/src/dictionary/detail/concatenate.cu b/cpp/src/dictionary/detail/concatenate.cu index 04bfe521eb47..9b2b275aa9c9 100644 --- a/cpp/src/dictionary/detail/concatenate.cu +++ b/cpp/src/dictionary/detail/concatenate.cu @@ -175,6 +175,9 @@ std::unique_ptr concatenate(host_span columns, return keys; }); + // TODO: Overload function to accept multiple vectors to do the concatenate at once, with a 2D + // kernel. The keys concatenate below and the indices concatenate further down are two separate + // launches over the same set of input columns and could be fused into a single batched call. // first, concatenate all the keys auto all_keys = cudf::detail::concatenate(keys_views, stream, cudf::get_current_device_resource_ref()); diff --git a/cpp/src/io/parquet/decode_fixed.cu b/cpp/src/io/parquet/decode_fixed.cu index 604c5d4c5bbb..a46c17ce26f1 100644 --- a/cpp/src/io/parquet/decode_fixed.cu +++ b/cpp/src/io/parquet/decode_fixed.cu @@ -81,6 +81,72 @@ __device__ static void scan_block_exclusive_sum( } } +/** + * @brief Write a batch of decoded dictionary indices directly as INT32 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 + * output DICTIONARY32 column. + * + * @tparam block_size Number of threads per block + * @tparam has_lists_t Whether the column has a list (repetition) level + * @tparam copy_mode_t Whether destination positions are direct or indirect (nz_idx) mapped + * @tparam state_buf Page state buffer type providing the decoded dictionary indices + * @param s Page decode state for the current page + * @param sb Page state buffers holding the decoded dictionary indices + * @param start First value position (within the page) to write in this batch + * @param end One-past-the-last value position to write in this batch + * @param t Thread index within the block + */ +template +__device__ void decode_dict_indices_as_int32( + page_state_s* s, state_buf* const sb, int start, int end, int t) +{ + constexpr int num_warps = block_size / cudf::detail::warp_size; + constexpr int max_batch_size = num_warps * cudf::detail::warp_size; + + int const leaf_level_index = s->setup.col.max_nesting_depth - 1; + auto const data_out = s->nesting_info[leaf_level_index].data_out; + + int const skipped_leaf_values = s->setup.page.skipped_leaf_values; + + int pos = start; + while (pos < end) { + int const batch_size = min(max_batch_size, end - pos); + int const target_pos = pos + batch_size; + int const thread_pos = pos + t; + + int const dst_pos = [&]() { + if constexpr (copy_mode_t == copy_mode::DIRECT) { + return thread_pos - s->setup.first_row; + } else { + int dst_pos = sb->nz_idx[rolling_index(thread_pos)]; + if constexpr (!has_lists_t) { dst_pos -= s->setup.first_row; } + return dst_pos; + } + }(); + + if (thread_pos < target_pos && dst_pos >= 0) { + int const src_pos = [&]() { + if constexpr (has_lists_t) { return thread_pos + skipped_leaf_values; } + 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)]; + if (idx >= num_keys) { + s->set_error_code(decode_error::DATA_STREAM_OVERRUN); + } else { + *dst = idx; + } + } + + pos += batch_size; + __syncthreads(); + } +} + template __device__ void decode_fixed_width_values( page_state_s* s, state_buf* const sb, int start, int end, int t) @@ -906,6 +972,12 @@ __device__ void skip_ahead_in_decoding(page_state_s* s, block.sync(); } +/** + * @brief Check if the kernel mask decodes dictionary-encoded data (has a dictionary stream). + * + * @tparam kernel_mask_t The decode kernel mask to test + * @return True for fixed-width, string and INT32-index dictionary masks + */ template CUDF_HOST_DEVICE constexpr bool has_dict() { @@ -914,7 +986,21 @@ CUDF_HOST_DEVICE constexpr bool has_dict() (kernel_mask_t == decode_kernel_mask::FIXED_WIDTH_DICT_LIST) || (kernel_mask_t == decode_kernel_mask::STRING_DICT) || (kernel_mask_t == decode_kernel_mask::STRING_DICT_NESTED) || - (kernel_mask_t == decode_kernel_mask::STRING_DICT_LIST); + (kernel_mask_t == decode_kernel_mask::STRING_DICT_LIST) || + (kernel_mask_t == decode_kernel_mask::DICT_INT32); +} + +/** + * @brief Check whether the kernel mask decodes parquet dictionary indices directly to an INT32 + * column. + * + * @tparam kernel_mask_t The decode kernel mask to test + * @return True for the DICT_INT32 mask + */ +template +CUDF_HOST_DEVICE constexpr bool is_dict_int32_output() +{ + return (kernel_mask_t == decode_kernel_mask::DICT_INT32); } template @@ -925,6 +1011,12 @@ CUDF_HOST_DEVICE constexpr bool has_bools() (kernel_mask_t == decode_kernel_mask::BOOLEAN_LIST); } +/** + * @brief Check if the kernel mask decodes a (non-list) nested column. + * + * @tparam kernel_mask_t The decode kernel mask to test + * @return True for the `*_NESTED` masks + */ template CUDF_HOST_DEVICE constexpr bool has_nesting() { @@ -937,6 +1029,12 @@ CUDF_HOST_DEVICE constexpr bool has_nesting() (kernel_mask_t == decode_kernel_mask::STRING_STREAM_SPLIT_NESTED); } +/** + * @brief Check if the kernel mask decodes a list column (has a repetition level). + * + * @tparam kernel_mask_t The decode kernel mask to test + * @return True for the `*_LIST` masks + */ template CUDF_HOST_DEVICE constexpr bool has_lists() { @@ -961,7 +1059,7 @@ CUDF_HOST_DEVICE constexpr bool is_split_decode() } /** - * @brief Kernel for computing fixed width non dictionary column data stored in the pages + * @brief Kernel for computing fixed width column data stored in the pages * * This function will write the page data and the page data's validity to the * output specified in the page's column chunk. If necessary, additional @@ -996,6 +1094,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8) constexpr bool split_decode_t = is_split_decode(); constexpr bool has_strings_t = (static_cast(kernel_mask_t) & STRINGS_MASK_NON_DELTA) != 0; + constexpr bool is_dict_int32_t = is_dict_int32_output(); constexpr int rolling_buf_size = decode_block_size_t * 2; constexpr int rle_run_buffer_size = rle_stream_required_run_buffer_size(); @@ -1170,7 +1269,10 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8) } auto decode_values = [&]() { - if constexpr (has_strings_t) { + if constexpr (is_dict_int32_t) { + decode_dict_indices_as_int32( + s, sb, valid_count, next_valid_count, t); + } else if constexpr (has_strings_t) { uint32_t* const str_offsets = s->setup.col.column_string_offset_base + page_string_offset_indices[page_idx]; string_output_offset = @@ -1199,10 +1301,14 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8) } // Zero-fill null positions after decoding valid values - if constexpr (has_strings_t || has_lists_t) { + if constexpr (has_strings_t || has_lists_t || is_dict_int32_t) { if (process_nulls) { - uint32_t const dtype_len = has_strings_t ? sizeof(cudf::size_type) : s->dtype_len; - int const num_values = [&]() { + uint32_t const dtype_len = [&]() -> uint32_t { + if constexpr (is_dict_int32_t) { return sizeof(int32_t); } + if constexpr (has_strings_t) { return sizeof(cudf::size_type); } + return s->dtype_len; + }(); + int const num_values = [&]() { if constexpr (has_lists_t) { auto const& ni = s->nesting_info[s->setup.col.max_nesting_depth - 1]; return ni.valid_map_offset - init_valid_map_offset; @@ -1366,6 +1472,9 @@ void decode_page_data(cudf::detail::hostdevice_span pages, case decode_kernel_mask::STRING_STREAM_SPLIT_LIST: launch_kernel(int_tag_t<128>{}, kernel_tag_t{}); break; + case decode_kernel_mask::DICT_INT32: + launch_kernel(int_tag_t<128>{}, kernel_tag_t{}); + break; default: CUDF_EXPECTS(false, "Kernel type not handled by this function"); break; } } diff --git a/cpp/src/io/parquet/page_hdr.cu b/cpp/src/io/parquet/page_hdr.cu index b69c26a4d5eb..71dc290f5901 100644 --- a/cpp/src/io/parquet/page_hdr.cu +++ b/cpp/src/io/parquet/page_hdr.cu @@ -232,8 +232,7 @@ __device__ decode_kernel_mask kernel_mask_for_page(PageInfo const& page, return is_list(chunk) ? decode_kernel_mask::STRING_LIST : is_nested(chunk) ? decode_kernel_mask::STRING_NESTED : decode_kernel_mask::STRING; - } else if (page.encoding == Encoding::PLAIN_DICTIONARY || - page.encoding == Encoding::RLE_DICTIONARY) { + } else if (is_dictionary_encoding(page.encoding)) { return is_list(chunk) ? decode_kernel_mask::STRING_DICT_LIST : is_nested(chunk) ? decode_kernel_mask::STRING_DICT_NESTED : decode_kernel_mask::STRING_DICT; @@ -249,8 +248,7 @@ __device__ decode_kernel_mask kernel_mask_for_page(PageInfo const& page, return is_list(chunk) ? decode_kernel_mask::FIXED_WIDTH_NO_DICT_LIST : is_nested(chunk) ? decode_kernel_mask::FIXED_WIDTH_NO_DICT_NESTED : decode_kernel_mask::FIXED_WIDTH_NO_DICT; - } else if (page.encoding == Encoding::PLAIN_DICTIONARY || - page.encoding == Encoding::RLE_DICTIONARY) { + } else if (is_dictionary_encoding(page.encoding)) { return is_list(chunk) ? decode_kernel_mask::FIXED_WIDTH_DICT_LIST : is_nested(chunk) ? decode_kernel_mask::FIXED_WIDTH_DICT_NESTED : decode_kernel_mask::FIXED_WIDTH_DICT; diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index 6adc4e1b270e..46c34c46e3af 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -84,6 +84,17 @@ CUDF_HOST_DEVICE constexpr bool is_supported_encoding(Encoding enc) } } +/** + * @brief Whether a page encoding references a dictionary page. + * + * Both PLAIN_DICTIONARY (legacy) and RLE_DICTIONARY mark a data page whose values are indices into + * the column chunk's dictionary page. + */ +CUDF_HOST_DEVICE constexpr bool is_dictionary_encoding(Encoding enc) +{ + return enc == Encoding::PLAIN_DICTIONARY or enc == Encoding::RLE_DICTIONARY; +} + /** * @brief Atomically OR `error` into `error_code`. */ @@ -223,7 +234,8 @@ enum class decode_kernel_mask { STRING_STREAM_SPLIT = (1 << 23), // Run decode kernel for BYTE_STREAM_SPLIT string data 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 + 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 }; constexpr uint32_t STRINGS_MASK_NON_DELTA = BitOr(decode_kernel_mask::STRING, @@ -664,7 +676,7 @@ struct EncPage { /** * @brief Test if the given column chunk is in a string column */ -__device__ constexpr bool is_string_col(ColumnChunkDesc const& chunk) +CUDF_HOST_DEVICE constexpr bool is_string_col(ColumnChunkDesc const& chunk) { // return true for non-hashed byte_array and fixed_len_byte_array that isn't representing // a decimal. diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index 307015ec2c3f..75394523b594 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -14,7 +14,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -262,6 +264,11 @@ void reader_impl::decode_page_data(read_mode mode, size_t skip_rows, size_t num_ decode_data(decode_kernel_mask::STRING_STREAM_SPLIT_LIST); } + // launch dict-index-as-int32 decoder for flat columns + if (BitAnd(kernel_mask, decode_kernel_mask::DICT_INT32) != 0) { + decode_data(decode_kernel_mask::DICT_INT32); + } + // launch delta byte array decoder if (BitAnd(kernel_mask, decode_kernel_mask::DELTA_BYTE_ARRAY) != 0) { decode_delta_byte_array(subpass.pages, @@ -524,11 +531,21 @@ reader_impl::reader_impl(std::size_t chunk_read_limit, options.is_enabled_use_jit_filter(), options.is_enabled_case_sensitive_names(), options.is_enabled_prepend_source_index_column(), - options.is_enabled_prepend_row_index_column()}, + options.is_enabled_prepend_row_index_column(), + options.is_enabled_output_dict_columns()}, _sources{std::move(sources)}, _output_chunk_read_limit{chunk_read_limit}, _input_pass_read_limit{pass_read_limit} { + // The direct parquet-dict → DICTIONARY32 transcode fast path only supports single-pass, + // non-chunked reads. + if (_options.output_dict_columns and (chunk_read_limit != 0 or pass_read_limit != 0)) { + CUDF_LOG_WARN( + "output_dict_columns: the direct parquet-dict transcode fast path is disabled for chunked / " + "multi-pass reads (non-zero chunk_read_limit or pass_read_limit); falling back to encoding " + "DICTIONARY32 columns at the output."); + } + // Open and parse the source dataset metadata CUDF_EXPECTS(file_metadatas.empty() or file_metadatas.size() == _sources.size(), "Encountered a mismatch in the number of provided data sources and metadatas"); @@ -695,6 +712,12 @@ table_with_metadata reader_impl::read_chunk_internal(read_mode mode) auto& subpass = *pass.subpass; auto const& read_info = subpass.output_chunk_read_info[subpass.current_output_chunk]; + // If the caller asked for direct parquet-dict → DICTIONARY32 transcode, detect per-column + // eligibility and mutate `_output_buffers` / `subpass.pages` before we allocate column buffers + // or dispatch decode kernels. This has to happen before `preprocess_chunk_strings` / + // `allocate_columns` because those branch on `subpass.kernel_mask` and on `out_buf.type`. + prepare_dict_transcode(mode); + // computes: // PageNestingInfo::batch_size for each level of nesting, for each page, taking row bounds into // account. PageInfo::skipped_values, which tells us where to start decoding in the input to @@ -744,6 +767,12 @@ 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. + // 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); + out_columns = cudf::structs::detail::enforce_null_consistency(std::move(out_columns), _stream, _mr); @@ -917,6 +946,24 @@ table_with_metadata reader_impl::finalize_output(read_mode mode, apply_decimal_width_cast(out_columns); + // Encode any remaining flat STRING columns to DICTIONARY32 via a post-hoc + // `dictionary::detail::encode`: columns not produced by the direct transcode fast path, or all of + // them when the fast path was disabled (e.g. under a filter). This is applied to the FINAL output + // table below -- after any filter is evaluated -- so the filter still operates on STRING columns + // and the dictionary is built over only the surviving rows. + auto const encode_output_dict_columns = + [&](std::unique_ptr tbl) -> std::unique_ptr
{ + if (not _options.output_dict_columns) { return tbl; } + auto columns = tbl->release(); + for (auto& col : columns) { + if (col and col->type().id() == type_id::STRING) { + col = + cudf::dictionary::detail::encode(col->view(), data_type{type_id::INT32}, _stream, _mr); + } + } + return std::make_unique
(std::move(columns)); + }; + if (!_output_metadata) { populate_metadata(out_metadata); // Finally, save the output table metadata into `_output_metadata` for reuse next time. @@ -984,15 +1031,16 @@ table_with_metadata reader_impl::finalize_output(read_mode mode, // Exclude columns present in filter only in output auto output_table = cudf::detail::apply_mask( only_output, *predicate, cudf::detail::mask_type::RETENTION, _stream, _mr); - return {std::move(output_table), std::move(out_metadata)}; + return {encode_output_dict_columns(std::move(output_table)), std::move(out_metadata)}; } else { auto output_table = cudf::filter( read_table->view(), final_filter_expr.value().get(), only_output, _stream, _mr); - return {std::move(output_table), std::move(out_metadata)}; + return {encode_output_dict_columns(std::move(output_table)), std::move(out_metadata)}; } } - return {std::make_unique
(std::move(out_columns)), std::move(out_metadata)}; + return {encode_output_dict_columns(std::make_unique
(std::move(out_columns))), + std::move(out_metadata)}; } table_with_metadata reader_impl::read() diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 76d88f52e310..46f4769158c4 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -188,6 +188,27 @@ class reader_impl { */ void preprocess_chunk_strings(read_mode mode, row_range const& read_info); + /** + * @brief Detect per-column eligibility for direct Parquet-dict → DICTIONARY32 transcode, and + * apply the required host-side mutations to `_output_buffers` and `subpass.pages`. + * + * Must be called after `prepare_data()`. Populates `_dict_transcode_eligible` with a bool per + * input column indicating whether the column will be assembled as a DICTIONARY32 output later in + * `assemble_dict_transcoded_columns`. That member is the sole signal of whether the fast path is + * active: `assemble_dict_transcoded_columns` no-ops when no column is eligible. + * + * @param mode Value indicating if the data sources are read all at once or chunk by chunk + */ + void prepare_dict_transcode(read_mode mode); + + /** + * @brief Assemble DICTIONARY32 output columns for input columns that were marked eligible by + * `prepare_dict_transcode`. + * + * @param out_columns The output columns vector to transcode in place. + */ + void assemble_dict_transcoded_columns(std::vector>& out_columns); + /** * @brief Copies over the relevant page mask information for the subpass */ @@ -542,6 +563,8 @@ class reader_impl { bool prepend_source_index_column = false; // Whether to prepend the file-local row index column to the output bool prepend_row_index_column = false; + // Whether to try outputting DICTIONARY32 columns for fully dict-encoded string columns + bool output_dict_columns = false; } _options; // name to reference converter to extract AST output filter @@ -600,6 +623,10 @@ class reader_impl { std::size_t _output_chunk_read_limit{0}; // output chunk size limit in bytes std::size_t _input_pass_read_limit{0}; // input pass memory usage limit in bytes + + // Per-input-column flag indicating whether that column was selected for direct + // Parquet-dict → DICTIONARY32 transcode. + std::vector _dict_transcode_eligible; }; } // 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 new file mode 100644 index 000000000000..796c48447662 --- /dev/null +++ b/cpp/src/io/parquet/reader_impl_dict_transcode.cu @@ -0,0 +1,394 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "reader_impl.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace cudf::io::parquet::detail { + +namespace { + +/** + * @brief Whether a column chunk is a plain BYTE_ARRAY string chunk. + * + * Narrows `is_string_col` (parquet_gpu.hpp) to BYTE_ARRAY only: `is_string_col` also accepts + * FIXED_LEN_BYTE_ARRAY, which is typically a binary payload and is excluded from transcode. This is + * a string-type classifier -- one of several inputs to eligibility, not the eligibility decision. + * + * @param chunk The column chunk descriptor to classify + * @return True if the chunk is a plain (non-categorical, non-decimal) BYTE_ARRAY string chunk + */ +[[nodiscard]] bool is_byte_array_string_chunk(ColumnChunkDesc const& chunk) +{ + return is_string_col(chunk) and chunk.physical_type == Type::BYTE_ARRAY; +} + +/** + * @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 + + /** + * @brief Whether the column satisfies every transcode-eligibility condition. + * + * @return True if the column is eligible for direct DICTIONARY32 transcode + */ + [[nodiscard]] bool is_eligible() const + { + return has_string_buffer and has_any_chunk and all_chunks_string and all_pages_dict; + } +}; + +/** + * @brief Fold a single chunk's properties into its column's eligibility state. + * + * @param e The per-column eligibility state to update in place + * @param chunk The column chunk descriptor to classify + */ +void update_from_chunk(column_eligibility& e, ColumnChunkDesc const& chunk) +{ + e.has_any_chunk = true; + 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; + } +} + +/** + * @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, + * - 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) + * @return A vector of per-input-column eligibility records, indexed by input column + */ +[[nodiscard]] std::vector compute_dict_transcode_eligibility( + pass_intermediate_data const& pass, + std::vector const& input_columns, + std::vector const& output_buffers) +{ + auto const num_input_cols = input_columns.size(); + std::vector elig(num_input_cols); + + // Mark columns whose output buffer is a flat string column. + 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; + return e; + }); + + // Fold per-chunk info into the per-column eligibility flags. + for (auto const& chunk : pass.chunks) { + auto const col_idx = chunk.src_col_index; + update_from_chunk(elig[col_idx], chunk); + } + + // Any non-dictionary data-page encoding disqualifies the whole column. Dictionary pages + // themselves (PAGEINFO_FLAGS_DICTIONARY) are skipped since they are not data pages. + for (auto const& page : pass.pages) { + if ((page.flags & PAGEINFO_FLAGS_DICTIONARY) != 0) { continue; } + auto const chunk_idx = page.chunk_idx; + auto const col_idx = pass.chunks[chunk_idx].src_col_index; + if (not is_dictionary_encoding(page.encoding)) { elig[col_idx].all_pages_dict = false; } + } + + return elig; +} + +/** + * @brief Build a STRING keys column from a chunk's dictionary entries. + * + * @param begin Pointer to the first `string_index_pair` entry for this chunk's dictionary + * @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 STRING column holding this chunk's dictionary keys (empty if `entry_count <= 0`) + */ +[[nodiscard]] std::unique_ptr make_keys_column_from_index_pairs( + string_index_pair const* begin, + size_type entry_count, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + if (entry_count <= 0) { return cudf::make_empty_column(data_type{type_id::STRING}); } + return cudf::strings::detail::make_strings_column(begin, begin + entry_count, stream, mr); +} + +} // namespace + +void reader_impl::prepare_dict_transcode(read_mode mode) +{ + CUDF_FUNC_RANGE(); + + _dict_transcode_eligible.assign(_input_columns.size(), false); + + if (not _options.output_dict_columns) { return; } + + // The fast path requires the whole column to live in a single subpass. For chunked / multi-pass + // reads (non-zero chunk or pass read limit) we skip it and let `finalize_output` produce the + // DICTIONARY32 columns via a post-hoc `dictionary::detail::encode` instead. + if (_output_chunk_read_limit != 0 or _input_pass_read_limit != 0) { return; } + + // Skip the fast path if custom row bounds are in effect. + if (uses_custom_row_bounds(mode)) { return; } + + // AST/JIT filters evaluate predicates on materialized STRING columns, so the direct transcode + // fast path cannot run under a filter. Skip it and let `finalize_output` encode the filtered + // STRING result to DICTIONARY32 via the post-hoc `dictionary::detail::encode` fallback. + if (_expr_conv.get_converted_expr().has_value()) { return; } + + auto& pass = *_pass_itm_data; + auto& subpass = *pass.subpass; + + if (pass.chunks.empty() or subpass.pages.size() == 0) { return; } + + auto const elig = compute_dict_transcode_eligibility(pass, _input_columns, _output_buffers); + std::transform( + elig.begin(), elig.end(), _dict_transcode_eligible.begin(), [](column_eligibility const& e) { + return e.is_eligible(); + }); + + auto const num_eligible = + std::count(_dict_transcode_eligible.begin(), _dict_transcode_eligible.end(), true); + if (num_eligible == 0) { return; } + + auto const num_input_cols = _input_columns.size(); + + // Change the output buffer type for eligible columns from STRING → INT32. + 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}; + }); + + // Rewrite per-page `kernel_mask` for eligible columns on the host subpass pages from + // STRING_DICT → DICT_INT32, then H2D so the device pages agree. + bool any_rewritten = false; + std::for_each(subpass.pages.host_begin(), subpass.pages.host_end(), [&](PageInfo& page) { + if ((page.flags & PAGEINFO_FLAGS_DICTIONARY) != 0) { return; } + 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) { + page.kernel_mask = decode_kernel_mask::DICT_INT32; + any_rewritten = true; + } + }); + + // No page was actually rewritten. Clear the eligibility flags so the member reflects the true + // "inactive" state. + if (not any_rewritten) { + _dict_transcode_eligible.assign(_input_columns.size(), false); + return; + } + + // Push the rewritten `kernel_mask`s back to device so subsequent decode kernels dispatch + // correctly. + subpass.pages.host_to_device_async(_stream); + subpass.kernel_mask = std::transform_reduce( + subpass.pages.host_begin(), + subpass.pages.host_end(), + uint32_t{0}, + std::bit_or<>{}, + [](PageInfo const& page) { return static_cast(page.kernel_mask); }); +} + +void reader_impl::assemble_dict_transcoded_columns( + std::vector>& out_columns) +{ + CUDF_FUNC_RANGE(); + + if (_pass_itm_data == nullptr) { return; } + + // Nothing to assemble unless `prepare_dict_transcode` marked at least one column eligible. + if (std::none_of(_dict_transcode_eligible.begin(), + _dict_transcode_eligible.end(), + [](bool eligible) { return eligible; })) { + return; + } + + auto const& pass = *_pass_itm_data; + + // For each eligible input column, collect its chunks in row-group order, build a per-chunk + // DICTIONARY32 segment (local 0-based indices + per-chunk keys column), and concatenate. + // + // IMPORTANT: Each segment carries row-group-local indices into its own keys column. We do NOT + // pre-shift indices into a global keyspace, because `cudf::dictionary::detail::concatenate` + // already re-maps the indices using `compute_children_offsets_fn`. Pre-shifting would cause + // double-offsetting and out-of-bounds reads in the `dispatch_compute_indices` kernel. + std::for_each( + cuda::counting_iterator{0}, + cuda::counting_iterator{_input_columns.size()}, + [&](size_t i) { + if (not _dict_transcode_eligible[i]) { return; } + + // Gather chunk indices for this input column in row-group order. + std::vector chunk_indices; + chunk_indices.reserve(pass.chunks.size() / std::max(_input_columns.size(), 1)); + std::copy_if(cuda::counting_iterator{0}, + cuda::counting_iterator{pass.chunks.size()}, + std::back_inserter(chunk_indices), + [&](size_t c) { return pass.chunks[c].src_col_index == static_cast(i); }); + if (chunk_indices.empty()) { return; } + + // `out_columns` is indexed by output-buffer (root column) ordinal, not input-column + // ordinal: a nested struct/list column contributes one entry to `_output_buffers` but one + // entry per leaf to `_input_columns`, so `i` and the corresponding root index can diverge + // as soon as any nested column precedes this one. Eligibility requires a flat (depth-1) + // 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`. + 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}; + }); + + 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"); + auto indices_owner = std::move(indices_col); + + // Single row group fast path: the Parquet dictionary page's entries become the keys as-is, + // in page order. If a file carries duplicate dictionary entries, the + // shortcut is skipped and the general multi-chunk path below runs instead: + // `emit_single_row_group_column` returns true when it fell back (keys not distinct), leaving + // `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 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 + out_columns[out_idx] = + cudf::make_dictionary_column(std::move(keys), std::move(indices_owner), _stream, _mr); + return false; + }; + + if (chunk_indices.size() == 1) { + bool const fallback_used = emit_single_row_group_column(); + if (not fallback_used) { return; } + // Keys were not distinct: fall through to the multi-row-group path, which deduplicates. + } + + // Multi-row-group path: the indices buffer is shared (aliased) by per-chunk DICTIONARY32 + // views below via the parent's offset/size, so it must stay alive until concatenate + // completes. + column_view const indices_view{indices_owner->view()}; + + // Per-chunk boundaries along the row axis: chunk k occupies rows + // [chunk_row_offsets[k], chunk_row_offsets[k+1]). + std::vector chunk_row_offsets(chunk_indices.size() + 1, 0); + std::transform( + chunk_indices.begin(), + chunk_indices.end(), + chunk_row_offsets.begin() + 1, + [&](size_t chunk_idx) { return static_cast(pass.chunks[chunk_idx].num_rows); }); + std::inclusive_scan( + chunk_row_offsets.begin() + 1, chunk_row_offsets.end(), chunk_row_offsets.begin() + 1); + CUDF_EXPECTS(chunk_row_offsets.back() == indices_view.size(), + "Row counts on pass chunks must sum to the indices column size"); + + // Pre-compute null counts for all segments in a single kernel launch. Building the + // column_views below requires a per-segment null count, and calling null_count(begin, end) + // inside the loop would launch one kernel per chunk. Batch them here instead. + std::vector seg_null_counts(chunk_indices.size(), 0); + if (indices_view.nullable()) { + std::vector indices_pairs; + indices_pairs.reserve(chunk_indices.size() * 2); + for (size_t k = 0; k < chunk_indices.size(); ++k) { + indices_pairs.push_back(chunk_row_offsets[k]); + indices_pairs.push_back(chunk_row_offsets[k + 1]); + } + seg_null_counts = + cudf::detail::segmented_null_count(indices_view.null_mask(), indices_pairs, _stream); + } + + // Build a per-chunk DICTIONARY32 *view* that aliases the shared decoded INT32 buffer (no + // copy): keys = this chunk's STRING column, indices = `indices_view`. The row range, null + // mask, and null count must all live on the *parent* view (via offset/size), not the indices + // child, because `get_indices_annotated()` rebuilds the indices from the child's `head()` + // plus the parent's offset/size/null_mask -- anything set on the child is ignored. A wrong + // null count (e.g. a hardcoded 0) would silently turn nulls into a valid index once + // `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()}}; + }); + + // `cudf::detail::concatenate` deduplicates + sorts keys and recomputes indices. + out_columns[out_idx] = cudf::detail::concatenate(dict_segment_views, _stream, _mr); + }); +} + +} // namespace cudf::io::parquet::detail diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 25938c6b5e05..06c0462ebed9 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -340,6 +340,7 @@ ConfigureTest( io/parquet_chunked_writer_test.cpp io/parquet_common.cpp io/parquet_misc_test.cpp + io/parquet_reader_dict_test.cpp io/parquet_reader_test.cpp io/parquet_test.cpp io/parquet_v2_test.cpp diff --git a/cpp/tests/io/parquet_reader_dict_test.cpp b/cpp/tests/io/parquet_reader_dict_test.cpp new file mode 100644 index 000000000000..80e2cf6dcf56 --- /dev/null +++ b/cpp/tests/io/parquet_reader_dict_test.cpp @@ -0,0 +1,488 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "parquet_common.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr cudf::size_type num_rows = 5000; +constexpr cudf::size_type cardinality = num_rows / 10; +constexpr cudf::size_type row_group_size = 1000; +constexpr unsigned int seed = 0xcece; +constexpr unsigned int list_strings_seed = seed ^ 0xA5701DUL; +constexpr cudf::size_type max_elements_per_list = 8; +constexpr double null_probability = 0.1; + +// Per-distinct-value prefixes deliberately mixing ASCII with multi-byte UTF-8 (accented Latin, +// Greek, CJK, and an emoji) so the transcode/fallback paths are exercised on non-ASCII keys. The +// numeric suffix in `make_value_string` keeps every distinct value a distinct string, preserving +// the intended cardinality. +std::array const utf8_prefixes{ + "str", "café", "naïve", "Ωμέγα", "日本語", "🚀rocket"}; + +// Map a dictionary value to a UTF-8 string. Distinct values map to distinct +// strings via the numeric suffix. +std::string make_value_string(int value) +{ + return std::string{utf8_prefixes[value % utf8_prefixes.size()]} + "_" + std::to_string(value); +} + +cudf::test::strings_column_wrapper make_low_cardinality_strings() +{ + std::mt19937 engine(seed); + std::uniform_int_distribution value_dist(0, cardinality - 1); + std::bernoulli_distribution null_dist(null_probability); + + std::vector strings(num_rows); + std::vector valids(num_rows); + for (cudf::size_type i = 0; i < num_rows; ++i) { + strings[i] = make_value_string(value_dist(engine)); + valids[i] = not null_dist(engine); + } + + return cudf::test::strings_column_wrapper(strings.begin(), strings.end(), valids.begin()); +} + +std::unique_ptr make_low_cardinality_lists_of_strings() +{ + std::mt19937 engine(list_strings_seed); + std::uniform_int_distribution value_dist(0, cardinality - 1); + std::uniform_int_distribution len_dist(0, max_elements_per_list); + + std::vector offsets; + offsets.reserve(num_rows + 1); + offsets.push_back(0); + std::vector child_strings; + for (cudf::size_type row = 0; row < num_rows; ++row) { + auto const len = len_dist(engine); + for (int e = 0; e < len; ++e) { + child_strings.push_back(make_value_string(value_dist(engine))); + } + offsets.push_back(offsets.back() + static_cast(len)); + } + + auto child = cudf::test::strings_column_wrapper(child_strings.begin(), child_strings.end()); + auto offsets_col = + cudf::test::fixed_width_column_wrapper(offsets.begin(), offsets.end()) + .release(); + + return cudf::make_lists_column( + num_rows, std::move(offsets_col), child.release(), 0, rmm::device_buffer{}); +} + +void write_parquet(cudf::table_view const& input, std::string const& filepath) +{ + // Produce row groups consisting of `row_group_size` rows, with a single (non-chunked) write. Row + // groups are built from whole page fragments, so `max_page_fragment_size` must also be lowered to + // `row_group_size` + // -- otherwise the default 5000-row fragment would force row groups to snap to multiples of 5000 + // instead of the requested size. + auto const options = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, input) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .row_group_size_rows(row_group_size) + .max_page_fragment_size(row_group_size) + .build(); + cudf::io::write_parquet(options); +} + +cudf::io::table_with_metadata read_parquet_as_dict(std::string const& filepath) +{ + auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .output_dict_columns(true) + .build(); + return cudf::io::read_parquet(read_opts); +} + +// A simple INT32 iota column, used as a non-string / filter key column alongside the strings. +cudf::test::fixed_width_column_wrapper make_int_key_column() +{ + std::vector keys(num_rows); + std::iota(keys.begin(), keys.end(), 0); + return cudf::test::fixed_width_column_wrapper(keys.begin(), keys.end()); +} + +// Build a string column whose first half is low-cardinality (a small dictionary that fits the +// writer's dictionary budget) and whose second half is all-distinct (a dictionary too large for +// the budget). cuDF chooses dictionary use per column-chunk, so under an ADAPTIVE policy with a +// small `max_dictionary_size` the resulting column carries a mix of dictionary-encoded and +// PLAIN-encoded chunks -- which makes it ineligible for the direct transcode fast path. +cudf::test::strings_column_wrapper make_mixed_encoding_strings() +{ + std::vector strings(num_rows); + for (cudf::size_type i = 0; i < num_rows; ++i) { + bool const low_cardinality_region = i < num_rows / 2; + int const value = low_cardinality_region ? (i % 16) : i; + strings[i] = make_value_string(value); + } + return cudf::test::strings_column_wrapper(strings.begin(), strings.end()); +} + +// Like `write_parquet`, but with an ADAPTIVE dictionary policy and a caller-supplied dictionary +// budget so the writer falls back to PLAIN encoding for row groups whose dictionary exceeds it. +void write_parquet_adaptive(cudf::table_view const& input, + std::string const& filepath, + size_t max_dict_size) +{ + auto const options = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, input) + .dictionary_policy(cudf::io::dictionary_policy::ADAPTIVE) + .max_dictionary_size(max_dict_size) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .row_group_size_rows(row_group_size) + .max_page_fragment_size(row_group_size) + .build(); + cudf::io::write_parquet(options); +} + +} // namespace + +struct ParquetReaderDictTest : public cudf::test::BaseFixture {}; + +// A flat string column that is fully dictionary-encoded in every row group should be returned +// as a DICTIONARY32 column when `output_dict_columns` is enabled, and the decoded keys +// should match the original input. +TEST_F(ParquetReaderDictTest, FlatStringDictTranscode) +{ + auto input_col = make_low_cardinality_strings(); + + auto const input_tbl = cudf::table_view{{input_col}}; + auto const filepath = temp_env->get_temp_filepath("FlatStringDictTranscode.parquet"); + write_parquet(input_tbl, filepath); + + auto const dict_input = cudf::dictionary::encode(input_col); + auto const dict_input_view = cudf::dictionary_column_view(dict_input->view()); + auto const decoded_input = cudf::dictionary::decode(dict_input_view); + + auto const read_table = read_parquet_as_dict(filepath).tbl; + ASSERT_EQ(read_table->num_rows(), num_rows); + ASSERT_EQ(read_table->num_columns(), 1); + + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::DICTIONARY32) + << "Expected the reader to produce a DICTIONARY32 column when output_dict_columns is on"; + + cudf::dictionary_column_view dict_read_view(read_col); + auto const decoded_read = cudf::dictionary::decode(dict_read_view); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(input_col, decoded_read->view()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(decoded_input->view(), decoded_read->view()); +} + +// When the option is not set, the reader should still produce a plain STRING column, regardless +// of whether the source file is fully dictionary-encoded. +TEST_F(ParquetReaderDictTest, FlatStringNoTranscodeByDefault) +{ + auto input_col = make_low_cardinality_strings(); + + auto const input_tbl = cudf::table_view{{input_col}}; + auto const filepath = temp_env->get_temp_filepath("FlatStringNoTranscodeByDefault.parquet"); + write_parquet(input_tbl, filepath); + + auto const read_opts = + cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}).build(); + auto const read_table = cudf::io::read_parquet(read_opts).tbl; + + ASSERT_EQ(read_table->num_columns(), 1); + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::STRING); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(input_col, read_col); +} + +// List is not eligible for Parquet-dictionary → DICTIONARY32 transcode (flat string columns +// only). With `output_dict_columns` enabled, the reader still round-trips as LIST. +TEST_F(ParquetReaderDictTest, ListOfStringsDictEncodedWithTryOutputDictOption) +{ + auto list_col = make_low_cardinality_lists_of_strings(); + + auto const input_tbl = cudf::table_view{{list_col->view()}}; + auto const filepath = + temp_env->get_temp_filepath("ListOfStringsDictEncodedWithTryOutputDictOption.parquet"); + write_parquet(input_tbl, filepath); + + auto const read_table = read_parquet_as_dict(filepath).tbl; + ASSERT_EQ(read_table->num_rows(), input_tbl.num_rows()); + ASSERT_EQ(read_table->num_columns(), 1); + + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::LIST) + << "List must remain LIST when output_dict_columns is on (transcode is flat-only)"; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(list_col->view(), read_col); +} + +// Edge case: empty input. A zero-row flat STRING column must round-trip through the transcode +// path without error and reproduce the empty input (whether it comes back as STRING or as an +// empty DICTIONARY32 via the best-effort fallback encode). +TEST_F(ParquetReaderDictTest, EmptyFlatStringDictTranscode) +{ + std::vector const empty; + auto const input_col = cudf::test::strings_column_wrapper(empty.begin(), empty.end()); + + auto const input_tbl = cudf::table_view{{input_col}}; + auto const filepath = temp_env->get_temp_filepath("EmptyFlatStringDictTranscode.parquet"); + + // Write directly: the chunked row-group loop in `write_parquet` would skip a zero-row table. + auto const write_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, input_tbl) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .build(); + cudf::io::write_parquet(write_opts); + + auto const read_table = read_parquet_as_dict(filepath).tbl; + ASSERT_EQ(read_table->num_rows(), 0); + ASSERT_EQ(read_table->num_columns(), 1); + + auto const read_col = read_table->view().column(0); + if (read_col.type().id() == cudf::type_id::DICTIONARY32) { + // An empty DICTIONARY32 has no keys, so `cudf::dictionary::decode` returns a type-EMPTY empty + // column (there is no key type to recover) rather than an empty STRING. Just confirm the + // round-trip stays empty. + auto const decoded = cudf::dictionary::decode(cudf::dictionary_column_view(read_col)); + EXPECT_EQ(decoded->size(), 0); + } else { + ASSERT_EQ(read_col.type().id(), cudf::type_id::STRING); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(input_col, read_col); + } +} + +// Edge case: sliced input. Writing a sliced (non-zero offset, reduced size) flat STRING column +// must transcode correctly -- the reader's DICTIONARY32 output must decode back to exactly the +// sliced rows (including nulls), not the underlying full column. +TEST_F(ParquetReaderDictTest, SlicedFlatStringDictTranscode) +{ + auto const full_col = make_low_cardinality_strings(); + + // Interior slice so the view carries a non-zero offset and a reduced size, spanning multiple + // row groups to also exercise the per-row-group key concatenation / index remapping path. + auto const slice_start = row_group_size + 7; + auto const slice_end = num_rows - 13; + auto const sliced = + cudf::slice(static_cast(full_col), {slice_start, slice_end}).front(); + + auto const input_tbl = cudf::table_view{{sliced}}; + auto const filepath = temp_env->get_temp_filepath("SlicedFlatStringDictTranscode.parquet"); + write_parquet(input_tbl, filepath); + + auto const read_table = read_parquet_as_dict(filepath).tbl; + ASSERT_EQ(read_table->num_rows(), slice_end - slice_start); + ASSERT_EQ(read_table->num_columns(), 1); + + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::DICTIONARY32) + << "Expected a DICTIONARY32 column for a fully dict-encoded sliced string input"; + + auto const decoded_read = cudf::dictionary::decode(cudf::dictionary_column_view(read_col)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(sliced, decoded_read->view()); +} + +// Non-happy test (Fast path ineligible): a filter combined with `output_dict_columns`. A filter +// forces the direct transcode fast path off (predicates evaluate on materialized STRING columns); +// `finalize_output` still encodes the surviving rows to DICTIONARY32 after the filter is applied. +// The key column is projected through unchanged. +TEST_F(ParquetReaderDictTest, FilterWithOutputDictColumns) +{ + auto key_col = make_int_key_column(); + auto str_col = make_low_cardinality_strings(); + + auto const input_tbl = cudf::table_view{{key_col, str_col}}; + auto const filepath = temp_env->get_temp_filepath("FilterWithOutputDictColumns.parquet"); + write_parquet(input_tbl, filepath); + + // Filter: key column (col 0) >= num_rows / 2. + auto literal_value = cudf::numeric_scalar(num_rows / 2); + auto literal = cudf::ast::literal(literal_value); + auto col_ref = cudf::ast::column_reference(0); + auto filter_expr = cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref, literal); + + // Expected result: apply the same predicate to the input table on host-visible data. + auto const predicate = cudf::compute_column(input_tbl, filter_expr); + auto const expected = cudf::apply_boolean_mask(input_tbl, predicate->view()); + ASSERT_LT(expected->num_rows(), num_rows) << "filter must remove some rows to be meaningful"; + + auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .output_dict_columns(true) + .filter(filter_expr) + .build(); + auto const read_table = cudf::io::read_parquet(read_opts).tbl; + ASSERT_EQ(read_table->num_columns(), 2); + ASSERT_EQ(read_table->num_rows(), expected->num_rows()); + + // Key column: unchanged INT32. + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected->view().column(0), read_table->view().column(0)); + + // String column: DICTIONARY32 via the post-filter fallback encode; decodes to the surviving rows. + auto const read_str = read_table->view().column(1); + ASSERT_EQ(read_str.type().id(), cudf::type_id::DICTIONARY32); + auto const decoded = cudf::dictionary::decode(cudf::dictionary_column_view(read_str)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected->view().column(1), decoded->view()); +} + +// Non-happy test (Fast path ineligible): a column whose chunks mix dictionary and PLAIN data pages. +// cuDF decides dictionary use per column-chunk, so a small dictionary budget over mixed-cardinality +// data yields some dictionary-encoded row groups and some PLAIN-encoded ones. The PLAIN pages +// disqualify the column from the fast path; `output_dict_columns` still produces a correct +// DICTIONARY32 via the fallback encode. +TEST_F(ParquetReaderDictTest, MixedDictAndPlainPagesFallback) +{ + auto input_col = make_mixed_encoding_strings(); + + auto const input_tbl = cudf::table_view{{input_col}}; + auto const filepath = temp_env->get_temp_filepath("MixedDictAndPlainPagesFallback.parquet"); + write_parquet_adaptive(input_tbl, filepath, /*max_dict_size=*/4 * 1024); + + auto const read_table = read_parquet_as_dict(filepath).tbl; + ASSERT_EQ(read_table->num_rows(), num_rows); + ASSERT_EQ(read_table->num_columns(), 1); + + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::DICTIONARY32) + << "output_dict_columns must still yield DICTIONARY32 via the fallback encode"; + auto const decoded = cudf::dictionary::decode(cudf::dictionary_column_view(read_col)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(input_col, decoded->view()); +} + +// Non-happy test (Fast path ineligible): `skip_rows` / `num_rows`. Custom row bounds force the fast +// path off; the fallback still emits a DICTIONARY32 that must decode to exactly the requested row +// window. +TEST_F(ParquetReaderDictTest, SkipRowsNumRowsDictTranscode) +{ + auto input_col = make_low_cardinality_strings(); + + auto const input_tbl = cudf::table_view{{input_col}}; + auto const filepath = temp_env->get_temp_filepath("SkipRowsNumRowsDictTranscode.parquet"); + write_parquet(input_tbl, filepath); + + cudf::size_type const skip = row_group_size + 25; + cudf::size_type const rows = 2 * row_group_size + 40; + + auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .output_dict_columns(true) + .skip_rows(skip) + .num_rows(rows) + .build(); + auto const read_table = cudf::io::read_parquet(read_opts).tbl; + ASSERT_EQ(read_table->num_columns(), 1); + ASSERT_EQ(read_table->num_rows(), rows); + + auto const read_col = read_table->view().column(0); + ASSERT_EQ(read_col.type().id(), cudf::type_id::DICTIONARY32); + auto const decoded = cudf::dictionary::decode(cudf::dictionary_column_view(read_col)); + + auto const expected = + cudf::slice(static_cast(input_col), {skip, skip + rows}).front(); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, decoded->view()); +} + +// Non-happy test (Fast path ineligible): `chunked_parquet_reader`. A chunked read sets an +// output-chunk byte limit, which disables the fast path. Each chunk must come back as DICTIONARY32 +// via the fallback; reassembling the decoded chunks must reproduce the original column. +TEST_F(ParquetReaderDictTest, ChunkedReadDictTranscode) +{ + auto input_col = make_low_cardinality_strings(); + + auto const input_tbl = cudf::table_view{{input_col}}; + auto const filepath = temp_env->get_temp_filepath("ChunkedReadDictTranscode.parquet"); + write_parquet(input_tbl, filepath); + + auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .output_dict_columns(true) + .build(); + // Small byte limit so the read is split across multiple output chunks. + auto reader = cudf::io::chunked_parquet_reader(/*chunk_read_limit=*/16 * 1024, read_opts); + + std::vector> decoded_chunks; + cudf::size_type total_rows = 0; + int num_chunks = 0; + while (reader.has_next()) { + auto chunk = reader.read_chunk(); + ASSERT_EQ(chunk.tbl->num_columns(), 1); + auto const read_col = chunk.tbl->view().column(0); + if (read_col.size() == 0) { continue; } + ASSERT_EQ(read_col.type().id(), cudf::type_id::DICTIONARY32); + decoded_chunks.push_back(cudf::dictionary::decode(cudf::dictionary_column_view(read_col))); + total_rows += read_col.size(); + ++num_chunks; + } + ASSERT_EQ(total_rows, num_rows); + EXPECT_GT(num_chunks, 1) << "byte limit should split the read into multiple chunks"; + + std::vector views; + views.reserve(decoded_chunks.size()); + for (auto const& c : decoded_chunks) { + views.push_back(c->view()); + } + auto const combined = cudf::concatenate(views); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(input_col, combined->view()); +} + +// Non-happy test (Fast path ineligible): a multi-column table mixing eligible and ineligible +// columns. Only the flat string column is transcoded to DICTIONARY32; the LIST column stays +// LIST (flat-only transcode) and the INT32 column is untouched (non-string). This also exercises +// the output-buffer indexing when an eligible column is preceded/followed by columns of differing +// nesting. +TEST_F(ParquetReaderDictTest, MultiColumnMixedEligibility) +{ + auto str_col = make_low_cardinality_strings(); // eligible -> DICTIONARY32 + auto list_col = make_low_cardinality_lists_of_strings(); // ineligible -> LIST + auto key_col = make_int_key_column(); // non-string -> INT32 + + auto const input_tbl = cudf::table_view{{str_col, list_col->view(), key_col}}; + auto const filepath = temp_env->get_temp_filepath("MultiColumnMixedEligibility.parquet"); + write_parquet(input_tbl, filepath); + + auto const read_table = read_parquet_as_dict(filepath).tbl; + ASSERT_EQ(read_table->num_rows(), num_rows); + ASSERT_EQ(read_table->num_columns(), 3); + + // Flat string column: transcoded to DICTIONARY32. + auto const read_str = read_table->view().column(0); + ASSERT_EQ(read_str.type().id(), cudf::type_id::DICTIONARY32); + auto const decoded_str = cudf::dictionary::decode(cudf::dictionary_column_view(read_str)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(str_col, decoded_str->view()); + + // List column: not eligible, remains LIST. + auto const read_list = read_table->view().column(1); + ASSERT_EQ(read_list.type().id(), cudf::type_id::LIST) + << "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. + 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); +}