Dictionary pruning on files without an offset index or encoding stats

cpp/benchmarks/io/parquet/experimental/hybrid_scan/dict_page_filter.cpp

@@ -85,8 +85,8 @@ void BM_filter_string_row_groups_with_dicts_common(nvbench::state& state,
       timer.start();
 
       // Get dictionary page byte ranges
-      dict_page_byte_ranges =
-        reader->dictionary_pages_byte_ranges(input_row_group_indices, read_opts);
+      dict_page_byte_ranges = cudf::io::parquet::experimental::dictionary_page_byte_ranges_to_read(
+        reader->dictionary_pages_byte_ranges(input_row_group_indices, read_opts));
       CUDF_EXPECTS(not dict_page_byte_ranges.empty(), "No dictionary page byte ranges found");
 
       // Fetch dictionary page data

cpp/benchmarks/io/parquet/experimental/hybrid_scan/hybrid_scan_composer.cpp

@@ -73,7 +73,8 @@ std::vector<cudf::size_type> apply_row_group_filters(
 
   if (filters.contains(hybrid_scan_filter_type::ROW_GROUPS_WITH_DICT_PAGES)) {
     auto const dict_page_byte_ranges =
-      reader.dictionary_pages_byte_ranges(current_row_group_indices, options);
+      cudf::io::parquet::experimental::dictionary_page_byte_ranges_to_read(
+        reader.dictionary_pages_byte_ranges(current_row_group_indices, options));
 
     if (not dict_page_byte_ranges.empty()) {
       auto [dictionary_page_buffers, dictionary_page_data, dict_read_tasks] =

cpp/examples/hybrid_scan_io/hybrid_scan_composer.cpp

@@ -129,12 +129,12 @@ std::vector<cudf::size_type> apply_row_group_filters(
     }
   }
 
-  // Get dictionary page byte ranges from the reader
-  std::vector<cudf::io::text::byte_range_info> dict_page_byte_ranges;
+  // Get dictionary page ranges from the reader
+  std::vector<cudf::io::parquet::experimental::dictionary_page_range> dict_page_ranges;
   if (filters.contains(hybrid_scan_filter_type::ROW_GROUPS_WITH_DICT_PAGES)) {
     if (verbose) { std::cout << "READER: Get dictionary page byte ranges...\n"; }
     timer.reset();
-    dict_page_byte_ranges = reader.dictionary_pages_byte_ranges(current_row_group_indices, options);
+    dict_page_ranges = reader.dictionary_pages_byte_ranges(current_row_group_indices, options);
     if (verbose) { timer.print_elapsed_millis(); }
   }
 
@@ -143,15 +143,16 @@ std::vector<cudf::size_type> apply_row_group_filters(
   // Filter row groups with dictionary pages
   std::vector<cudf::size_type> dictionary_page_filtered_row_group_indices;
   dictionary_page_filtered_row_group_indices.reserve(current_row_group_indices.size());
-  if (not dict_page_byte_ranges.empty()) {
+  if (not dict_page_ranges.empty()) {
     if (verbose) { std::cout << "READER: Filter row groups with dictionary pages...\n"; }
     timer.reset();
 
-    // Fetch dictionary page buffers and corresponding device spans from the input file buffer
+    // Fetch dictionary pages, trimming any upper-bound range to exactly one dictionary page (or an
+    // empty span when the chunk has none) so the reader never reads data-page bytes as dictionary
+    // data.
     nvtxRangePush("fetch_dict_page_byte_ranges");
-    auto [dictionary_page_buffers, dictionary_page_data, dict_read_tasks] =
-      fetch_byte_ranges_async(datasource, dict_page_byte_ranges, stream, temp_mr);
-    dict_read_tasks.get();
+    auto [dictionary_page_buffers, dictionary_page_data] =
+      fetch_dictionary_pages(datasource, dict_page_ranges, stream, temp_mr);
     nvtxRangePop();
 
     dictionary_page_filtered_row_group_indices = reader.filter_row_groups_with_dictionary_pages(

cpp/examples/hybrid_scan_io/io_utils.cpp

@@ -3,14 +3,24 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include "io_utils.hpp"
+
 #include <cudf/io/datasource.hpp>
+#include <cudf/io/experimental/hybrid_scan.hpp>
 #include <cudf/io/parquet_io_utils.hpp>
 #include <cudf/io/text/byte_range_info.hpp>
+#include <cudf/utilities/span.hpp>
 
+#include <rmm/device_buffer.hpp>
 #include <rmm/resource_ref.hpp>
 
 #include <cuda/stream>
 
+#include <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
 /**
  * @file io_utils.cpp
  * @brief Definitions for IO utilities for hybrid_scan examples
@@ -40,3 +50,61 @@ fetch_byte_ranges_async(cudf::io::datasource& datasource,
   // Using libcudf utility but may have custom implementation in the future
   return cudf::io::parquet::fetch_byte_ranges_to_device_async(datasource, byte_ranges, stream, mr);
 }
+
+std::pair<std::vector<rmm::device_buffer>, std::vector<cudf::device_span<uint8_t const>>>
+fetch_dictionary_pages(cudf::io::datasource& datasource,
+                       cudf::host_span<cudf::io::parquet::experimental::dictionary_page_range const>
+                         dictionary_page_ranges,
+                       cuda::stream_ref stream,
+                       rmm::device_async_resource_ref mr,
+                       int64_t max_upper_bound_size)
+{
+  using cudf::io::parquet::experimental::dictionary_page_extent;
+
+  auto const read_ranges = cudf::io::parquet::experimental::dictionary_page_byte_ranges_to_read(
+    dictionary_page_ranges, max_upper_bound_size);
+
+  auto buffers = std::vector<rmm::device_buffer>{};
+  auto spans   = std::vector<cudf::device_span<uint8_t const>>{};
+  buffers.reserve(read_ranges.size());
+  spans.reserve(read_ranges.size());
+
+  // Keep the host reads alive until the async copies below have completed.
+  auto host_reads = std::vector<std::unique_ptr<cudf::io::datasource::buffer>>{};
+
+  for (std::size_t i = 0; i < read_ranges.size(); ++i) {
+    auto const read_size = read_ranges[i].size();
+
+    // A chunk the reader will not prune with has an empty range; keep an empty span in its place.
+    if (read_size == 0) {
+      buffers.emplace_back();
+      spans.emplace_back();
+      continue;
+    }
+
+    auto host_buffer = datasource.host_read(read_ranges[i].offset(), read_size);
+    auto const bytes = cudf::host_span<uint8_t const>{host_buffer->data(), host_buffer->size()};
+
+    // An exact range is already one page; an upper-bound range has to be measured and trimmed, and
+    // may hold no dictionary page at all.
+    auto const page_size =
+      (dictionary_page_ranges[i].extent == dictionary_page_extent::upper_bound_if_present)
+        ? cudf::io::parquet::experimental::dictionary_page_length(bytes).value_or(0)
+        : static_cast<int64_t>(bytes.size());
+
+    if (page_size == 0) {
+      buffers.emplace_back();
+      spans.emplace_back();
+      continue;
+    }
+
+    auto const page_bytes = static_cast<std::size_t>(page_size);
+    auto device_buffer    = rmm::device_buffer{bytes.data(), page_bytes, stream, mr};
+    spans.emplace_back(static_cast<uint8_t const*>(device_buffer.data()), page_bytes);
+    buffers.emplace_back(std::move(device_buffer));
+    host_reads.emplace_back(std::move(host_buffer));
+  }
+
+  stream.sync();  // host_reads are freed on return, so the copies must finish first
+  return {std::move(buffers), std::move(spans)};
+}

cpp/examples/hybrid_scan_io/io_utils.hpp

@@ -6,15 +6,20 @@
 #pragma once
 
 #include <cudf/io/datasource.hpp>
+#include <cudf/io/experimental/hybrid_scan.hpp>
+#include <cudf/io/parquet.hpp>
 #include <cudf/io/text/byte_range_info.hpp>
 
 #include <rmm/device_buffer.hpp>
 #include <rmm/resource_ref.hpp>
 
 #include <cuda/stream>
 
+#include <cstdint>
 #include <future>
+#include <limits>
 #include <tuple>
+#include <utility>
 #include <vector>
 
 /**
@@ -58,3 +63,29 @@ fetch_byte_ranges_async(cudf::io::datasource& datasource,
                         cudf::host_span<cudf::io::text::byte_range_info const> byte_ranges,
                         cuda::stream_ref stream,
                         rmm::device_async_resource_ref mr);
+
+/**
+ * @brief Fetches dictionary pages, trimming every upper-bound range to exactly one dictionary page
+ *
+ * Reads each range on the host (capping a range that only bounds its page at `max_upper_bound_size`
+ * bytes), measures a real dictionary page with `dictionary_page_length`, and copies only the
+ * verified page bytes to the device, leaving an empty span for a chunk with no dictionary page.
+ * Positions are preserved so the returned spans stay aligned with `dictionary_page_ranges`, which
+ * is what `filter_row_groups_with_dictionary_pages` expects.
+ *
+ * @param datasource Input datasource
+ * @param dictionary_page_ranges Dictionary page ranges from `dictionary_pages_byte_ranges`
+ * @param stream CUDA stream
+ * @param mr Device memory resource
+ * @param max_upper_bound_size Most bytes to read of a range that only bounds its dictionary page
+ *
+ * @return Owning device buffers and one device span per input range
+ */
+std::pair<std::vector<rmm::device_buffer>, std::vector<cudf::device_span<uint8_t const>>>
+fetch_dictionary_pages(cudf::io::datasource& datasource,
+                       cudf::host_span<cudf::io::parquet::experimental::dictionary_page_range const>
+                         dictionary_page_ranges,
+                       cuda::stream_ref stream,
+                       rmm::device_async_resource_ref mr,
+                       int64_t max_upper_bound_size =
+                         cudf::io::parquet::experimental::default_max_dictionary_page_read_size);

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

@@ -17,6 +17,7 @@
 #include <cuda/stream>
 
 #include <memory>
+#include <optional>
 #include <span>
 #include <utility>
 #include <vector>
@@ -58,6 +59,73 @@ enum class use_data_page_mask : bool {
   NO  = false  ///< Do not compute or use a data page mask
 };
 
+/**
+ * @brief How closely a dictionary page byte range describes the page it points at
+ *
+ * An `upper_bound_if_present` range begins at the dictionary page if the column chunk has one, and
+ * ends no earlier than that page does. A writer is allowed to leave out where the page ends, and to
+ * say that a chunk is dictionary encoded when it holds no dictionary page at all, so a range of
+ * this kind is a bound on a page that may not be there.
+ */
+enum class dictionary_page_extent : bool {
+  exact,                  ///< The range is exactly the dictionary page
+  upper_bound_if_present  ///< The range bounds a dictionary page that may not be there
+};
+
+/**
+ * @brief Byte range of a column chunk's dictionary page, and how closely it describes that page
+ *
+ * A caller is free to read less than an `upper_bound_if_present` range, which is how it caps what
+ * it spends looking for a page that may not be there. The reader still wants a span holding exactly
+ * one dictionary page, so a caller that reads such a range measures the page in it with
+ * `dictionary_page_length`, and passes an empty span for a chunk whose page is not there or does
+ * not fit in what was read.
+ */
+struct dictionary_page_range {
+  byte_range_info byte_range;     ///< Byte range to read from the file
+  dictionary_page_extent extent;  ///< How closely `byte_range` describes the dictionary page
+};
+
+/**
+ * @brief Default cap on the bytes read of a range that only bounds its dictionary page
+ *
+ * One mebibyte is what writers commonly cap a dictionary at, and the slack on top of that
+ * covers the page header and compression framing. A column chunk whose dictionary page does
+ * not fit is not pruned.
+ */
+constexpr int64_t default_max_dictionary_page_read_size = (1024 * 1024) + (64 * 1024);
+
+/**
+ * @brief Byte ranges to read for the specified dictionary page ranges
+ *
+ * No more than `max_upper_bound_size` bytes are read of a range that only bounds its dictionary
+ * page, which is how a caller caps what it spends looking for a page that may not be there. What is
+ * read of such a range still has to be trimmed to the dictionary page before it is handed to the
+ * reader, see `dictionary_page_range`.
+ *
+ * @param dictionary_page_ranges Dictionary page ranges from `dictionary_pages_byte_ranges`
+ * @param max_upper_bound_size Most bytes to read of a range that only bounds its dictionary page. A
+ *        column chunk whose dictionary page is longer than this is not pruned.
+ * @return Byte ranges to read, one per input dictionary page range
+ */
+[[nodiscard]] std::vector<byte_range_info> dictionary_page_byte_ranges_to_read(
+  cudf::host_span<dictionary_page_range const> dictionary_page_ranges,
+  int64_t max_upper_bound_size = default_max_dictionary_page_read_size);
+
+/**
+ * @brief Length of the dictionary page at the front of the specified bytes, header included
+ *
+ * What was read of a range that only bounds its dictionary page begins at that page and runs past
+ * it. The page's own header says how long the page is, so this reads that header to find where the
+ * page ends, which is what turns such a range into the one page the reader takes.
+ *
+ * @param page_bytes Bytes read for a dictionary page range, from the start of the range
+ * @return Length of the dictionary page, or `std::nullopt` if these bytes do not begin with a whole
+ *         dictionary page, which is the case for a column chunk that has none to prune with
+ */
+[[nodiscard]] std::optional<int64_t> dictionary_page_length(
+  cudf::host_span<uint8_t const> page_bytes);
+
 /**
  * @brief Shareable, pre-parsed Parquet file metadata for the Hybrid Scan reader.
  *
@@ -217,18 +285,42 @@ class hybrid_scan_metadata {
  * // Update current row group indices to now track the stats-filtered row group indices
  * current_row_group_indices = stats_filtered_row_group_indices;
  *
- * // Get byte ranges of dictionary pages for the current row groups
- * auto dict_page_byte_ranges =
+ * // Get the dictionary page ranges for the current row groups
+ * auto dict_page_ranges =
  *   reader->dictionary_pages_byte_ranges(current_row_group_indices, options);
  *
  * // Optional: Prune row groups if we have valid dictionary pages
  * auto dict_filtered_row_group_indices = std::vector<size_type>{};
  *
- * if (dict_page_byte_ranges.size()) {
- *   // Fetch dictionary page byte ranges into device buffers and create spans
- *   auto [dict_page_buffers, dict_page_data, dict_page_tasks] =
- *     parquet::fetch_byte_ranges_to_device_async(datasource, dict_page_byte_ranges, stream, mr);
- *   dict_page_tasks.get();
+ * if (dict_page_ranges.size()) {
+ *   // Decide how much of each range to read. A range that only bounds its dictionary page can be
+ *   // much larger than the page it bounds, so read no more of it than a dictionary page is worth.
+ *   auto const dict_page_byte_ranges = dictionary_page_byte_ranges_to_read(dict_page_ranges);
+ *
+ *   // Hand the reader exactly one dictionary page per chunk. An `exact` range is already one page,
+ *   // but an `upper_bound_if_present` range runs past its page and may hold none at all, so it is
+ *   // read on the host, measured with `dictionary_page_length`, and copied to the device trimmed
+ *   // to that page, or left as an empty span when the chunk has no dictionary page. Passing the
+ *   // untrimmed range would let the reader read the following data-page bytes as dictionary data.
+ *   // The reader matches spans to ranges by position, so an empty span is kept in place.
+ *   auto dict_page_buffers = std::vector<rmm::device_buffer>{};
+ *   auto dict_page_data    = std::vector<device_span<uint8_t const>>{};
+ *   auto host_reads        = std::vector<std::unique_ptr<datasource::buffer>>{};
+ *   for (auto i = 0uz; i < dict_page_byte_ranges.size(); ++i) {
+ *     auto const& read_range = dict_page_byte_ranges[i];
+ *     auto host_bytes        = datasource.host_read(read_range.offset(), read_range.size());
+ *     auto const bytes = host_span<uint8_t const>{host_bytes->data(), host_bytes->size()};
+ *     auto const page_size =
+ *       (dict_page_ranges[i].extent == dictionary_page_extent::upper_bound_if_present)
+ *         ? dictionary_page_length(bytes).value_or(0)
+ *         : static_cast<int64_t>(bytes.size());
+ *     // Copy the first `page_size` bytes to the device (an empty buffer when there is no page)
+ *     dict_page_buffers.emplace_back(bytes.data(), page_size, stream, mr);
+ *     dict_page_data.emplace_back(
+ *       static_cast<uint8_t const*>(dict_page_buffers.back().data()), page_size);
+ *     host_reads.emplace_back(std::move(host_bytes));  // keep alive until the copies complete
+ *   }
+ *   stream.synchronize();
  *
  *   // Prune row groups using dictionaries
  *   dict_filtered_row_group_indices = reader->filter_row_groups_with_dictionary_pages(
@@ -506,18 +598,22 @@ class hybrid_scan_reader {
     cuda::stream_ref stream) const;
 
   /**
-   * @brief Get byte ranges of column chunk dictionary pages for row group pruning
+   * @brief Get the ranges of column chunk dictionary pages for row group pruning
    *
    * @param row_group_indices Input row groups indices
    * @param options Parquet reader options
-   * @return Vector of byte ranges to column chunk dictionary pages subject to the filter predicate
+   * @return Vector of dictionary page ranges of column chunks subject to the filter predicate
    */
-  [[nodiscard]] std::vector<byte_range_info> dictionary_pages_byte_ranges(
+  [[nodiscard]] std::vector<dictionary_page_range> dictionary_pages_byte_ranges(
     std::span<size_type const> row_group_indices, parquet_reader_options const& options) const;
 
   /**
    * @brief Filter the row groups using column chunk dictionary pages
    *
+   * Each span must hold exactly one dictionary page, or nothing at all for a column chunk that has
+   * no dictionary page to prune with. See `dictionary_page_range` for trimming a range that only
+   * bounds its page.
+   *
    * @param dictionary_page_data Device spans of dictionary page data of column chunks with an
    * (in)equality predicate, in the same order as the byte ranges returned by
    * `dictionary_pages_byte_ranges` including empty spans against empty byte ranges

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

@@ -194,16 +194,20 @@ class hybrid_scan_multifile {
    *
    * @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
+   * @return Pair of flattened dictionary page ranges of column chunks subject to the filter
    *         predicate and their corresponding source indices
    */
-  [[nodiscard]] std::pair<std::vector<byte_range_info>, std::vector<size_type>>
+  [[nodiscard]] std::pair<std::vector<dictionary_page_range>, 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
    *
+   * Each span must hold exactly one dictionary page, or nothing at all for a column chunk that has
+   * no dictionary page to prune with. See `dictionary_page_range` for trimming a range that only
+   * bounds its page.
+   *
    * @param dictionary_page_data Device spans of dictionary page data of column chunks with an
    * (in)equality predicate, in the same order as the byte ranges returned by
    * `dictionary_pages_byte_ranges` including empty spans against empty byte ranges