Transcode parquet to cuDF dictionaries for flat STRING columns

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

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

cpp/src/dictionary/detail/concatenate.cu

@@ -175,6 +175,9 @@ std::unique_ptr<column> concatenate(host_span<column_view const> 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());

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 <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->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<state_buf::nz_buf_size>(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<int32_t*>(data_out) + dst_pos;
+      auto const num_keys = static_cast<uint32_t>(s->stream.dict_size / sizeof(string_index_pair));
+      auto const idx      = sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)];
+      if (idx >= num_keys) {
+        s->set_error_code(decode_error::DATA_STREAM_OVERRUN);
+      } else {
+        *dst = idx;
+      }
+    }
+
+    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,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 <decode_kernel_mask kernel_mask_t>
 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 <decode_kernel_mask kernel_mask_t>
+CUDF_HOST_DEVICE constexpr bool is_dict_int32_output()
+{
+  return (kernel_mask_t == decode_kernel_mask::DICT_INT32);
 }
 
 template <decode_kernel_mask kernel_mask_t>
@@ -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 <decode_kernel_mask kernel_mask_t>
 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 <decode_kernel_mask kernel_mask_t>
 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<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>();
@@ -1170,7 +1269,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->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<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;
     default: CUDF_EXPECTS(false, "Kernel type not handled by this function"); break;
   }
 }

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;

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.