Relax parquet page index requirements in hybrid scan

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

@@ -47,6 +47,19 @@ struct metadata : public metadata_base {
 
 class aggregate_reader_metadata : public aggregate_reader_metadata_base {
  private:
+  /**
+   * @brief Check whether selected columns have column and offset indexes
+   *
+   * Schema indices are mapped to each source before locating the column chunks.
+   *
+   * @param row_group_indices Row group indices, one vector per source
+   * @param schema_indices Schema indices from the first source
+   * @return A pair indicating column-index and offset-index presence, respectively
+   */
+  [[nodiscard]] std::pair<bool, bool> page_index_presence(
+    std::span<std::vector<size_type> const> row_group_indices,
+    std::span<size_type const> schema_indices) const;
+
   /**
    * @brief Filters the row groups using dictionary pages
    *
@@ -287,6 +300,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(

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

@@ -117,9 +117,7 @@ void hybrid_scan_reader_impl::prepare_row_groups(
 
   // Check for offset indexes.
   _has_offset_index =
-    std::all_of(_file_itm_data.row_groups.cbegin(),
-                _file_itm_data.row_groups.cend(),
-                [](auto const& row_group) { return row_group.has_offset_index(); });
+    _extended_metadata->has_offset_index(_file_itm_data.row_groups, _input_columns);
 
   if (_file_itm_data.global_num_rows > 0 && not _file_itm_data.row_groups.empty() &&
       not _input_columns.empty()) {

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

@@ -861,14 +861,6 @@ std::unique_ptr<cudf::column> aggregate_reader_metadata::build_row_mask_with_pag
                "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);
-
-  // Return if page index is not present
-  CUDF_EXPECTS(has_page_index,
-               "Page pruning requires the Parquet page index for all output columns",
-               std::runtime_error);
-
   // Total number of rows
   auto const total_rows = total_rows_in_row_groups(row_group_indices);
   CUDF_EXPECTS(std::cmp_less_equal(total_rows, std::numeric_limits<size_type>::max()),
@@ -885,12 +877,33 @@ std::unique_ptr<cudf::column> aggregate_reader_metadata::build_row_mask_with_pag
       .get_stats_columns_mask();
 
   // Return early if no columns will participate in stats based page filtering
-  if (stats_columns_mask.empty()) {
-    auto const scalar_true =
-      cudf::numeric_scalar<bool>(true, true, stream, cudf::get_current_device_resource_ref());
-    return cudf::make_column_from_scalar(scalar_true, total_rows, stream, mr);
+  if (stats_columns_mask.empty()) { return build_all_true_row_mask(row_group_indices, stream, mr); }
+
+  // Check if we have page index available for all participating columns
+  std::vector<size_type> stats_column_schemas;
+  stats_column_schemas.reserve(num_columns);
+  std::for_each(cuda::counting_iterator<std::size_t>{0},
+                cuda::counting_iterator{num_columns},
+                [&](auto const col_idx) {
+                  auto const& dtype = output_dtypes[col_idx];
+                  if (stats_columns_mask[col_idx] and
+                      (not cudf::is_compound(dtype) or dtype.id() == cudf::type_id::STRING)) {
+                    stats_column_schemas.push_back(output_column_schemas[col_idx]);
+                  }
+                });
+  // Return early if no participating columns
+  if (stats_column_schemas.empty()) {
+    return build_all_true_row_mask(row_group_indices, stream, mr);
   }
 
+  // We need both column and offset indexes to be present for each participating column.
+  auto const [has_column_index, has_offset_index] =
+    page_index_presence(row_group_indices, stats_column_schemas);
+  CUDF_EXPECTS(has_column_index and has_offset_index,
+               "Filter column page pruning using page-statistics requires both column and "
+               "offset indexes to be present",
+               std::runtime_error);
+
   // Optimization for single column filter: Directly build the row mask from page statistics
   if (num_columns == 1) {
     page_stats_to_row_mask_converter const stats_col{static_cast<size_type>(total_rows),
@@ -1005,11 +1018,20 @@ thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(
     return thrust::host_vector<bool>(0, stream);
   }
 
-  auto const has_page_index = compute_has_page_index(per_file_metadata, row_group_indices);
+  // Collect column schema indices from the input columns.
+  auto column_schema_indices = std::vector<size_type>(input_columns.size());
+  std::transform(
+    input_columns.begin(), input_columns.end(), column_schema_indices.begin(), [](auto const& col) {
+      return col.schema_idx;
+    });
 
-  // Return early if page index is not present
-  if (not has_page_index) {
-    CUDF_LOG_WARN("Encountered missing Parquet page index for one or more output columns");
+  // Mapping a row mask to data pages only requires page row locations from the offset index.
+  auto const has_offset_index =
+    page_index_presence(row_group_indices, column_schema_indices).second;
+  if (not has_offset_index) {
+    CUDF_LOG_WARN(
+      "Encountered missing Parquet offset index for one or more output columns. Skipping page "
+      "pruning.");
     return thrust::host_vector<bool>(0, stream);
   }
 
@@ -1019,13 +1041,6 @@ thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(
                "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(
-    input_columns.begin(), input_columns.end(), column_schema_indices.begin(), [](auto const& col) {
-      return col.schema_idx;
-    });
-
   // Compute page row offsets and column chunk page offsets for each column
   auto const num_columns = input_columns.size();
   std::vector<size_type> page_row_offsets;

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

@@ -22,28 +22,6 @@
 
 namespace cudf::io::parquet::experimental::detail {
 
-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)
-{
-  // For all parquet data sources
-  return std::all_of(
-    cuda::counting_iterator<std::size_t>{0},
-    cuda::counting_iterator{row_group_indices.size()},
-    [&](auto const src_index) {
-      // For all row groups in this parquet data source
-      auto const& rg_indices = row_group_indices[src_index];
-      return std::all_of(rg_indices.begin(), rg_indices.end(), [&](auto const& rg_index) {
-        auto const& row_group = file_metadatas[src_index].row_groups[rg_index];
-        return std::any_of(
-          row_group.columns.begin(), row_group.columns.end(), [&](auto const& col) {
-            return col.offset_index.has_value() and col.column_index.has_value();
-          });
-      });
-    });
-}
-
 std::pair<cudf::detail::host_vector<size_type>, cudf::detail::host_vector<size_type>>
 compute_page_row_offsets_and_colchunk_page_offsets(
   std::span<metadata_base const> per_file_metadata,
@@ -79,36 +57,34 @@ compute_page_row_offsets_and_colchunk_page_offsets(
       std::optional<size_type> colchunk_iter_offset{};
       std::for_each(rg_indices.cbegin(), rg_indices.cend(), [&](auto rg_idx) {
         auto const& row_group = per_file_metadata[src_idx].row_groups[rg_idx];
-        if (not colchunk_iter_offset.has_value() or
-            row_group.columns[colchunk_iter_offset.value()].schema_idx != schema_idx) {
-          colchunk_iter_offset = find_colchunk_iter_offset(row_group, schema_idx);
-        }
+        colchunk_iter_offset =
+          parquet::detail::find_colchunk_iter_offset(row_group, schema_idx, colchunk_iter_offset);
         auto const& colchunk_iter = row_group.columns.begin() + colchunk_iter_offset.value();
 
-        // Compute page row offsets if this column chunk has column and offset indexes
-        if (colchunk_iter->offset_index.has_value()) {
-          // Get the offset index of the column chunk
-          auto const& offset_index       = colchunk_iter->offset_index.value();
-          auto const row_group_num_pages = offset_index.page_locations.size();
-
-          col_chunk_page_offsets.push_back(col_chunk_page_offsets.back() + row_group_num_pages);
-
-          // For all pages in this column chunk, update page row offsets.
-          std::for_each(
-            cuda::counting_iterator<std::size_t>{0},
-            cuda::counting_iterator{row_group_num_pages},
-            [&](auto const page_idx) {
-              int64_t const first_row_idx = offset_index.page_locations[page_idx].first_row_index;
-              // For the last page, this is simply the total number of rows in the column chunk
-              int64_t const last_row_idx =
-                (page_idx < row_group_num_pages - 1)
-                  ? offset_index.page_locations[page_idx + 1].first_row_index
-                  : row_group.num_rows;
-
-              // Update the page row offsets.
-              page_row_offsets.push_back(page_row_offsets.back() + last_row_idx - first_row_idx);
-            });
-        }
+        CUDF_EXPECTS(colchunk_iter->offset_index.has_value(),
+                     "Offset index not found for column chunk",
+                     std::invalid_argument);
+
+        auto const& offset_index       = colchunk_iter->offset_index.value();
+        auto const row_group_num_pages = offset_index.page_locations.size();
+
+        col_chunk_page_offsets.push_back(col_chunk_page_offsets.back() + row_group_num_pages);
+
+        // For all pages in this column chunk, update page row offsets.
+        std::for_each(
+          cuda::counting_iterator<std::size_t>{0},
+          cuda::counting_iterator{row_group_num_pages},
+          [&](auto const page_idx) {
+            int64_t const first_row_idx = offset_index.page_locations[page_idx].first_row_index;
+            // For the last page, this is simply the total number of rows in the column chunk
+            int64_t const last_row_idx =
+              (page_idx < row_group_num_pages - 1)
+                ? offset_index.page_locations[page_idx + 1].first_row_index
+                : row_group.num_rows;
+
+            // Update the page row offsets.
+            page_row_offsets.push_back(page_row_offsets.back() + last_row_idx - first_row_idx);
+          });
       });
     });
 
@@ -139,13 +115,13 @@ std::pair<std::vector<size_type>, size_type> compute_page_row_offsets(
                   std::optional<size_type> colchunk_iter_offset{};
                   std::for_each(rg_indices.begin(), rg_indices.end(), [&](auto const& rg_idx) {
                     auto const& row_group = per_file_metadata[src_idx].row_groups[rg_idx];
-                    // Find the column chunk with the given schema index
-                    if (not colchunk_iter_offset.has_value() or
-                        row_group.columns[colchunk_iter_offset.value()].schema_idx != schema_idx) {
-                      colchunk_iter_offset = find_colchunk_iter_offset(row_group, schema_idx);
-                    }
+                    colchunk_iter_offset  = parquet::detail::find_colchunk_iter_offset(
+                      row_group, schema_idx, colchunk_iter_offset);
                     auto const& colchunk_iter =
                       row_group.columns.begin() + colchunk_iter_offset.value();
+                    CUDF_EXPECTS(colchunk_iter->offset_index.has_value(),
+                                 "Offset index not found for column chunk",
+                                 std::invalid_argument);
                     auto const& offset_index       = colchunk_iter->offset_index.value();
                     auto const row_group_num_pages = offset_index.page_locations.size();
                     std::for_each(cuda::counting_iterator<std::size_t>{0},

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

@@ -22,18 +22,6 @@ namespace cudf::io::parquet::experimental::detail {
 
 using metadata_base = parquet::detail::metadata;
 
-/**
- * @brief Compute if the page index is present in all parquet data sources for all columns
- *
- * @param file_metadatas Span of parquet footer metadata
- * @param row_group_indices Span of input row group indices
- * @return Boolean indicating if the page index is present in all parquet data sources for all
- * columns
- */
-[[nodiscard]] bool compute_has_page_index(
-  std::span<metadata_base const> file_metadatas,
-  std::span<std::vector<size_type> const> row_group_indices);
-
 /**
  * @brief Compute page row offsets and column chunk page (count) offsets for a given column schema
  * index

cpp/src/io/parquet/predicate_pushdown.cpp

@@ -146,14 +146,9 @@ bool aggregate_reader_metadata::any_row_group_stats_available(
 
       auto const& first_row_group =
         per_file_metadata[src_idx].row_groups[row_group_indices.front()];
-      auto const num_col_chunks    = static_cast<size_type>(first_row_group.columns.size());
       auto const mapped_schema_idx = map_schema_index(schema_idx, static_cast<int>(src_idx));
-      auto const cached_offset     = colchunk_offset.value_or(-1);
-
-      if (cached_offset < 0 or cached_offset >= num_col_chunks or
-          first_row_group.columns[cached_offset].schema_idx != mapped_schema_idx) {
-        colchunk_offset = find_colchunk_iter_offset(first_row_group, mapped_schema_idx);
-      }
+      colchunk_offset =
+        find_colchunk_iter_offset(first_row_group, mapped_schema_idx, colchunk_offset);
 
       if (colchunk_has_stats(first_row_group.columns[colchunk_offset.value()])) { return true; }
     }

cpp/src/io/parquet/reader_impl_helpers.cpp

@@ -57,8 +57,15 @@ 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)
+size_type find_colchunk_iter_offset(RowGroup const& row_group,
+                                    size_type schema_idx,
+                                    std::optional<size_type> cached_offset)
 {
+  if (cached_offset.has_value() and cached_offset.value() >= 0 and
+      std::cmp_less(cached_offset.value(), row_group.columns.size()) and
+      row_group.columns[cached_offset.value()].schema_idx == schema_idx) {
+    return cached_offset.value();
+  }
   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;
@@ -719,8 +726,9 @@ void aggregate_reader_metadata::column_info_for_row_group(row_group_info& rg_inf
     auto const max_def_level = schema.max_definition_level;
     auto const max_rep_level = schema.max_repetition_level;
 
-    // Return early if any columns lack the offset index.
-    if (not col_chunk.offset_index.has_value()) { return; }
+    // Skip this column chunk if it does not have an offset index. This is because the decode
+    // paths can use column-index-derived information only together with offset index data.
+    if (not col_chunk.offset_index.has_value()) { continue; }
 
     auto const& offset_index = col_chunk.offset_index.value();
 
@@ -822,13 +830,32 @@ void aggregate_reader_metadata::column_info_for_row_group(row_group_info& rg_inf
         }
       }
 
+      // If column-index metadata is insufficient to derive all value information, leave those
+      // fields unset. Later decoding derives the missing values from page headers and levels, and
+      // scans string data when its byte size is unavailable.
+
       chunk_info.pages.push_back(std::move(pg_info));
     }
   }
 
   rg_info.column_chunks = std::move(chunks);
 }
 
+bool aggregate_reader_metadata::has_offset_index(
+  std::span<row_group_info const> row_groups,
+  std::span<input_column_info const> input_columns) const
+{
+  for (auto const& rg_info : row_groups) {
+    auto const& row_group = per_file_metadata[rg_info.source_index].row_groups[rg_info.index];
+    for (auto const& input_column : input_columns) {
+      auto const schema_idx      = map_schema_index(input_column.schema_idx, rg_info.source_index);
+      auto const colchunk_offset = find_colchunk_iter_offset(row_group, schema_idx);
+      if (not row_group.columns[colchunk_offset].offset_index.has_value()) { return false; }
+    }
+  }
+  return true;
+}
+
 void aggregate_reader_metadata::initialize_internals(bool use_arrow_schema,
                                                      bool has_cols_from_mismatched_srcs)
 {
@@ -1216,8 +1243,9 @@ 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 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;
+  auto const& row_group      = per_file_metadata[src_idx].row_groups[row_group_index];
+  auto const colchunk_offset = find_colchunk_iter_offset(row_group, schema_idx);
+  return row_group.columns[colchunk_offset].meta_data;
 }
 
 std::vector<std::unordered_map<std::string, int64_t>>