Extend Parquet reader output_dict_columns to fixed-width columns and sized indices

cpp/benchmarks/CMakeLists.txt

@@ -297,7 +297,8 @@ ConfigureNVBench(
 # * parquet reader benchmark ----------------------------------------------------------------------
 ConfigureNVBench(
   PARQUET_READER_NVBENCH io/parquet/parquet_reader_input.cpp io/parquet/parquet_reader_encoding.cpp
-  io/parquet/parquet_reader_options.cpp io/parquet/reader_common.cpp
+  io/parquet/parquet_reader_options.cpp io/parquet/parquet_reader_dict_transcode.cpp
+  io/parquet/reader_common.cpp
 )
 
 # ##################################################################################################

cpp/benchmarks/io/parquet/parquet_reader_dict_transcode.cpp

@@ -0,0 +1,79 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <benchmarks/common/generate_input.hpp>
+#include <benchmarks/common/memory_stats.hpp>
+#include <benchmarks/io/cuio_common.hpp>
+#include <benchmarks/io/nvbench_helpers.hpp>
+
+#include <cudf/io/parquet.hpp>
+#include <cudf/utilities/default_stream.hpp>
+
+#include <nvbench/nvbench.cuh>
+
+// Size of the data in the benchmark dataframe; chosen to be low enough to allow benchmarks to
+// run on most GPUs, but large enough to allow highest throughput
+constexpr std::size_t data_size = 512 << 20;
+
+// Measures reads of dictionary-encoded low-cardinality columns with and without the direct
+// Parquet-dictionary -> DICTIONARY32 transcode (`output_dict_columns`), over the column types the
+// fast path accepts: flat strings and flat fixed-width INT32/INT64/TIMESTAMP_DAYS columns.
+template <output_dict OutputDict>
+void BM_parquet_read_dict_transcode(nvbench::state& state,
+                                    nvbench::type_list<nvbench::enum_type<OutputDict>>)
+{
+  auto constexpr output_dict_columns = OutputDict == output_dict::YES;
+
+  auto const cardinality = static_cast<cudf::size_type>(state.get_int64("cardinality"));
+  auto const run_length  = static_cast<cudf::size_type>(state.get_int64("run_length"));
+
+  auto const data_types = std::vector<cudf::type_id>{cudf::type_id::INT32,
+                                                     cudf::type_id::INT64,
+                                                     cudf::type_id::TIMESTAMP_DAYS,
+                                                     cudf::type_id::STRING};
+  data_profile const profile =
+    data_profile_builder().cardinality(cardinality).avg_run_length(run_length);
+  auto const tbl  = create_random_table(data_types, table_size_bytes{data_size}, profile);
+  auto const view = tbl->view();
+
+  cuio_source_sink_pair source_sink(io_type::HOST_BUFFER);
+  auto const write_options =
+    cudf::io::parquet_writer_options::builder(source_sink.make_sink_info(), view)
+      .dictionary_policy(cudf::io::dictionary_policy::ALWAYS)
+      .build();
+  cudf::io::write_parquet(write_options);
+
+  cudf::io::parquet_reader_options read_options =
+    cudf::io::parquet_reader_options::builder(source_sink.make_source_info())
+      .output_dict_columns(output_dict_columns)
+      .build();
+
+  auto mem_stats_logger = cudf::memory_stats_logger();
+  state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().value()));
+  state.exec(
+    nvbench::exec_tag::sync | nvbench::exec_tag::timer, [&](nvbench::launch& launch, auto& timer) {
+      drop_page_cache_if_enabled(read_options.get_source().filepaths());
+      timer.start();
+      auto const result = cudf::io::read_parquet(read_options);
+      timer.stop();
+      CUDF_EXPECTS(result.tbl->num_rows() == view.num_rows(),
+                   "Benchmark did not read the entire table");
+      CUDF_EXPECTS(result.tbl->num_columns() == view.num_columns(), "Unexpected number of columns");
+    });
+
+  auto const elapsed_time = state.get_summary("nv/cold/time/gpu/mean").get_float64("value");
+  state.add_element_count(static_cast<double>(data_size) / elapsed_time, "bytes_per_second");
+  state.add_buffer_size(
+    mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage");
+  state.add_buffer_size(source_sink.size(), "encoded_file_size", "encoded_file_size");
+}
+
+NVBENCH_BENCH_TYPES(BM_parquet_read_dict_transcode,
+                    NVBENCH_TYPE_AXES(nvbench::enum_type_list<output_dict::NO, output_dict::YES>))
+  .set_name("parquet_read_dict_transcode")
+  .set_type_axes_names({"output_dict_columns"})
+  .set_min_samples(4)
+  .add_int64_axis("cardinality", {1000})
+  .add_int64_axis("run_length", {4});

