Add multifile dictionary pruning support for hybrid scan parquet reader

cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp

@@ -167,6 +167,35 @@ class hybrid_scan_multifile {
   secondary_filters_byte_ranges(cudf::host_span<std::vector<size_type> const> row_group_indices,
                                 parquet_reader_options const& options) const;
 
+  /**
+   * @brief Get byte ranges of column chunk dictionary pages for row group pruning
+   *
+   * @param row_group_indices Input row group indices, one per source
+   * @param options Parquet reader options
+   * @return Vector of byte ranges of column chunks with dictionary pages subject to the filter
+   *         predicate, ordered source-major then row-group then dictionary column
+   */
+  [[nodiscard]] std::vector<byte_range_info> dictionary_pages_byte_ranges(
+    cudf::host_span<std::vector<size_type> const> row_group_indices,
+    parquet_reader_options const& options) const;
+
+  /**
+   * @brief Filter the row groups using column chunk dictionary pages
+   *
+   * @param dictionary_page_data Device spans of dictionary page data of column chunks with an
+   *                             (in)equality predicate, ordered to match the dictionary page byte
+   *                             ranges returned by `dictionary_pages_byte_ranges`
+   * @param row_group_indices Input row group indices, one per source
+   * @param options Parquet reader options
+   * @param stream CUDA stream used for device memory operations and kernel launches
+   * @return Filtered per-source row group indices (one inner vector per source)
+   */
+  [[nodiscard]] std::vector<std::vector<size_type>> filter_row_groups_with_dictionary_pages(
+    cudf::host_span<cudf::device_span<uint8_t const> const> dictionary_page_data,
+    cudf::host_span<std::vector<size_type> const> row_group_indices,
+    parquet_reader_options const& options,
+    rmm::cuda_stream_view stream) const;
+
  private:
   std::unique_ptr<detail::hybrid_scan_reader_impl> _impl;
 };

cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp

@@ -477,38 +477,46 @@ std::vector<byte_range_info> aggregate_reader_metadata::get_dictionary_page_byte
   // Flag to check if we have at least one valid dictionary page
   auto have_dictionary_pages = false;
 
