Relax parquet page index requirements in hybrid scan

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

@@ -31,7 +31,6 @@ 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;
@@ -62,8 +61,107 @@ 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 std::cmp_greater(min_offset, 0) and std::cmp_greater(max_offset, min_offset)
+           ? byte_range_info{min_offset, max_offset - min_offset}
+           : byte_range_info{};
+}
+
+std::pair<bool, bool> compute_page_index_presence(
+  std::span<metadata_base const> file_metadatas,
+  std::span<std::vector<size_type> const> row_group_indices,
+  std::span<size_type const> schema_indices)
+{
+  auto has_column = true;
+  auto has_offset = true;
+
+  auto file_metadata_iter = file_metadatas.begin();
+  for (auto const& rg_indices : row_group_indices) {
+    auto const& file_metadata = *file_metadata_iter++;
+    std::vector<std::optional<size_type>> cached_offsets(schema_indices.size());
+    for (auto const rg_index : rg_indices) {
+      auto const& row_group   = file_metadata.row_groups[rg_index];
+      auto cached_offset_iter = cached_offsets.begin();
+      for (auto const schema_idx : schema_indices) {
+        auto& colchunk_offset = *cached_offset_iter++;
+        auto const has_colchunk =
+          parquet::detail::find_colchunk_iter_offset(row_group, schema_idx, colchunk_offset);
+        auto const has_column_index =
+          has_colchunk and row_group.columns[colchunk_offset.value()].column_index.has_value();
+        auto const has_offset_index =
+          has_colchunk and row_group.columns[colchunk_offset.value()].offset_index.has_value();
+        if (has_column_index and has_offset_index) {
+          auto const& col_chunk = row_group.columns[colchunk_offset.value()];
+          CUDF_EXPECTS(col_chunk.column_index->min_values.size() ==
+                         col_chunk.offset_index->page_locations.size(),
+                       "Column index and offset index page counts must match");
+        }
+        has_column &= has_column_index;
+        has_offset &= has_offset_index;
+      }
+    }
+  }
+  return {has_column, has_offset};
+}
+
 }  // namespace
 
+bool has_column_index(std::span<metadata_base const> file_metadatas,
+                      std::span<std::vector<size_type> const> row_group_indices,
+                      std::span<size_type const> schema_indices)
+{
+  return compute_page_index_presence(file_metadatas, row_group_indices, schema_indices).first;
+}
+
+bool has_offset_index(std::span<metadata_base const> file_metadatas,
+                      std::span<std::vector<size_type> const> row_group_indices,
+                      std::span<size_type const> schema_indices)
+{
+  return compute_page_index_presence(file_metadatas, row_group_indices, schema_indices).second;
+}
+
+bool has_page_index(std::span<metadata_base const> file_metadatas,
+                    std::span<std::vector<size_type> const> row_group_indices,
+                    std::span<size_type const> schema_indices)
+{
+  auto const [has_column, has_offset] =
+    compute_page_index_presence(file_metadatas, row_group_indices, schema_indices);
+  return has_column and has_offset;
+}
+
 metadata::metadata(cudf::host_span<uint8_t const> footer_bytes)
 {
   CUDF_FUNC_RANGE();
@@ -145,16 +243,7 @@ std::vector<text::byte_range_info> aggregate_reader_metadata::page_index_byte_ra
                  per_file_metadata.end(),
                  std::back_inserter(page_index_byte_ranges),
                  [](auto const& file_metadata) -> text::byte_range_info {
-                   auto const& row_groups = file_metadata.row_groups;
-                   if (row_groups.empty() or row_groups.front().columns.empty()) { return {}; }
-
-                   auto const min_offset = row_groups.front().columns.front().column_index_offset;
-                   auto const& last_col  = row_groups.back().columns.back();
-                   auto const max_offset =
-                     last_col.offset_index_offset + last_col.offset_index_length;
-
-                   if (max_offset <= min_offset) { return {}; }
-                   return {min_offset, max_offset - min_offset};
+                   return page_index_byte_range(file_metadata);
                  });
 
   return page_index_byte_ranges;