cpp/include/cudf/io/parquet.hpp

@@ -343,16 +343,23 @@ class parquet_reader_options {
   }
 
   /**
-   * @brief Returns whether the reader returns flat string columns as DICTIONARY32 encoded columns
+   * @brief Returns whether the reader returns flat string and fixed-width columns as DICTIONARY32
+   * encoded columns
    *
-   * When true, the reader outputs STRING columns as DICTIONARY32 encoded columns. A DICTIONARY32
-   * column consists of an INT32 indices child and a STRING keys child.
+   * When true, the reader outputs eligible flat columns as DICTIONARY32 encoded columns: a signed
+   * integer indices child (INT8/INT16/INT32, sized to the largest row-group dictionary with the
+   * same width rule as cudf::dictionary::encode) plus a keys child of the column's logical type.
+   * Eligible are flat STRING columns and flat fixed-width columns whose physical storage is INT32
+   * or INT64 and whose decode is a plain copy (no decimal, timestamp-unit, or width conversion),
+   * when every data page of the column is dictionary encoded.
    *
    * When AST/JIT filters are set, the direct transcode fast path is disabled.
-   * String columns are materialized, then operated on by the filter. The filtered results are then
-   * encoded as DICTIONARY32 columns.
+   * String columns are materialized, then operated on by the filter, and the filtered results are
+   * encoded as DICTIONARY32 columns. Fixed-width columns participate only in the direct fast path
+   * and are returned as plain columns whenever it does not apply (filters, chunked or bounded
+   * reads, or pages that are not dictionary encoded).
    *
-   * @return `true` if the reader returns flat string columns as DICTIONARY32 encoded columns
+   * @return `true` if the reader returns eligible flat columns as DICTIONARY32 encoded columns
    */
   [[nodiscard]] bool is_enabled_output_dict_columns() const { return _output_dict_columns; }
 
@@ -651,9 +658,11 @@ class parquet_reader_options {
   void enable_prepend_row_index_column(bool val) { _prepend_row_index_column = val; }
 
   /**
-   * @brief Sets to enable/disable trying to output DICTIONARY32 columns for flat string columns.
+   * @brief Sets to enable/disable trying to output DICTIONARY32 columns for eligible flat string
+   * and fixed-width columns (see is_enabled_output_dict_columns for eligibility and fallback
+   * behavior).
    *
-   * @param val Boolean indicating whether to output DICTIONARY32 columns for flat string columns
+   * @param val Boolean indicating whether to output eligible flat columns as DICTIONARY32
    */
   void enable_output_dict_columns(bool val) { _output_dict_columns = val; }
 };
