Fix Parquet page-index read ranges for optional indexes

cpp/benchmarks/io/parquet/parquet_reader_metadata.cpp

@@ -11,19 +11,26 @@
 #include <cudf/ast/expressions.hpp>
 #include <cudf/detail/utilities/integer_utils.hpp>
 #include <cudf/io/datasource.hpp>
+#include <cudf/io/experimental/hybrid_scan.hpp>
 #include <cudf/io/parquet.hpp>
+#include <cudf/io/parquet_io_utils.hpp>
 #include <cudf/io/parquet_metadata.hpp>
 #include <cudf/scalar/scalar.hpp>
 #include <cudf/utilities/default_stream.hpp>
 
 #include <cuda/iterator>
 
 #include <nvbench/nvbench.cuh>
+#include <src/io/parquet/compact_protocol_reader.hpp>
+#include <src/io/parquet/compact_protocol_writer.hpp>
 
 #include <algorithm>
+#include <atomic>
 #include <cctype>
 #include <cstddef>
+#include <cstring>
 #include <limits>
+#include <memory>
 #include <string>
 #include <utility>
 #include <vector>
@@ -84,6 +91,87 @@ auto write_file_data(cudf::size_type num_cols,
   return source_sink;
 }
 
+/** @brief Counts logical host-read bytes without including filesystem or page-cache effects. */
+class PageIndexCountingDatasource : public cudf::io::datasource {
+ public:
+  explicit PageIndexCountingDatasource(std::vector<char> const& data)
+    : source_{cudf::io::datasource::create(cudf::host_span<std::byte const>{
+        reinterpret_cast<std::byte const*>(data.data()), data.size()})}
+  {
+  }
+
+  std::unique_ptr<buffer> host_read(std::size_t offset, std::size_t size) override
+  {
+    auto result = source_->host_read(offset, size);
+    bytes_read_ += result->size();
+    return result;
+  }
+
+  std::size_t host_read(std::size_t offset, std::size_t size, uint8_t* dst) override
+  {
+    auto const result = source_->host_read(offset, size, dst);
+    bytes_read_ += result;
+    return result;
+  }
+
+  [[nodiscard]] std::size_t size() const override { return source_->size(); }
+  [[nodiscard]] std::size_t bytes_read() const { return bytes_read_.load(); }
+  void reset() { bytes_read_ = 0; }
+
+ private:
+  std::unique_ptr<cudf::io::datasource> source_;
+  std::atomic<std::size_t> bytes_read_{0};
+};
+
+/** @brief Reuses the metadata benchmark input, removing only optional index references. */
+std::vector<char> make_optional_index_data(cudf::size_type num_cols,
+                                           cudf::size_type num_row_groups,
+                                           std::string const& layout)
+{
+  auto source_sink         = write_file_data(num_cols, num_row_groups, io_type::HOST_BUFFER, true);
+  auto sources             = cudf::io::make_datasources(source_sink.make_source_info());
+  auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*sources.front());
+  cudf::io::parquet::FileMetaData metadata;
+  cudf::io::parquet::detail::CompactProtocolReader cp(footer_buffer->data(), footer_buffer->size());
+  cp.read(&metadata);
+  CUDF_EXPECTS(layout == "none" or layout == "offset_only" or layout == "mixed" or layout == "both",
+               "Unexpected page index layout");
+  for (auto& rg : metadata.row_groups) {
+    for (auto& col : rg.columns) {
+      if (layout == "none" or layout == "offset_only") {
+        col.column_index_offset = 0;
+        col.column_index_length = 0;
+      }
+      if (layout == "none") {
+        col.offset_index_offset = 0;
+        col.offset_index_length = 0;
+      }
+    }
+  }
+  if (layout == "mixed") {
+    metadata.row_groups.front().columns.front().column_index_offset = 0;
+    metadata.row_groups.front().columns.front().column_index_length = 0;
+  }
+
+  auto const original = sources.front()->host_read(0, sources.front()->size());
+  auto const begin    = reinterpret_cast<char const*>(original->data());
+  std::vector<char> data(begin, begin + original->size());
+  cudf::io::parquet::file_ender_s ender;
+  CUDF_EXPECTS(data.size() >= sizeof(ender), "Invalid Parquet benchmark input");
+  std::memcpy(&ender, data.data() + data.size() - sizeof(ender), sizeof(ender));
+  CUDF_EXPECTS(ender.footer_len <= data.size() - sizeof(ender), "Invalid Parquet benchmark footer");
+  data.resize(data.size() - sizeof(ender) - ender.footer_len);
+  // Keep unused index bytes in place so every layout has the same data-page offsets.
+  std::vector<uint8_t> footer;
+  cudf::io::parquet::detail::CompactProtocolWriter writer(&footer);
+  writer.write(metadata);
+  data.insert(data.end(), footer.begin(), footer.end());
+  ender.footer_len       = static_cast<uint32_t>(footer.size());
+  auto const ender_bytes = reinterpret_cast<char const*>(&ender);
+  data.insert(data.end(), ender_bytes, ender_bytes + sizeof(ender));
+  return data;
+}
+
 // Combines `operands` into a balanced AST tree using `op`: pairing adjacent operands gives a tree
 // of depth ceil(log2(n)) rather than the n-deep chain a left fold would produce.
 [[nodiscard]] cudf::ast::expression const* reduce_balanced(
@@ -343,6 +431,57 @@ void BM_parquet_filter_name_resolution(nvbench::state& state)
     mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage");
 }
 