@@ -184,17 +273,13 @@ void aggregate_reader_metadata::setup_page_indexes(
     CUDF_EXPECTS(not row_groups.empty() and not row_groups.front().columns.empty(),
                  "No column chunks in Parquet schema to read page index for");
 
-    // Set the first ColumnChunk's offset of ColumnIndex as the adjusted zero offset
-    int64_t const min_offset = row_groups.front().columns.front().column_index_offset;
+    auto const expected_byte_range = page_index_byte_range(file_metadata);
 
-    // Check if the page index buffer is valid
-    {
-      auto const& last_col  = row_groups.back().columns.back();
-      auto const max_offset = last_col.offset_index_offset + last_col.offset_index_length;
-      CUDF_EXPECTS(max_offset > min_offset, "Encountered an invalid page index buffer");
-    }
+    CUDF_EXPECTS(not expected_byte_range.is_empty() and
+                   std::cmp_equal(pgidx_bytes.size(), expected_byte_range.size()),
+                 "Encountered an invalid page index buffer");
 
-    file_metadata.setup_page_index(pgidx_bytes, min_offset);
+    file_metadata.setup_page_index(pgidx_bytes, expected_byte_range.offset());
   });
 }
 
@@ -261,6 +346,22 @@ std::size_t aggregate_reader_metadata::total_rows_in_row_groups(
     });
 }
 
+std::unique_ptr<cudf::column> aggregate_reader_metadata::build_all_true_row_mask(
+  std::span<std::vector<size_type> const> row_group_indices,
+  rmm::cuda_stream_view stream,
+  rmm::device_async_resource_ref mr) const
+{
+  CUDF_FUNC_RANGE();
+  auto const num_rows = total_rows_in_row_groups(row_group_indices);
+  CUDF_EXPECTS(num_rows < std::numeric_limits<cudf::size_type>::max(),
+               "Total rows in row groups exceed the cudf's column size limit. Retry with a smaller "
+               "set of row groups",
+               std::invalid_argument);
+  auto true_scalar =
+    cudf::numeric_scalar<bool>(true, true, stream, cudf::get_current_device_resource_ref());
+  return cudf::make_column_from_scalar(true_scalar, num_rows, stream, mr);
+}
+
 std::tuple<std::vector<input_column_info>,
            std::vector<inline_column_buffer>,
            std::vector<cudf::size_type>>