@@ -955,13 +964,17 @@ class parquet_reader_options_builder {
   }
 
   /**
-   * @brief Sets options for enabling/disabling output of DICTIONARY32 columns for flat string
-   * columns.
+   * @brief Sets options for enabling/disabling output of DICTIONARY32 columns for eligible flat
+   * string and fixed-width columns.
    *
-   * @param val Boolean value whether to output flat string columns as DICTIONARY32 encoded columns
+   * @param val Boolean value whether to output eligible flat columns as DICTIONARY32 encoded
+   * columns
    *
-   * @note When enabled, the output columns will be of type DICTIONARY32. When disabled, the output
-   * columns will be of type STRING.
+   * @note When enabled, eligible columns are returned as DICTIONARY32 with a signed integer
+   * indices child sized to the dictionary and a keys child of the logical type; string columns
+   * are always delivered as DICTIONARY32 (post-read encoded when the direct path does not apply),
+   * while fixed-width columns fall back to their plain type
+   * (see parquet_reader_options::is_enabled_output_dict_columns).
    *
    * @return this for chaining
    */

cpp/src/dictionary/detail/concatenate.cu

@@ -11,15 +11,19 @@
 #include <cudf/detail/iterator.cuh>
 #include <cudf/detail/row_operator/equality.cuh>
 #include <cudf/detail/row_operator/hashing.cuh>
+#include <cudf/detail/unary.hpp>
 #include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/dictionary/detail/concatenate.hpp>
+#include <cudf/dictionary/detail/encode.hpp>
 #include <cudf/dictionary/dictionary_column_view.hpp>
 #include <cudf/dictionary/dictionary_factories.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/memory_resource.hpp>
+#include <cudf/utilities/traits.hpp>
 #include <cudf/utilities/type_checks.hpp>
+#include <cudf/utilities/type_dispatcher.hpp>
 
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
@@ -37,6 +41,7 @@
 #include <thrust/transform_scan.h>
 
 #include <algorithm>
