Add multifile dictionary pruning support for hybrid scan parquet reader

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

@@ -454,6 +454,35 @@ class hybrid_scan_multifile {
    */
   [[nodiscard]] bool has_next_table_chunk() const;
 
+  /**
+   * @brief Get byte ranges of column chunk dictionary pages for row group pruning
+   *
+   * @param row_group_indices Span of vectors of input row group indices, one per source
+   * @param options Parquet reader options
+   * @return Pair of flattened byte ranges to column chunk dictionary pages subject to the filter
+   *         predicate and their corresponding source indices
+   */
+  [[nodiscard]] std::pair<std::vector<byte_range_info>, std::vector<size_type>>
+  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 Span of vectors of 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 Vector of vectors of filtered row group indices, one 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

@@ -31,6 +31,7 @@ using metadata_base                  = parquet::detail::metadata;
 using io::detail::inline_column_buffer;
 using parquet::detail::CompactProtocolReader;
 using parquet::detail::equality_literals_collector;
+using parquet::detail::find_colchunk_iter_offset;
 using parquet::detail::input_column_info;
 using parquet::detail::row_group_info;
 using text::byte_range_info;
@@ -429,7 +430,8 @@ std::vector<byte_range_info> aggregate_reader_metadata::get_bloom_filter_bytes(
   return bloom_filter_bytes;
 }
 
-std::vector<byte_range_info> aggregate_reader_metadata::get_dictionary_page_bytes(
+std::pair<std::vector<byte_range_info>, std::vector<cudf::size_type>>
+aggregate_reader_metadata::dictionary_pages_byte_ranges(
   std::span<std::vector<cudf::size_type> const> row_group_indices,
   std::span<data_type const> output_dtypes,
   std::span<cudf::size_type const> output_column_schemas,
@@ -447,7 +449,7 @@ std::vector<byte_range_info> aggregate_reader_metadata::get_dictionary_page_byte
                   std::back_inserter(dictionary_col_schemas),
                   [](auto& dict_literals) { return not dict_literals.empty(); });
 
-  // No (in)equality literals found, return empty vector
+  // No (in)equality literals found, return empty vectors
   if (dictionary_col_schemas.empty()) { return {}; }
 
   // Compute total number of input row groups
@@ -463,38 +465,41 @@ 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;
 
+  // Association between each dictionary page byte range and its source
+  std::vector<cudf::size_type> dictionary_page_source_map;
+  dictionary_page_source_map.reserve(num_chunks);
+
+  // 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
         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) +
-                             " not found in row group",
-                           std::invalid_argument);
-              colchunk_iter_offset = std::distance(row_group.columns.begin(), colchunk_iter);
+          cuda::counting_iterator<std::size_t>{0},
+          cuda::counting_iterator{dictionary_col_schemas.size()},
+          [&](auto const col) {
+            // Map the schema index to this source
+            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) {
+              colchunk_offset = find_colchunk_iter_offset(row_group, mapped_schema_idx);
             }
-            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
@@ -550,13 +555,14 @@ std::vector<byte_range_info> aggregate_reader_metadata::get_dictionary_page_byte
             }
 
             dictionary_page_bytes.emplace_back(dictionary_offset, dictionary_size);
+            dictionary_page_source_map.emplace_back(static_cast<size_type>(src_index));
           });
       });
     });
 
   if (not have_dictionary_pages) { return {}; }
 
-  return dictionary_page_bytes;
+  return {std::move(dictionary_page_bytes), std::move(dictionary_page_source_map)};
 }
 
 std::vector<std::vector<cudf::size_type>>

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