@@ -481,8 +582,7 @@ aggregate_reader_metadata::dictionary_pages_byte_ranges(
       auto const& rg_indices = row_group_indices[src_index];
       // 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];
-        auto const num_col_chunks = static_cast<size_type>(row_group.columns.size());
+        auto const& row_group = per_file_metadata[src_index].row_groups[rg_index];
         // For all dictionary column chunks
         std::for_each(
           cuda::counting_iterator<std::size_t>{0},
@@ -491,47 +591,39 @@ aggregate_reader_metadata::dictionary_pages_byte_ranges(
             // 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& colchunk_offset = colchunk_offsets[col];
+            CUDF_EXPECTS(parquet::detail::find_colchunk_iter_offset(
+                           row_group, mapped_schema_idx, colchunk_offset),
+                         "Column chunk with schema index " + std::to_string(mapped_schema_idx) +
+                           " not found in row group",
+                         std::invalid_argument);
 
             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
-            auto const has_page_index_and_only_dict_encoded_pages = [&]() {
-              auto const has_page_index =
-                col_chunk.offset_index.has_value() and col_chunk.column_index.has_value();
-
-              if (has_page_index and not col_meta.encoding_stats.has_value()) {
+            // Make sure that all column chunk pages are dictionary encoded
+            auto const only_dict_encoded_pages = [&]() {
+              if (not col_meta.encoding_stats.has_value()) {
                 CUDF_LOG_WARN(
                   "Skipping the column chunk because it does not have encoding stats "
                   "needed to determine if all pages are dictionary encoded");
                 return false;
               }
 
-              return has_page_index and
-                     std::all_of(
-                       col_meta.encoding_stats.value().cbegin(),
-                       col_meta.encoding_stats.value().cend(),
-                       [](auto const& page_encoding_stats) {
-                         return page_encoding_stats.page_type == PageType::DICTIONARY_PAGE or
-                                page_encoding_stats.encoding == Encoding::PLAIN_DICTIONARY or
-                                page_encoding_stats.encoding == Encoding::RLE_DICTIONARY;
-                       });
+              return std::all_of(
+                col_meta.encoding_stats.value().cbegin(),
+                col_meta.encoding_stats.value().cend(),
+                [](auto const& page_encoding_stats) {
+                  return page_encoding_stats.page_type == PageType::DICTIONARY_PAGE or
+                         page_encoding_stats.encoding == Encoding::PLAIN_DICTIONARY or
+                         page_encoding_stats.encoding == Encoding::RLE_DICTIONARY;
+                });
             }();
 
             auto dictionary_offset = int64_t{0};
             auto dictionary_size   = int64_t{0};
 
-            if (has_page_index_and_only_dict_encoded_pages) {
-              auto const& offset_index = col_chunk.offset_index.value();
-              auto const num_pages     = offset_index.page_locations.size();
-
+            if (only_dict_encoded_pages) {
               // There is a bug in older versions of parquet-mr where the first data page offset
               // really points to the dictionary page. The first possible offset in a file is 4
               // (after the "PAR1" header), so check to see if the dictionary_page_offset is > 0.
@@ -544,11 +636,14 @@ aggregate_reader_metadata::dictionary_pages_byte_ranges(
                 // dictionary_page_offset is 0, so check to see if the data_page_offset does not
                 // match the first offset in the offset index.  If they don't match, then
                 // data_page_offset points to the dictionary page.
-                if (num_pages > 0 &&
-                    col_meta.data_page_offset < offset_index.page_locations[0].offset) {
+                auto const offset_index = col_chunk.offset_index;
+                auto const num_pages =
+                  offset_index.has_value() ? offset_index->page_locations.size() : size_type{0};
+                if (num_pages > 0 and
+                    col_meta.data_page_offset < offset_index->page_locations[0].offset) {
                   dictionary_offset = col_meta.data_page_offset;
                   dictionary_size =
-                    offset_index.page_locations[0].offset - col_meta.data_page_offset;
+                    offset_index->page_locations[0].offset - col_meta.data_page_offset;
                   have_dictionary_pages = true;
                 }
               }

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

@@ -31,6 +31,27 @@ using parquet::detail::equality_literals_collector;
 using parquet::detail::input_column_info;
 using parquet::detail::row_group_info;
 
+/**
+ * @brief Checks whether column indexes are present for selected columns.
+ */
+[[nodiscard]] bool has_column_index(std::span<metadata_base const> file_metadatas,
+                                    std::span<std::vector<size_type> const> row_group_indices,
+                                    std::span<size_type const> schema_indices);
+
+/**
+ * @brief Checks whether offset indexes are present for selected columns.
+ */
+[[nodiscard]] bool has_offset_index(std::span<metadata_base const> file_metadatas,
+                                    std::span<std::vector<size_type> const> row_group_indices,
+                                    std::span<size_type const> schema_indices);
+
+/**
+ * @brief Checks whether column and offset indexes are present for selected columns.
+ */
+[[nodiscard]] bool has_page_index(std::span<metadata_base const> file_metadatas,
+                                  std::span<std::vector<size_type> const> row_group_indices,
+                                  std::span<size_type const> schema_indices);
+
 /**
  * @brief Class for parsing dataset metadata
  */
@@ -287,6 +308,20 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base {
     std::reference_wrapper<ast::expression const> filter,
     rmm::cuda_stream_view stream) const;
 
+  /**
+   * @brief Builds a row mask with all rows set to true
+   *
+   * @param row_group_indices Input row groups indices
+   * @param stream CUDA stream used for device memory operations and kernel launches
+   * @param mr Device memory resource used to allocate the returned column's device memory
+   *
+   * @return A boolean column representing a mask of rows with all rows set to true
+   */
+  [[nodiscard]] std::unique_ptr<cudf::column> build_all_true_row_mask(
+    std::span<std::vector<size_type> const> row_group_indices,
+    rmm::cuda_stream_view stream,
+    rmm::device_async_resource_ref mr) const;
+
   /**
    * @brief Builds a row mask based on the data pages that survive page-level statistics based on
    * predicate filter

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

@@ -397,14 +397,7 @@ std::unique_ptr<cudf::column> hybrid_scan_reader_impl::build_all_true_row_mask(
 {
   CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered");
 
-  auto const num_rows = total_rows_in_row_groups(row_group_indices);
-  CUDF_EXPECTS(num_rows < std::numeric_limits<cudf::size_type>::max(),
-               "Total rows in row groups exceed the cudf's column size limit. Retry with a smaller "
-               "set of row groups",
-               std::invalid_argument);
-  auto true_scalar =
-    cudf::numeric_scalar<bool>(true, true, stream, cudf::get_current_device_resource_ref());
-  return cudf::make_column_from_scalar(true_scalar, num_rows, stream, mr);
+  return _extended_metadata->build_all_true_row_mask(row_group_indices, stream, mr);
 }
 
 std::unique_ptr<cudf::column> hybrid_scan_reader_impl::build_row_mask_with_page_index_stats(