+  // Cache each dictionary column's chunk offset across sources and row groups
+  std::vector<std::optional<size_type>> colchunk_offsets(dictionary_col_schemas.size());
+
   // For all sources
   std::for_each(
     cuda::counting_iterator<std::size_t>{0},
     cuda::counting_iterator{row_group_indices.size()},
     [&](auto const src_index) {
       // Get all row group indices in the data source
       auto const& rg_indices = row_group_indices[src_index];
-      std::optional<size_type> colchunk_iter_offset{};
       // For all row groups
       std::for_each(rg_indices.cbegin(), rg_indices.cend(), [&](auto const rg_index) {
-        auto const& row_group = per_file_metadata[src_index].row_groups[rg_index];
-        // For all column chunks
+        auto const& row_group     = per_file_metadata[src_index].row_groups[rg_index];
+        auto const num_col_chunks = static_cast<size_type>(row_group.columns.size());
+        // For all dictionary column chunks (kept inner to preserve the source-major emission order)
         std::for_each(
-          dictionary_col_schemas.begin(),
-          dictionary_col_schemas.end(),
-          [&](auto const& schema_idx) {
-            // Get the column chunk iterator
-            if (not colchunk_iter_offset.has_value() or
-                row_group.columns[colchunk_iter_offset.value()].schema_idx != schema_idx) {
-              auto const& colchunk_iter = std::find_if(
-                row_group.columns.begin(), row_group.columns.end(), [schema_idx](auto const& col) {
-                  return col.schema_idx == schema_idx;
-                });
-              CUDF_EXPECTS(colchunk_iter != row_group.columns.end(),
-                           "Column chunk with schema index " + std::to_string(schema_idx) +
+          cuda::counting_iterator<std::size_t>{0},
+          cuda::counting_iterator{dictionary_col_schemas.size()},
+          [&](auto const col) {
+            // Map the schema index to this source (for `allow_mismatched_pq_schemas`)
+            auto const mapped_schema_idx =
+              map_schema_index(dictionary_col_schemas[col], static_cast<int>(src_index));
+            auto& colchunk_offset    = colchunk_offsets[col];
+            auto const cached_offset = colchunk_offset.value_or(-1);
+            if (cached_offset < 0 or cached_offset >= num_col_chunks or
+                row_group.columns[cached_offset].schema_idx != mapped_schema_idx) {
+              auto const it = std::find_if(
+                row_group.columns.begin(),
+                row_group.columns.end(),
+                [mapped_schema_idx](auto const& c) { return c.schema_idx == mapped_schema_idx; });
+              CUDF_EXPECTS(it != row_group.columns.end(),
+                           "Column chunk with schema index " + std::to_string(mapped_schema_idx) +
                              " not found in row group",
                            std::invalid_argument);
-              colchunk_iter_offset = std::distance(row_group.columns.begin(), colchunk_iter);
+              colchunk_offset =
+                static_cast<size_type>(std::distance(row_group.columns.begin(), it));
             }
-            auto const colchunk_iter = row_group.columns.begin() + colchunk_iter_offset.value();
-            auto const& col_chunk    = *colchunk_iter;
-            auto const& col_meta     = col_chunk.meta_data;
+
+            auto const& col_chunk = row_group.columns[colchunk_offset.value()];
+            auto const& col_meta  = col_chunk.meta_data;
 
             // Make sure that we have page index and the column chunk doesn't have any
             // non-dictionary encoded pages

cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp

@@ -297,6 +297,19 @@ hybrid_scan_reader_impl::secondary_filters_byte_ranges(
   return {bloom_filter_bytes, dictionary_page_bytes};
 }
 
+std::vector<byte_range_info> hybrid_scan_reader_impl::dictionary_pages_byte_ranges(
+  cudf::host_span<std::vector<size_type> const> row_group_indices,
+  parquet_reader_options const& options)
+{
+  CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered");
+  auto [expr_conv, output_dtypes] = prepare_filter_and_output_types(options);
+
+  return _extended_metadata->get_dictionary_page_bytes(row_group_indices,
+                                                       output_dtypes,
+                                                       _output_column_schemas,
+                                                       expr_conv.get_converted_expr().value());
+}
+
 std::vector<std::vector<size_type>>
 hybrid_scan_reader_impl::filter_row_groups_with_dictionary_pages(
   cudf::host_span<cudf::device_span<uint8_t const> const> dictionary_page_data,

cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp

@@ -112,6 +112,13 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
   secondary_filters_byte_ranges(cudf::host_span<std::vector<size_type> const> row_group_indices,
                                 parquet_reader_options const& options);
 
+  /**
+   * @copydoc cudf::io::experimental::hybrid_scan_multifile::dictionary_pages_byte_ranges
+   */
+  [[nodiscard]] std::vector<byte_range_info> dictionary_pages_byte_ranges(
+    cudf::host_span<std::vector<size_type> const> row_group_indices,
+    parquet_reader_options const& options);
+
   /**
    * @copydoc cudf::io::experimental::hybrid_scan::filter_row_groups_with_dictionary_pages
    */

cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp

@@ -83,4 +83,23 @@ hybrid_scan_multifile::secondary_filters_byte_ranges(
   return _impl->secondary_filters_byte_ranges(row_group_indices, options);
 }
 
+std::vector<text::byte_range_info> hybrid_scan_multifile::dictionary_pages_byte_ranges(
+  cudf::host_span<std::vector<size_type> const> row_group_indices,
+  parquet_reader_options const& options) const
+{
+  CUDF_FUNC_RANGE();
+  return _impl->dictionary_pages_byte_ranges(row_group_indices, options);
+}
+
+std::vector<std::vector<size_type>> hybrid_scan_multifile::filter_row_groups_with_dictionary_pages(
+  cudf::host_span<cudf::device_span<uint8_t const> const> dictionary_page_data,
+  cudf::host_span<std::vector<size_type> const> row_group_indices,
+  parquet_reader_options const& options,
+  rmm::cuda_stream_view stream) const
+{
+  CUDF_FUNC_RANGE();
+  return _impl->filter_row_groups_with_dictionary_pages(
+    dictionary_page_data, row_group_indices, options, stream);
+}
+
 }  // namespace cudf::io::parquet::experimental

cpp/tests/io/experimental/hybrid_scan_common.hpp

@@ -9,6 +9,7 @@
 
 #include <cudf_test/column_wrapper.hpp>
 
+#include <cudf/column/column_view.hpp>
 #include <cudf/io/parquet.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/types.hpp>
@@ -18,6 +19,8 @@
 #include <cuda/iterator>
 
 #include <format>
+#include <string>
+#include <vector>
 
 /**
  * @brief Creates a strings column with a constant stringified value between 0 and 9999
@@ -92,6 +95,11 @@ cudf::test::fixed_width_column_wrapper<T> descending_low_cardinality()
  *
  * @param str_col_value Value for the constant string column used when IsConstantStrings is true
  * @param compression Compression type
+ * @param column_names Top-level column names assigned in `column_order` order (default
+ *        {"col0", "col1", "col2"})
+ * @param column_order Physical emit order of the base [col0, col1, col2] columns (default
+ *        {0, 1, 2}). Reordering emits the same logical columns at different schema positions, which
+ *        is used to build mismatched per-source schemas for the row-group filtering tests.
  * @param stream CUDA stream
  *
  * @return Tuple of table and Parquet host buffer
@@ -101,9 +109,11 @@ template <typename T,
           bool IsConstantStrings = true,
           bool IsNullable        = false>
 auto create_parquet_with_stats(
-  cudf::size_type str_col_value          = 100,
-  cudf::io::compression_type compression = cudf::io::compression_type::AUTO,
-  rmm::cuda_stream_view stream           = cudf::get_default_stream())
+  cudf::size_type str_col_value             = 100,
+  cudf::io::compression_type compression    = cudf::io::compression_type::AUTO,
+  std::vector<std::string> column_names     = {"col0", "col1", "col2"},
+  std::vector<cudf::size_type> column_order = {0, 1, 2},
+  rmm::cuda_stream_view stream              = cudf::get_default_stream())
 {
   static_assert(NumTableConcats >= 1, "Concatenated table must contain at least one table");
 
@@ -157,12 +167,23 @@ auto create_parquet_with_stats(
     output = table_view{{columns[0]->view(), columns[1]->view(), columns[2]->view()}};
   }
 
+  // Reorder the base [col0, col1, col2] columns into the requested physical order, naming them in
+  // that new order. Reordering lets callers emit the same logical columns at different schema
+  // positions across sources (a "mismatched schema"), which exercises the per-source schema-index
+  // mapping in the row-group filtering paths. The defaults leave the table and names unchanged.
+  std::vector<cudf::column_view> reordered_columns;
+  reordered_columns.reserve(column_order.size());
+  for (auto const col_idx : column_order) {
+    reordered_columns.emplace_back(output.column(col_idx));
+  }
+  output = table_view{reordered_columns};
+
   auto table = cudf::concatenate(std::vector<table_view>(NumTableConcats, output));
   output     = table->view();
   cudf::io::table_input_metadata output_metadata(output);
-  output_metadata.column_metadata[0].set_name("col0");
-  output_metadata.column_metadata[1].set_name("col1");
-  output_metadata.column_metadata[2].set_name("col2");
+  for (std::size_t i = 0; i < column_names.size(); ++i) {
+    output_metadata.column_metadata[i].set_name(column_names[i]);
+  }
 
   std::vector<char> buffer;
   cudf::io::parquet_writer_options out_opts =