+/**
+ * @brief Measures metadata parsing, index range calculation and index loading for both readers.
+ *
+ * Reports logical host-read bytes per invocation alongside metadata latency. Input generation is
+ * excluded from timing; host-buffer sources isolate metadata work from storage and cache behavior.
+ */
+void BM_parquet_page_index_metadata(nvbench::state& state)
+{
+  auto const num_cols       = static_cast<cudf::size_type>(state.get_int64("num_cols"));
+  auto const num_row_groups = static_cast<cudf::size_type>(state.get_int64("num_row_groups"));
+  auto const hybrid         = state.get_string("reader") == "hybrid";
+  auto const data = make_optional_index_data(num_cols, num_row_groups, state.get_string("layout"));
+  auto source     = std::make_unique<PageIndexCountingDatasource>(data);
+  auto const counter = source.get();
+  std::vector<std::unique_ptr<cudf::io::datasource>> sources;
+  sources.emplace_back(std::move(source));
+  auto const options = cudf::io::parquet_reader_options::builder().use_arrow_schema(false).build();
+  state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().get()));
+
+  state.exec(
+    nvbench::exec_tag::sync | nvbench::exec_tag::timer, [&](nvbench::launch&, auto& timer) {
+      counter->reset();
+      timer.start();
+      if (hybrid) {
+        auto const footer = cudf::io::parquet::fetch_footer_to_host(*counter);
+        auto const reader = cudf::io::parquet::experimental::hybrid_scan_reader{*footer, options};
+        auto const range  = reader.page_index_byte_range();
+        if (not range.is_empty()) {
+          auto const indexes = cudf::io::parquet::fetch_page_index_to_host(*counter, range);
+          reader.setup_page_index(*indexes);
+        }
+      } else {
+        auto const metadata = cudf::io::read_parquet_footers(sources);
+        CUDF_EXPECTS(std::cmp_equal(metadata.front().row_groups.size(), num_row_groups),
+                     "Unexpected number of row groups");
+      }
+      timer.stop();
+    });
+
+  state.add_buffer_size(counter->bytes_read(), "host_bytes_read", "Logical host bytes read");
+  state.add_buffer_size(data.size(), "file_size", "Parquet file size");
+}
+
+NVBENCH_BENCH(BM_parquet_page_index_metadata)
+  .set_name("parquet_page_index_metadata")
+  .set_min_samples(4)
+  .add_string_axis("layout", {"none", "offset_only", "mixed", "both"})
+  .add_string_axis("reader", {"regular", "hybrid"})
+  .add_int64_axis("num_cols", {4, 16})
+  .add_int64_axis("num_row_groups", {10, 100});
+
 NVBENCH_BENCH(BM_parquet_read_footer)
   .set_name("parquet_read_footer")
   .set_min_samples(4)

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

@@ -31,6 +31,7 @@ using io::detail::inline_column_buffer;
 using parquet::detail::CompactProtocolReader;
 using parquet::detail::equality_literals_collector;
 using parquet::detail::input_column_info;
+using parquet::detail::page_index_byte_range;
 using parquet::detail::row_group_info;
 using text::byte_range_info;
 
@@ -76,45 +77,6 @@ namespace {
   return static_cast<cudf::size_type>(total_row_groups);
 }
 
-// Compute the page index (column index and/or offset index) byte range
-[[nodiscard]] byte_range_info page_index_byte_range(FileMetaData const& file_metadata)
-{
-  auto const& row_groups = file_metadata.row_groups;
-  if (row_groups.empty() or row_groups.front().columns.empty()) { return {}; }
-
-  // Helpers to check if a column chunk has a column index or offset index
-  auto const has_column_index = [](ColumnChunk const& col) {
-    return col.column_index_offset > 0 and col.column_index_length > 0;
-  };
-  auto const has_offset_index = [](ColumnChunk const& col) {
-    return col.offset_index_offset > 0 and col.offset_index_length > 0;
-  };
-
-  auto const min_offset = [&]() -> int64_t {
-    auto const& first_col = row_groups.front().columns.front();
-    if (has_column_index(first_col)) {
-      return first_col.column_index_offset;
-    } else if (has_offset_index(first_col)) {
-      return first_col.offset_index_offset;
-    }
-    return int64_t{0};
-  }();
-
-  auto const max_offset = [&]() -> int64_t {
-    auto const& last_col = row_groups.back().columns.back();
-    if (has_offset_index(last_col)) {
-      return last_col.offset_index_offset + last_col.offset_index_length;
-    } else if (has_column_index(last_col)) {
-      return last_col.column_index_offset + last_col.column_index_length;
-    }
-    return int64_t{0};
-  }();
-
-  return (min_offset > 0 and max_offset > min_offset)
-           ? byte_range_info{min_offset, max_offset - min_offset}
-           : byte_range_info{};
-}
-
 }  // namespace
 
 metadata::metadata(cudf::host_span<uint8_t const> footer_bytes)