+#include <numeric>
 #include <vector>
 
 namespace cudf {
@@ -133,22 +138,57 @@ struct compute_children_offsets_fn {
 /**
  * @brief Functor for mapping the old indices values to the new indices values
  *        based on the new keys arrangement after concatenation
+ *
+ * @tparam IndexType Integral type of the (already widened) indices column
  */
+template <typename IndexType>
 struct map_indices_fn {
   cuda::std::span<offsets_pair const> d_offsets;
   cuda::std::span<size_type const> d_keys_remap;
   column_device_view d_indices;
 
-  __device__ size_type operator()(size_type idx) const
+  __device__ IndexType operator()(size_type idx) const
   {
-    if (d_indices.is_null(idx)) { return 0; }
+    if (d_indices.is_null(idx)) { return IndexType{0}; }
     auto cmp = [] __device__(auto const& lhs, auto const& rhs) { return lhs.second < rhs.second; };
     auto col_iter = thrust::upper_bound(
                       thrust::seq, d_offsets.begin(), d_offsets.end(), offsets_pair{0, idx}, cmp) -
                     1;
-    auto col_idx    = cuda::std::distance(d_offsets.begin(), col_iter);
-    auto key_offset = d_offsets[col_idx].first;
-    return d_keys_remap[key_offset + d_indices.element<size_type>(idx)];
+    auto col_idx         = cuda::std::distance(d_offsets.begin(), col_iter);
+    auto key_offset      = d_offsets[col_idx].first;
+    auto const old_index = static_cast<size_type>(d_indices.element<IndexType>(idx));
+    return static_cast<IndexType>(d_keys_remap[key_offset + old_index]);
+  }
+};
+
+/**
+ * @brief Dispatches map_indices_fn on the indices type so that the indices are
+ *        read and written with their actual width (INT8/INT16/INT32/...).
+ */
+struct remap_indices_fn {
+  template <typename IndexType>
+  void operator()(column_device_view const& d_indices,
+                  mutable_column_view const& output,
+                  cuda::std::span<offsets_pair const> d_offsets,
+                  cuda::std::span<size_type const> d_keys_remap,
+                  rmm::cuda_stream_view stream,
+                  rmm::device_async_resource_ref mr) const
+    requires(cudf::is_index_type<IndexType>())
+  {
+    auto policy = rmm::exec_policy_nosync(stream, mr);
+    auto iota   = cuda::counting_iterator<size_type>{0};
+    thrust::transform(policy,
+                      iota,
+                      iota + output.size(),
+                      output.begin<IndexType>(),
+                      map_indices_fn<IndexType>{d_offsets, d_keys_remap, d_indices});
+  }
+
+  template <typename IndexType, typename... Args>
+  void operator()(Args&&...) const
+    requires(!cudf::is_index_type<IndexType>())
+  {
+    CUDF_FAIL("dictionary indices must be an integral index type");
   }
 };
 
@@ -231,15 +271,31 @@ std::unique_ptr<column> concatenate(host_span<column_view const> columns,
   thrust::gather(
     policy, d_indices.begin(), d_indices.end(), all_keys_remap.begin(), final_remap.begin());
 
-  // next, concatenate the indices
+  // next, concatenate the indices.
+  // The output indices type is the widest of the input indices types, widened further if the
+  // concatenated keys no longer fit in it (e.g. two INT8 dictionaries with 200 distinct keys).
+  auto indices_type = std::accumulate(
+    columns.begin(), columns.end(), data_type{type_id::INT8}, [](data_type widest, auto cv) {
+      // an empty dictionary column may carry no children at all, but a sliced-to-empty
+      // view still has an indices child whose type participates in the selection
+      if (cv.num_children() == 0) { return widest; }
+      auto const t = dictionary_column_view(cv).indices().type();
+      return cudf::size_of(t) > cudf::size_of(widest) ? t : widest;
+    });
+  auto const needed_type = get_indices_type_for_size(keys_column->size());
+  if (cudf::size_of(needed_type) > cudf::size_of(indices_type)) { indices_type = needed_type; }
+
+  std::vector<std::unique_ptr<column>> widened_indices;  // keeps casted indices alive
   std::vector<column_view> indices_views(columns.size());
-  std::transform(columns.begin(), columns.end(), indices_views.begin(), [](auto cv) {
-    auto dict_view = dictionary_column_view(cv);
-    if (dict_view.is_empty()) {
-      return column_view{data_type{type_id::INT32}, 0, nullptr, nullptr, 0};
-    }
-    return dict_view.get_indices_annotated();  // nicely includes validity mask and view offset
-  });
+  std::transform(
+    columns.begin(), columns.end(), indices_views.begin(), [&](auto cv) -> column_view {
+      auto dict_view = dictionary_column_view(cv);
+      if (dict_view.is_empty()) { return column_view{indices_type, 0, nullptr, nullptr, 0}; }
+      auto indices = dict_view.get_indices_annotated();  // includes validity mask and view offset
+      if (indices.type() == indices_type) { return indices; }
+      widened_indices.emplace_back(cudf::detail::cast(indices, indices_type, stream, temp_mr));
+      return widened_indices.back()->view();
+    });
   auto all_indices = cudf::detail::concatenate(indices_views, stream, mr);
 
   // remap the input indices values to the new indices for the new keys order
@@ -248,9 +304,14 @@ std::unique_ptr<column> concatenate(host_span<column_view const> columns,
   auto output_view      = indices_column->mutable_view();
   auto input_view       = column_device_view::create(all_indices->view(), stream, temp_mr);
   auto children_offsets = child_offsets_fn.create_children_offsets(stream, temp_mr);
-  auto map_fn           = map_indices_fn{children_offsets, final_remap, *input_view};
-  thrust::transform(
-    policy, iota, iota + all_indices->size(), output_view.begin<size_type>(), map_fn);
+  cudf::type_dispatcher(all_indices->type(),
+                        remap_indices_fn{},
+                        *input_view,
+                        output_view,
+                        cuda::std::span<offsets_pair const>{children_offsets},
+                        cuda::std::span<size_type const>{final_remap},
+                        stream,
+                        temp_mr);
 
   // remove the bitmask from the all_indices
   auto null_count = all_indices->null_count();  // get before release()

cpp/src/io/parquet/decode_fixed.cu

@@ -83,10 +83,11 @@ __device__ static void scan_block_exclusive_sum(
 }
 
 /**
- * @brief Write a batch of decoded dictionary indices directly as INT32 output values.
+ * @brief Write a batch of decoded dictionary indices directly as signed integer output values.
  *
  * Used by the Parquet-dict → DICTIONARY32 transcode path: instead of materializing the dictionary
- * keys, the per-row dictionary indices are emitted verbatim as the INT32 indices child of the
+ * keys, the per-row dictionary indices are emitted verbatim as the signed integer indices child
+ * (INT8/INT16/INT32, per the chunk's dict_index_bytes) of the
  * output DICTIONARY32 column.
  *
  * @tparam block_size Number of threads per block
@@ -133,13 +134,22 @@ __device__ void decode_dict_indices_as_int32(
         return thread_pos;
       }();
 
-      auto* dst           = reinterpret_cast<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)];
+      // string chunks index a string_index_pair table; fixed-width chunks index the raw
+      // dictionary page, whose entries are the physical value width
+      auto const entry_size = s->setup.col.physical_type == Type::BYTE_ARRAY
+                                ? sizeof(string_index_pair)
+                                : static_cast<size_t>(s->output_cvt.dtype_len_in);
+      auto const num_keys   = static_cast<uint32_t>(s->stream.dict_size / entry_size);
+      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;
+        // the index width follows dictionary::get_indices_type_for_size for the key count
+        switch (s->output_cvt.dtype_len) {
+          case 1: reinterpret_cast<int8_t*>(data_out)[dst_pos] = static_cast<int8_t>(idx); break;
+          case 2: reinterpret_cast<int16_t*>(data_out)[dst_pos] = static_cast<int16_t>(idx); break;
+          default: reinterpret_cast<int32_t*>(data_out)[dst_pos] = idx; break;
+        }
       }
     }
 
@@ -1306,7 +1316,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8)
   if constexpr (has_strings_t || has_lists_t || is_dict_int32_t) {
     if (process_nulls) {
       uint32_t const dtype_len = [&]() -> uint32_t {
-        if constexpr (is_dict_int32_t) { return sizeof(int32_t); }
+        if constexpr (is_dict_int32_t) { return s->setup.col.dict_index_bytes; }
         if constexpr (has_strings_t) { return sizeof(cudf::size_type); }
         return s->output_cvt.dtype_len;
       }();

cpp/src/io/parquet/page_decode.cuh

@@ -1257,6 +1257,12 @@ inline __device__ bool setup_local_page_info(auto* const s,
         } else if (data_type == Type::INT96) {
           s->output_cvt.dtype_len = 8;  // Convert to 64-bit timestamp
         }
+        // Parquet-dict -> DICTIONARY32 transcode: the output holds dictionary indices, so both
+        // the per-page output offset and the value width follow the index width, not the
+        // logical type width. (String pages coincidentally match via the string special case.)
+        if (s->setup.page.kernel_mask == decode_kernel_mask::DICT_INT32) {
+          s->output_cvt.dtype_len = s->setup.col.dict_index_bytes;
+        }
       }
 
       // during the decoding step we need to offset the global output buffers
@@ -1301,7 +1307,10 @@ inline __device__ bool setup_local_page_info(auto* const s,
                   idx < max_depth - 1 ? sizeof(cudf::size_type) : s->output_cvt.dtype_len;
                 // if this is a string column, then dtype_len is a lie. data will be offsets rather
                 // than (ptr,len) tuples.
-                if (is_string_col(s->setup.col)) { len = sizeof(cudf::size_type); }
+                if (is_string_col(s->setup.col) &&
+                    s->setup.page.kernel_mask != decode_kernel_mask::DICT_INT32) {
+                  len = sizeof(cudf::size_type);
+                }
                 nesting_info->data_out += (output_offset * len);
               }
               if (nesting_info->string_out != nullptr) {