@@ -229,13 +229,15 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base {
    * @param output_column_schemas schema indices of output columns
    * @param filter AST expression to filter row groups based on dictionary pages
    *
-   * @return Byte ranges of dictionary pages, one input column chunk with (in)equality predicate
+   * @return A pair of vectors containing dictionary page byte ranges and corresponding source
+   *         indices
    */
-  [[nodiscard]] std::vector<cudf::io::text::byte_range_info> get_dictionary_page_bytes(
-    std::span<std::vector<size_type> const> row_group_indices,
-    std::span<data_type const> output_dtypes,
-    std::span<cudf::size_type const> output_column_schemas,
-    std::reference_wrapper<ast::expression const> filter);
+  [[nodiscard]] std::pair<std::vector<cudf::io::text::byte_range_info>,
+                          std::vector<cudf::size_type>>
+  dictionary_pages_byte_ranges(std::span<std::vector<size_type> const> row_group_indices,
+                               std::span<data_type const> output_dtypes,
+                               std::span<cudf::size_type const> output_column_schemas,
+                               std::reference_wrapper<ast::expression const> filter);
 
   /**
    * @brief Filter the row groups using dictionaries based on predicate filter

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

@@ -292,14 +292,30 @@ hybrid_scan_reader_impl::secondary_filters_byte_ranges(
                                                _output_column_schemas,
                                                expr_conv.get_converted_expr().value());
   auto const dictionary_page_bytes =
-    _extended_metadata->get_dictionary_page_bytes(row_group_indices,
-                                                  output_dtypes,
-                                                  _output_column_schemas,
-                                                  expr_conv.get_converted_expr().value());
+    _extended_metadata
+      ->dictionary_pages_byte_ranges(row_group_indices,
+                                     output_dtypes,
+                                     _output_column_schemas,
+                                     expr_conv.get_converted_expr().value())
+      .first;
 
   return {bloom_filter_bytes, dictionary_page_bytes};
 }
 
+std::pair<std::vector<byte_range_info>, std::vector<cudf::size_type>>
+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->dictionary_pages_byte_ranges(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(
   std::span<cudf::device_span<uint8_t const> const> dictionary_page_data,

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

@@ -115,6 +115,13 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
   secondary_filters_byte_ranges(std::span<std::vector<size_type> const> row_group_indices,
                                 parquet_reader_options const& options);
 
+  /**
+   * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::dictionary_pages_byte_ranges
+   */
+  [[nodiscard]] std::pair<std::vector<byte_range_info>, std::vector<cudf::size_type>>
+  dictionary_pages_byte_ranges(cudf::host_span<std::vector<size_type> const> row_group_indices,
+                               parquet_reader_options const& options);
+
   /**
    * @copydoc cudf::io::parquet::experimental::hybrid_scan::filter_row_groups_with_dictionary_pages
    */
@@ -283,7 +290,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
    * @throws std::invalid_argument if @p row_group_indices.size() is all empty or not equal to the
    * number of input datasources
    *
-   * @param row_group_indices Input row group indices, one per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param total_row_groups Total number of row groups across all sources
    * @param pass_read_limit Memory limit to read and decompress row
    * group data

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

@@ -295,4 +295,24 @@ std::vector<std::vector<std::vector<size_type>>> hybrid_scan_multifile::construc
   return source_passes;
 }
 
+std::pair<std::vector<text::byte_range_info>, std::vector<size_type>>
+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/src/io/parquet/experimental/page_index_filter.cu

@@ -855,6 +855,12 @@ std::unique_ptr<cudf::column> aggregate_reader_metadata::build_row_mask_with_pag
   // Return if empty row group indices
   if (row_group_indices.empty()) { return cudf::make_empty_column(cudf::type_id::BOOL8); }
 
+  // TODO(#22900): remove this guard once this path maps schema indices per source. It currently
+  // reuses one source's schema index for every source, so it is correct only when schemas match.
+  CUDF_EXPECTS(schema_idx_maps.empty(),
+               "Page index statistics filtering does not support mismatched Parquet schemas yet",
+               std::invalid_argument);
+
   // Check if we have page index for all columns in all row groups
   auto const has_page_index = compute_has_page_index(per_file_metadata, row_group_indices);
 
@@ -1007,6 +1013,12 @@ thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(
     return thrust::host_vector<bool>(0, stream);
   }
 
+  // TODO(#22900): remove this guard once this path maps schema indices per source. It currently
+  // reuses one source's schema index for every source, so it is correct only when schemas match.
+  CUDF_EXPECTS(schema_idx_maps.empty(),
+               "Data page masking does not support mismatched Parquet schemas yet",
+               std::invalid_argument);
+
   // Collect column schema indices from the input columns.
   auto column_schema_indices = std::vector<size_type>(input_columns.size());
   std::transform(

cpp/src/io/parquet/experimental/page_index_filter_utils.cu

@@ -22,18 +22,7 @@
 
 namespace cudf::io::parquet::experimental::detail {
 
-size_type find_colchunk_iter_offset(RowGroup const& row_group, size_type 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) + " not found in row group",
-    std::invalid_argument);
-  return std::distance(row_group.columns.begin(), colchunk_iter);
-}
+using parquet::detail::find_colchunk_iter_offset;
 
 bool compute_has_page_index(std::span<metadata_base const> file_metadatas,
                             std::span<std::vector<size_type> const> row_group_indices)

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