cpp/src/io/parquet/reader_impl_helpers.cpp

@@ -37,6 +37,7 @@
 #include <functional>
 #include <future>
 #include <iterator>
+#include <limits>
 #include <numeric>
 #include <optional>
 #include <regex>
@@ -46,6 +47,34 @@
 
 namespace cudf::io::parquet::detail {
 
+// Compute the page index (column index and/or offset index) byte range
+text::byte_range_info page_index_byte_range(FileMetaData const& file_metadata)
+{
+  int64_t min_offset       = std::numeric_limits<int64_t>::max();
+  int64_t max_offset       = 0;
+  auto const include_index = [&](int64_t offset, int32_t length) {
+    if (offset > 0 and length > 0) {
+      CUDF_EXPECTS(offset <= std::numeric_limits<int64_t>::max() - length,
+                   "Parquet page index range exceeds the supported offset range",
+                   std::invalid_argument);
+      min_offset = std::min(min_offset, offset);
+      max_offset = std::max(max_offset, offset + length);
+    }
+  };
+
+  // Indexes are optional for each column chunk. The first and last chunks need not have either
+  // index, so inspect all chunks to include every index that setup_page_index will parse.
+  for (auto const& row_group : file_metadata.row_groups) {
+    for (auto const& column : row_group.columns) {
+      include_index(column.column_index_offset, column.column_index_length);
+      include_index(column.offset_index_offset, column.offset_index_length);
+    }
+  }
+
+  return max_offset > min_offset ? text::byte_range_info{min_offset, max_offset - min_offset}
+                                 : text::byte_range_info{};
+}
+
 std::size_t derive_pass_read_limit(std::size_t chunk_read_limit)
 {
   if (chunk_read_limit == 0) { return 0; }
@@ -529,19 +558,22 @@ metadata::metadata(datasource* source, bool read_page_indexes)
   auto const has_strings = std::any_of(
     schema.begin(), schema.end(), [](auto const& elem) { return elem.type == Type::BYTE_ARRAY; });
 
-  if (read_page_indexes and has_strings and not row_groups.empty() and
-      not row_groups.front().columns.empty()) {
-    // column index and offset index are encoded back to back.
-    // the first column of the first row group will have the first column index, the last
-    // column of the last row group will have the final offset index.
-    int64_t const min_offset = row_groups.front().columns.front().column_index_offset;
-    auto const& last_col     = row_groups.back().columns.back();
-    int64_t const max_offset = last_col.offset_index_offset + last_col.offset_index_length;
-
-    if (max_offset > min_offset) {
-      size_t const length     = max_offset - min_offset;
-      auto const page_idx_buf = source->host_read(min_offset, length);
-      setup_page_index({page_idx_buf->data(), length}, min_offset);
+  // Without offset indexes the decode paths cannot use column-index-derived information.
+  auto const has_offset_index =
+    std::any_of(row_groups.begin(), row_groups.end(), [](auto const& rg) {
+      return std::any_of(rg.columns.begin(), rg.columns.end(), [](auto const& col) {
+        return col.offset_index_offset > 0 and col.offset_index_length > 0;
+      });
+    });
+
+  if (read_page_indexes and has_strings and has_offset_index) {
+    auto const page_index_range = page_index_byte_range(*this);
+    if (not page_index_range.is_empty()) {
+      auto const page_idx_buf =
+        source->host_read(page_index_range.offset(), page_index_range.size());
+      CUDF_EXPECTS(std::cmp_equal(page_idx_buf->size(), page_index_range.size()),
+                   "Encountered an invalid page index buffer");
+      setup_page_index({page_idx_buf->data(), page_idx_buf->size()}, page_index_range.offset());
     }
   }
 

cpp/src/io/parquet/reader_impl_helpers.hpp

@@ -11,6 +11,7 @@
 #include <cudf/io/datasource.hpp>
 #include <cudf/io/parquet.hpp>
 #include <cudf/io/parquet_schema.hpp>
+#include <cudf/io/text/byte_range_info.hpp>
 #include <cudf/types.hpp>
 
 #include <cstddef>
@@ -27,6 +28,16 @@
 
 namespace cudf::io::parquet::detail {
 
+/**
+ * @brief Computes the byte range containing the column and/or offset indexes.
+ *
+ * @throws std::invalid_argument if an index end exceeds the supported offset range
+ *
+ * @param file_metadata Parquet file metadata
+ * @return Page-index byte range, or an empty range when no indexes are available
+ */
+[[nodiscard]] text::byte_range_info page_index_byte_range(FileMetaData const& file_metadata);
+
 /**
  * @brief page location and size info
  */