🚧 [DO NOT MERGE] Transcode parquet to cuDF dictionaries for STRING

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

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.
    */

cpp/src/io/parquet/decode_fixed.cu

@@ -81,6 +81,49 @@ __device__ static void scan_block_exclusive_sum(
   }
 }
 
+template <int block_size, bool has_lists_t, copy_mode copy_mode_t, typename state_buf>
+__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<state_buf::nz_buf_size>(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<int32_t*>(data_out) + dst_pos;
+      *dst      = sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)];
+    }
+
+    pos += batch_size;
+    __syncthreads();
+  }
+}
+
 template <int block_size, bool has_lists_t, copy_mode copy_mode_t, typename state_buf>
 __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 <decode_kernel_mask kernel_mask_t>
+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 <decode_kernel_mask kernel_mask_t>
@@ -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 <decode_kernel_mask kernel_mask_t>
@@ -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 <decode_kernel_mask kernel_mask_t>
@@ -988,6 +1044,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8)
   constexpr bool split_decode_t = is_split_decode<kernel_mask_t>();
   constexpr bool has_strings_t =
     (static_cast<uint32_t>(kernel_mask_t) & STRINGS_MASK_NON_DELTA) != 0;
+  constexpr bool is_dict_int32_t = is_dict_int32_output<kernel_mask_t>();
 
   constexpr int rolling_buf_size    = decode_block_size_t * 2;
   constexpr int rle_run_buffer_size = rle_stream_required_run_buffer_size<decode_block_size_t>();
@@ -1177,7 +1234,10 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8)
     }
 
     auto decode_values = [&]<copy_mode copy_mode_t>() {
-      if constexpr (has_strings_t) {
+      if constexpr (is_dict_int32_t) {
+        decode_dict_indices_as_int32<decode_block_size_t, has_lists_t, copy_mode_t>(
+          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<PageInfo> pages,
     case decode_kernel_mask::STRING_STREAM_SPLIT_LIST:
       launch_kernel(int_tag_t<128>{}, kernel_tag_t<decode_kernel_mask::STRING_STREAM_SPLIT_LIST>{});
       break;
+    case decode_kernel_mask::DICT_INT32:
+      launch_kernel(int_tag_t<128>{}, kernel_tag_t<decode_kernel_mask::DICT_INT32>{});
+      break;
+    case decode_kernel_mask::DICT_INT32_NESTED:
+      launch_kernel(int_tag_t<128>{}, kernel_tag_t<decode_kernel_mask::DICT_INT32_NESTED>{});
+      break;
+    case decode_kernel_mask::DICT_INT32_LIST:
+      launch_kernel(int_tag_t<128>{}, kernel_tag_t<decode_kernel_mask::DICT_INT32_LIST>{});
+      break;
     default: CUDF_EXPECTS(false, "Kernel type not handled by this function"); break;
   }
 }

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,

cpp/src/io/parquet/reader_impl.cpp

@@ -13,7 +13,9 @@
 #include <cudf/detail/structs/utilities.hpp>
 #include <cudf/detail/transform.hpp>
 #include <cudf/detail/utilities/stream_pool.hpp>
+#include <cudf/dictionary/detail/encode.hpp>
 #include <cudf/io/parquet_schema.hpp>
+#include <cudf/logger.hpp>
 #include <cudf/null_mask.hpp>
 #include <cudf/stream_compaction.hpp>
 #include <cudf/strings/detail/utilities.hpp>
@@ -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.