@@ -22,15 +22,6 @@ namespace cudf::io::parquet::experimental::detail {
 
 using metadata_base = parquet::detail::metadata;
 
-/**
- * @brief Find the offset of the column chunk with the given schema index in the row group
- *
- * @param row_group Row group
- * @param schema_idx Schema index
- * @return Offset of the column chunk iterator
- */
-[[nodiscard]] size_type find_colchunk_iter_offset(RowGroup const& row_group, size_type schema_idx);
-
 /**
  * @brief Compute if the page index is present in all parquet data sources for all columns
  *

cpp/src/io/parquet/predicate_pushdown.cpp

@@ -23,7 +23,6 @@
 #include <thrust/iterator/counting_iterator.h>
 
 #include <algorithm>
-#include <format>
 #include <functional>
 #include <numeric>
 #include <optional>
@@ -155,17 +154,7 @@ bool aggregate_reader_metadata::any_row_group_stats_available(
 
       if (cached_offset < 0 or cached_offset >= num_col_chunks or
           first_row_group.columns[cached_offset].schema_idx != mapped_schema_idx) {
-        auto const it = std::find_if(
-          first_row_group.columns.begin(),
-          first_row_group.columns.end(),
-          [mapped_schema_idx](ColumnChunk const& c) { return c.schema_idx == mapped_schema_idx; });
-        CUDF_EXPECTS(
-          it != first_row_group.columns.end(),
-          std::format(
-            "Column chunk with schema index {} not found in source {}", mapped_schema_idx, src_idx),
-          std::invalid_argument);
-        colchunk_offset =
-          static_cast<size_type>(std::distance(first_row_group.columns.begin(), it));
+        colchunk_offset = find_colchunk_iter_offset(first_row_group, mapped_schema_idx);
       }
 
       if (colchunk_has_stats(first_row_group.columns[colchunk_offset.value()])) { return true; }

cpp/src/io/parquet/reader_impl_helpers.cpp

@@ -57,6 +57,18 @@ std::size_t derive_pass_read_limit(std::size_t chunk_read_limit)
   return pass_read_limit;
 }
 
+size_type find_colchunk_iter_offset(RowGroup const& row_group, size_type 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(),
+               std::format("Column chunk with schema index {} not found in row group", schema_idx),
+               std::invalid_argument);
+  return std::distance(row_group.columns.begin(), colchunk_iter);
+}
+
 namespace flatbuf = cudf::io::parquet::flatbuf;
 
 namespace {
@@ -1215,14 +1227,8 @@ ColumnChunkMetaData const& aggregate_reader_metadata::get_column_metadata(size_t
   // Map schema index to the provided source file index
   schema_idx = map_schema_index(schema_idx, src_idx);
 
-  auto col =
-    std::find_if(per_file_metadata[src_idx].row_groups[row_group_index].columns.begin(),
-                 per_file_metadata[src_idx].row_groups[row_group_index].columns.end(),
-                 [schema_idx](ColumnChunk const& col) { return col.schema_idx == schema_idx; });
-  CUDF_EXPECTS(col != std::end(per_file_metadata[src_idx].row_groups[row_group_index].columns),
-               "Found no metadata for schema index",
-               std::range_error);
-  return col->meta_data;
+  auto const& row_group = per_file_metadata[src_idx].row_groups[row_group_index];
+  return row_group.columns[find_colchunk_iter_offset(row_group, schema_idx)].meta_data;
 }
 
 std::vector<std::unordered_map<std::string, int64_t>>

cpp/src/io/parquet/reader_impl_helpers.hpp

@@ -108,6 +108,18 @@ struct row_group_info {
  */
 [[nodiscard]] std::size_t derive_pass_read_limit(std::size_t chunk_read_limit);
 
+/**
+ * @brief Find the offset of the column chunk with the given schema index in the specified row group
+ *
+ * @note For mismatched schemas, `schema_idx` must be pre-mapped to the row group's source using
+ * `map_schema_index`.
+ *
+ * @param row_group Row group
+ * @param schema_idx Schema index, already mapped to the row group's source
+ * @return Offset of the column chunk within the row group's columns
+ */
+[[nodiscard]] size_type find_colchunk_iter_offset(RowGroup const& row_group, size_type schema_idx);
+
 /**
  * @brief Class for parsing dataset metadata
  */