diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 428878e94566..fae5d5a7ebb1 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -550,6 +550,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 6e48e5a061a9..52de9f197502 100644 --- a/cpp/include/cudf/io/parquet.hpp +++ b/cpp/include/cudf/io/parquet.hpp @@ -103,6 +103,8 @@ class parquet_reader_options { type_id _decimal_width{type_id::EMPTY}; // Whether to use JIT compilation for filtering bool _use_jit_filter = false; + // Best-effort: try to output DICTIONARY32 columns for fully dict-encoded string columns + bool _try_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; @@ -298,6 +300,18 @@ class parquet_reader_options { */ [[nodiscard]] bool is_enabled_case_sensitive_names() const { return _case_sensitive_names; } + /** + * @brief Returns whether the reader should try to output DICTIONARY32 columns. + * + * When true, the reader may output DICTIONARY32 columns for fully dict-encoded + * string columns instead of fully decoded STRING columns. A DICTIONARY32 column + * consists of an INT32 indices child and a STRING keys child. + * Best-effort: falls back to STRING if the column has mixed encoding. + * + * @return `true` if the reader should try to output DICTIONARY32 columns + */ + [[nodiscard]] bool is_enabled_try_output_dict_columns() const { return _try_output_dict_columns; } + /** * @brief Set a new source location * @@ -533,6 +547,13 @@ class parquet_reader_options { * @param val Boolean indicating whether to enable case-sensitive matching. */ void enable_case_sensitive_names(bool val) { _case_sensitive_names = val; } + + /** + * @brief Sets to enable/disable trying to output DICTIONARY32 columns. + * + * @param val Boolean indicating whether to try to output DICTIONARY32 columns + */ + void enable_try_output_dict_columns(bool val) { _try_output_dict_columns = val; } }; /** @@ -796,6 +817,18 @@ class parquet_reader_options_builder { return *this; } + /** + * @brief Sets to enable/disable trying to output DICTIONARY32 columns. + * + * @param val Boolean value whether to try to output DICTIONARY32 columns + * @return this for chaining + */ + parquet_reader_options_builder& try_output_dict_columns(bool val) + { + options._try_output_dict_columns = val; + return *this; + } + /** * @brief move parquet_reader_options member once it's built. */ diff --git a/cpp/src/io/parquet/decode_fixed.cu b/cpp/src/io/parquet/decode_fixed.cu index e622cfbe3df3..6cb240ae8a27 100644 --- a/cpp/src/io/parquet/decode_fixed.cu +++ b/cpp/src/io/parquet/decode_fixed.cu @@ -81,6 +81,49 @@ __device__ static void scan_block_exclusive_sum( } } +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->col.max_nesting_depth - 1; + auto const data_out = s->nesting_info[leaf_level_index].data_out; + + int const skipped_leaf_values = s->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->first_row; + } else { + int dst_pos = sb->nz_idx[rolling_index(thread_pos)]; + if constexpr (!has_lists_t) { dst_pos -= s->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; + *dst = sb->dict_idx[rolling_index(src_pos)]; + } + + 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,7 +949,18 @@ 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) || + (kernel_mask_t == decode_kernel_mask::DICT_INT32_NESTED) || + (kernel_mask_t == decode_kernel_mask::DICT_INT32_LIST); +} + +template +CUDF_HOST_DEVICE constexpr bool is_dict_int32_output() +{ + return (kernel_mask_t == decode_kernel_mask::DICT_INT32) || + (kernel_mask_t == decode_kernel_mask::DICT_INT32_NESTED) || + (kernel_mask_t == decode_kernel_mask::DICT_INT32_LIST); } template @@ -926,7 +980,8 @@ CUDF_HOST_DEVICE constexpr bool has_nesting() (kernel_mask_t == decode_kernel_mask::BYTE_STREAM_SPLIT_FIXED_WIDTH_NESTED) || (kernel_mask_t == decode_kernel_mask::STRING_NESTED) || (kernel_mask_t == decode_kernel_mask::STRING_DICT_NESTED) || - (kernel_mask_t == decode_kernel_mask::STRING_STREAM_SPLIT_NESTED); + (kernel_mask_t == decode_kernel_mask::STRING_STREAM_SPLIT_NESTED) || + (kernel_mask_t == decode_kernel_mask::DICT_INT32_NESTED); } template @@ -938,7 +993,8 @@ CUDF_HOST_DEVICE constexpr bool has_lists() (kernel_mask_t == decode_kernel_mask::BYTE_STREAM_SPLIT_FIXED_WIDTH_LIST) || (kernel_mask_t == decode_kernel_mask::STRING_LIST) || (kernel_mask_t == decode_kernel_mask::STRING_DICT_LIST) || - (kernel_mask_t == decode_kernel_mask::STRING_STREAM_SPLIT_LIST); + (kernel_mask_t == decode_kernel_mask::STRING_STREAM_SPLIT_LIST) || + (kernel_mask_t == decode_kernel_mask::DICT_INT32_LIST); } template @@ -988,6 +1044,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(); @@ -1177,7 +1234,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->col.column_string_offset_base + page_string_offset_indices[page_idx]; string_output_offset = @@ -1206,10 +1266,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->col.max_nesting_depth - 1]; return ni.valid_map_offset - init_valid_map_offset; @@ -1371,6 +1435,15 @@ 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; + case decode_kernel_mask::DICT_INT32_NESTED: + launch_kernel(int_tag_t<128>{}, kernel_tag_t{}); + break; + case decode_kernel_mask::DICT_INT32_LIST: + 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/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index abea80f9d764..13c4f562f62e 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -221,7 +221,10 @@ 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 + DICT_INT32_NESTED = (1 << 27), // Run decode kernel for nested dict string → INT32 indices + DICT_INT32_LIST = (1 << 28), // Run decode kernel for list dict string → INT32 indices }; constexpr uint32_t STRINGS_MASK_NON_DELTA = BitOr(decode_kernel_mask::STRING, diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index b6a1d3d767e6..e5d249b9841f 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -13,7 +13,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -262,6 +264,21 @@ 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 dict-index-as-int32 decoder for nested columns + if (BitAnd(kernel_mask, decode_kernel_mask::DICT_INT32_NESTED) != 0) { + decode_data(decode_kernel_mask::DICT_INT32_NESTED); + } + + // launch dict-index-as-int32 decoder for list columns + if (BitAnd(kernel_mask, decode_kernel_mask::DICT_INT32_LIST) != 0) { + decode_data(decode_kernel_mask::DICT_INT32_LIST); + } + // launch delta byte array decoder if (BitAnd(kernel_mask, decode_kernel_mask::DELTA_BYTE_ARRAY) != 0) { decode_delta_byte_array(subpass.pages, @@ -513,11 +530,20 @@ reader_impl::reader_impl(std::size_t chunk_read_limit, options.get_num_bytes(), options.get_row_groups(), options.is_enabled_use_jit_filter(), - options.is_enabled_case_sensitive_names()}, + options.is_enabled_case_sensitive_names(), + options.is_enabled_try_output_dict_columns()}, _sources{std::move(sources)}, _output_chunk_read_limit{chunk_read_limit}, _input_pass_read_limit{pass_read_limit} { + // Direct parquet-dict → DICTIONARY32 transcode currently only supports single-pass, non-chunked + // reads. Splitting rowgroups across passes/subpasses would require aligning dictionary keys + // across passes, which we don't support yet. + CUDF_EXPECTS( + not _options.try_output_dict_columns or (chunk_read_limit == 0 and pass_read_limit == 0), + "try_output_dict_columns is only supported for single-pass reads; it cannot be combined " + "with a non-zero chunk_read_limit or pass_read_limit."); + // 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"); @@ -694,6 +720,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(); + // 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 @@ -716,6 +748,11 @@ table_with_metadata reader_impl::read_chunk_internal(read_mode mode) // Allocate memory buffers for the output columns. allocate_columns(mode, read_info.skip_rows, read_info.num_rows); + // Zero-init the INT32 index buffers of dict-transcoded columns before launching decode, so + // that null positions (which the DICT_INT32 kernel does not write to) carry well-defined + // indices in the produced DICTIONARY32 output. + zero_init_dict_transcoded_index_buffers(); + // Parse data into the output buffers. decode_page_data(mode, read_info.skip_rows, read_info.num_rows); @@ -743,6 +780,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 and shifting + // per-chunk indices so the concatenation refers 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); @@ -878,6 +921,22 @@ table_with_metadata reader_impl::finalize_output(read_mode mode, apply_decimal_width_cast(out_columns); + // When the user requested DICTIONARY32 output for flat string columns, the direct transcode + // fast path in `prepare_dict_transcode`/`assemble_dict_transcoded_columns` has already + // assembled DICTIONARY32 columns for all *eligible* flat STRING columns (i.e. those whose + // chunks were fully dictionary-encoded). For columns that were *not* eligible (e.g. chunks + // with mixed or non-dictionary encodings, nested schemas, or columns added as empty columns + // above), fall back to a post-hoc `dictionary::detail::encode` so the user still gets a + // DICTIONARY32 column from every flat string column in the output table. + if (_options.try_output_dict_columns) { + for (auto& col : out_columns) { + if (col and col->type().id() == type_id::STRING) { + col = cudf::dictionary::detail::encode( + col->view(), data_type{type_id::INT32}, _stream, _mr); + } + } + } + if (!_output_metadata) { populate_metadata(out_metadata); // Finally, save the output table metadata into `_output_metadata` for reuse next time. diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index e96a3f16bb55..30710870c3eb 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -189,6 +189,41 @@ 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` so that the + * subsequent allocate/decode path produces INT32 indices for eligible columns. + * + * Must be called after `prepare_data()` (so that `pass.chunks`, `pass.pages` and + * `subpass.pages` are populated on the host) and before `preprocess_chunk_strings()` / + * `allocate_columns()` / `decode_page_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`. + */ + void prepare_dict_transcode(); + + /** + * @brief Zero-initialize the INT32 output buffers of dict-transcoded columns so that null rows + * carry a well-defined dictionary index (the `DICT_INT32` kernel skips null positions). + * + * Must be called after `allocate_columns` and before `decode_page_data`. + */ + void zero_init_dict_transcoded_index_buffers(); + + /** + * @brief Assemble DICTIONARY32 output columns for input columns that were marked eligible by + * `prepare_dict_transcode`. Each chunk's INT32 indices produced by the `DICT_INT32` kernel are + * shifted by the cumulative number of keys from prior chunks, and the per-chunk keys (built + * from `pass.str_dict_index`) are concatenated into a single keys child. + * + * Non-eligible flat STRING columns are left untouched here and are expected to go through the + * post-hoc `cudf::dictionary::encode` fallback in `finalize_output`. + * + * @param out_columns The output columns vector to mutate in place. + */ + void assemble_dict_transcoded_columns(std::vector>& out_columns); + /** * @brief Copies over the relevant page mask information for the subpass */ @@ -459,6 +494,8 @@ class reader_impl { bool use_jit_filter = false; // Whether to use case-sensitive matching for column names bool case_sensitive_names = true; + // Whether to try outputting DICTIONARY32 columns for fully dict-encoded string columns + bool try_output_dict_columns = false; } _options; // name to reference converter to extract AST output filter @@ -517,6 +554,11 @@ 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 in `prepare_dict_transcode()`. Populated before decode + // and consumed in `assemble_dict_transcoded_columns()`. + 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..be2d3a9f1767 --- /dev/null +++ b/cpp/src/io/parquet/reader_impl_dict_transcode.cu @@ -0,0 +1,330 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "reader_impl.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace cudf::io::parquet::detail { + +namespace { + +// Host-side counterpart of `is_string_col` in `parquet_gpu.hpp`. Kept narrow: for direct +// Parquet-dict → DICTIONARY32 transcode we only accept pure BYTE_ARRAY columns without a +// DECIMAL logical type and without the strings-to-categorical flag. FIXED_LEN_BYTE_ARRAY is +// deliberately excluded because it is typically a binary payload. +[[nodiscard]] bool is_host_byte_array_string_chunk(ColumnChunkDesc const& chunk) +{ + if (chunk.physical_type != Type::BYTE_ARRAY) { return false; } + if (chunk.is_strings_to_cat) { return false; } + if (chunk.logical_type.has_value() and chunk.logical_type->type == LogicalType::DECIMAL) { + return false; + } + return true; +} + +// Both PLAIN_DICTIONARY (legacy) and RLE_DICTIONARY are valid encodings for data pages that +// reference a parquet dictionary page. +[[nodiscard]] bool is_dict_data_page_encoding(Encoding enc) +{ + return enc == Encoding::PLAIN_DICTIONARY or enc == Encoding::RLE_DICTIONARY; +} + +// Per-input-column eligibility flags. Each column must satisfy all of these conditions to be +// eligible for direct Parquet-dict → DICTIONARY32 transcode. +struct column_eligibility { + bool has_string_buffer = false; + bool has_any_chunk = false; + bool all_chunks_string = true; + bool all_pages_dict = true; + + [[nodiscard]] bool is_eligible() const + { + return has_string_buffer and has_any_chunk and all_chunks_string and all_pages_dict; + } +}; + +// Classify a chunk against its column's eligibility state. +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_host_byte_array_string_chunk(chunk) or chunk.num_dict_pages < 1) { + e.all_chunks_string = false; + } +} + +// Build a STRING keys column covering the dictionary entries of a single chunk of a single input +// column. `begin` points into the pass-wide `string_index_pair` buffer. +[[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() +{ + CUDF_FUNC_RANGE(); + + _dict_transcode_eligible.assign(_input_columns.size(), false); + + if (not _options.try_output_dict_columns) { return; } + if (_pass_itm_data == nullptr or _pass_itm_data->subpass == nullptr) { return; } + + auto& pass = *_pass_itm_data; + auto& subpass = *pass.subpass; + + if (pass.chunks.empty() or subpass.pages.size() == 0) { return; } + + // Determine per-input-column eligibility. 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 (PLAIN|RLE)_DICTIONARY encoding, + // - the chunk has a flat (non-list, non-nested) schema. + // + // We scan host-side `pass.chunks` and `pass.pages` here rather than `subpass.pages` because + // `subpass.pages` may be a subset. For single-pass single-subpass reads (the only configuration + // in which `try_output_dict_columns` is supported), `subpass.pages == pass.pages`. + auto const num_input_cols = _input_columns.size(); + std::vector elig(num_input_cols); + + // Seed from the output buffer type: only flat STRING columns have a single leaf buffer whose + // type is STRING and can be flipped in place. + std::for_each( + cuda::counting_iterator{0}, cuda::counting_iterator{num_input_cols}, [&](size_t i) { + auto const& input_col = _input_columns[i]; + if (input_col.nesting_depth() != 1) { return; } + auto const& out_buf = _output_buffers[input_col.nesting[0]]; + if (out_buf.type.id() == type_id::STRING) { elig[i].has_string_buffer = true; } + }); + + // Fold per-chunk info into the per-column eligibility flags. + for (auto const& chunk : pass.chunks) { + auto const col_idx = chunk.src_col_index; + if (col_idx < 0 or static_cast(col_idx) >= num_input_cols) { continue; } + 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; + if (chunk_idx < 0 or static_cast(chunk_idx) >= pass.chunks.size()) { continue; } + auto const col_idx = pass.chunks[chunk_idx].src_col_index; + if (col_idx < 0 or static_cast(col_idx) >= num_input_cols) { continue; } + if (not is_dict_data_page_encoding(page.encoding)) { elig[col_idx].all_pages_dict = false; } + } + + 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; } + + // Flip the output buffer type for eligible columns from STRING → INT32. The subsequent + // `allocate_columns` call will then allocate an INT32 buffer that the DICT_INT32 kernel can + // write directly into. + 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. Only the flat variant is + // considered here (eligibility requires `max_nesting_depth == 1`). + 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; + if (chunk_idx < 0 or static_cast(chunk_idx) >= pass.chunks.size()) { return; } + auto const col_idx = pass.chunks[chunk_idx].src_col_index; + if (col_idx < 0 or static_cast(col_idx) >= num_input_cols) { return; } + 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; + } + }); + + if (not any_rewritten) { return; } + + // Push the rewritten `kernel_mask`s back to device so subsequent decode kernels dispatch + // correctly. Then refresh the aggregated `subpass.kernel_mask` on the host by re-OR'ing all + // page kernel masks. + 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); }); + _stream.synchronize(); +} + +void reader_impl::zero_init_dict_transcoded_index_buffers() +{ + CUDF_FUNC_RANGE(); + + if (not _options.try_output_dict_columns) { return; } + if (_dict_transcode_eligible.empty()) { return; } + + // The `DICT_INT32` kernel only writes to positions with valid definition levels, leaving null + // slots untouched. Since `allocate_columns` does not zero-initialize fixed-width buffers by + // default, the INT32 output buffer for a transcoded column may contain uninitialized bytes at + // null positions. Zero them here so null rows carry a well-defined (valid) index into the + // dictionary keys -- a requirement for `cudf::dictionary::detail::concatenate` to correctly + // remap indices below. + std::for_each( + cuda::counting_iterator{0}, + cuda::counting_iterator{_input_columns.size()}, + [&](size_t i) { + if (not _dict_transcode_eligible[i]) { return; } + auto& out_buf = _output_buffers[_input_columns[i].nesting[0]]; + if (out_buf.type.id() != type_id::INT32) { return; } + if (out_buf.data() == nullptr or out_buf.size == 0) { return; } + CUDF_CUDA_TRY(cudaMemsetAsync( + out_buf.data(), 0, static_cast(out_buf.size) * sizeof(int32_t), _stream.value())); + }); +} + +void reader_impl::assemble_dict_transcoded_columns( + std::vector>& out_columns) +{ + CUDF_FUNC_RANGE(); + + if (not _options.try_output_dict_columns) { return; } + if (_dict_transcode_eligible.empty()) { return; } + if (_pass_itm_data == nullptr) { 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; } + + // 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}; + }); + + // Grab ownership of the decoded INT32 indices column; we slice it into zero-copy + // per-chunk views to feed into the per-segment dictionary column builders. + auto& indices_col = out_columns[i]; + 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); + 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"); + + // Build a DICTIONARY32 segment for every chunk. Each segment carries row-group-local + // indices into its own keys child; `cudf::detail::concatenate` rewrites the indices + // against the unified, deduplicated keys. We deep-copy each sliced view into a fresh + // offset-zero indices column before handing it to `make_dictionary_column`: the + // `make_dictionary_column(column_view, column_view, ...)` path would otherwise + // double-apply the slice's offset when constructing the internal indices child. + std::vector> dict_segments(chunk_indices.size()); + std::transform(cuda::counting_iterator{0}, + cuda::counting_iterator{chunk_indices.size()}, + dict_segments.begin(), + [&](size_t k) { + auto const chunk_idx = chunk_indices[k]; + auto const& chunk = pass.chunks[chunk_idx]; + + auto const seg_view = cudf::detail::slice( + indices_view, chunk_row_offsets[k], chunk_row_offsets[k + 1], _stream); + auto seg_indices = std::make_unique(seg_view, _stream, _mr); + + auto seg_keys = make_keys_column_from_index_pairs( + chunk.str_dict_index, chunk_key_counts[k], _stream, _mr); + + return cudf::make_dictionary_column( + std::move(seg_keys), std::move(seg_indices), _stream, _mr); + }); + + // Concatenate all segments into the final DICTIONARY32 column for this input column. + // `cudf::dictionary::detail::concatenate` deduplicates + sorts keys and recomputes indices. + if (dict_segments.size() == 1) { + out_columns[i] = std::move(dict_segments.front()); + } else { + std::vector segment_views(dict_segments.size()); + std::transform( + dict_segments.begin(), dict_segments.end(), segment_views.begin(), [](auto const& seg) { + return seg->view(); + }); + out_columns[i] = cudf::detail::concatenate(segment_views, _stream, _mr); + } + }); +} + +} // namespace cudf::io::parquet::detail diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index f13a3af71b69..e0eadec6a4b2 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -333,6 +333,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..7837b39331aa --- /dev/null +++ b/cpp/tests/io/parquet_reader_dict_test.cpp @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "parquet_common.hpp" + +#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 double null_probability = 0.1; + +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] = "str_" + std::to_string(value_dist(engine)); + valids[i] = not null_dist(engine); + } + + return cudf::test::strings_column_wrapper(strings.begin(), strings.end(), valids.begin()); +} + +void write_parquet(cudf::table_view const& input, std::string const& filepath) +{ + auto const options = + cudf::io::chunked_parquet_writer_options::builder(cudf::io::sink_info{filepath}) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .build(); + + cudf::io::chunked_parquet_writer writer(options); + for (auto offset = 0; offset < input.num_rows(); offset += row_group_size) { + auto const length = std::min(row_group_size, input.num_rows() - offset); + auto const chunk = cudf::slice(input, {offset, offset + length}); + writer.write(chunk.front()); + } + writer.close(); +} + +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}) + .try_output_dict_columns(true) + .build(); + return cudf::io::read_parquet(read_opts); +} + +} // 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 `try_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 try_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); +}