Add Hybrid scan page pruning when offset index is absent

cpp/examples/hybrid_scan_io/hybrid_scan_composer.cpp

@@ -430,22 +430,6 @@ std::unique_ptr<cudf::table> hybrid_scan(
   }
 }
 
-// Specialization for two-step read without page index
-template <bool single_step_read, bool use_page_index>
-  requires(not single_step_read and not use_page_index)
-std::unique_ptr<cudf::table> inline hybrid_scan(
-  io_source const& io_source,
-  std::optional<cudf::ast::operation const> filter_expression,
-  std::unordered_set<hybrid_scan_filter_type> const& filters,
-  bool verbose,
-  cuda::stream_ref stream,
-  rmm::device_async_resource_ref mr)
-{
-  static_assert(single_step_read or use_page_index,
-                "Hybrid scan requires parquet page index for two-step parquet read");
-  return nullptr;
-}
-
 // Instantiations for hybrid_scan template
 
 template std::unique_ptr<cudf::table> hybrid_scan<true, false>(
@@ -464,6 +448,14 @@ template std::unique_ptr<cudf::table> hybrid_scan<true, true>(
   cuda::stream_ref,
   rmm::device_async_resource_ref);
 
+template std::unique_ptr<cudf::table> hybrid_scan<false, false>(
+  io_source const&,
+  std::optional<cudf::ast::operation const>,
+  std::unordered_set<hybrid_scan_filter_type> const&,
+  bool,
+  cuda::stream_ref,
+  rmm::device_async_resource_ref);
+
 template std::unique_ptr<cudf::table> hybrid_scan<false, true>(
   io_source const&,
   std::optional<cudf::ast::operation const>,

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

@@ -595,7 +595,7 @@ class hybrid_scan_reader {
    *
    * @param row_group_indices Input row groups indices
    * @param column_chunk_data Device spans of column chunk data of filter columns
-   * @param[in,out] row_mask Mutable boolean column indicating surviving rows from page pruning
+   * @param[in,out] row_mask Mutable boolean column indicating surviving rows
    * @param mask_data_pages Whether to build and use a data page mask using the row mask
    * @param options Parquet reader options
    * @param stream CUDA stream used for device memory operations and kernel launches
@@ -668,6 +668,7 @@ class hybrid_scan_reader {
     parquet_reader_options const& options,
     cuda::stream_ref stream,
     rmm::device_async_resource_ref mr) const;
+
   /**
    * @brief Setup chunking information for filter columns and preprocess the input data pages
    *
@@ -676,7 +677,7 @@ class hybrid_scan_reader {
    * @param pass_read_limit Limit on the memory used for reading and decompressing data. `0` if
    * there is no limit
    * @param row_group_indices Input row groups indices
-   * @param row_mask Boolean column indicating which rows need to be read
+   * @param row_mask Boolean column indicating surviving rows
    * @param mask_data_pages Whether to build and use a data page mask using the row mask
    * @param column_chunk_data Device spans of column chunk data of filter columns
    * @param options Parquet reader options

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

@@ -274,7 +274,7 @@ class hybrid_scan_multifile {
    * @param column_chunk_data Flattened device spans of filter column chunk data returned in the
    * same order as `filter_column_chunks_byte_ranges`
    * @param[in,out] row_mask Mutable boolean column spanning all selected rows across all sources
-   * and indicating surviving rows from page pruning
+   * indicating surviving rows
    * @param mask_data_pages Whether to build and use a data page mask using the row mask
    * @param options Parquet reader options
    * @param stream CUDA stream used for device memory operations and kernel launches
@@ -389,8 +389,8 @@ class hybrid_scan_multifile {
    * @param pass_read_limit Limit on the memory used for reading and decompressing data. `0` if
    * there is no limit
    * @param row_group_indices Span of vectors of input row group indices, one per source
-   * @param row_mask Boolean column spanning all selected rows across all sources and indicating
-   * which rows need to be read
+   * @param row_mask Boolean column spanning all selected rows across all sources
+   * indicating surviving rows
    * @param mask_data_pages Whether to build and use a data page mask using the row mask
    * @param column_chunk_data Flattened device spans of filter column chunk data returned in the
    * same order as `filter_column_chunks_byte_ranges`

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

@@ -33,12 +33,13 @@ using parquet::detail::pass_intermediate_data;
 void hybrid_scan_reader_impl::handle_chunking(
   read_mode mode,
   std::span<cudf::device_span<uint8_t const> const> column_chunk_data,
-  host_span<bool const> data_page_mask)
+  host_span<bool const> data_page_mask,
+  std::optional<cudf::column_view> row_mask)
 {
   // if this is our first time in here, setup the first pass.
   if (!_pass_itm_data) {
     // setup the next pass
-    setup_next_pass(column_chunk_data, data_page_mask);
+    setup_next_pass(column_chunk_data, data_page_mask, row_mask);
   }
 
   auto& pass = *_pass_itm_data;
@@ -76,7 +77,8 @@ void hybrid_scan_reader_impl::handle_chunking(
 
 void hybrid_scan_reader_impl::setup_next_pass(
   std::span<cudf::device_span<uint8_t const> const> column_chunk_data,
-  std::span<bool const> data_page_mask)
+  std::span<bool const> data_page_mask,
+  std::optional<cudf::column_view> row_mask)
 {
   auto const num_passes = _file_itm_data.num_passes();
   CUDF_EXPECTS(num_passes == 1,
@@ -126,7 +128,18 @@ void hybrid_scan_reader_impl::setup_next_pass(
       set_sparse_pass_page_mask(column_chunk_data);
     } else {
       setup_compressed_data(column_chunk_data);
-      set_pass_page_mask(data_page_mask);
+      // When offset index is absent, compute and use the data page mask using the decoded page
+      // headers from `setup_compressed_data`.
+      auto const data_page_mask_pghdr = [&]() {
+        if (not _has_offset_index and row_mask.has_value()) {
+          return compute_data_page_mask_with_page_headers(row_mask.value());
+        }
+        return thrust::host_vector<bool>{};
+      }();
+      set_pass_page_mask(
+        data_page_mask_pghdr.empty()
+          ? data_page_mask
+          : std::span<bool const>{data_page_mask_pghdr.data(), data_page_mask_pghdr.size()});
     }
 
     // detect malformed columns.

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

@@ -338,24 +338,18 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base {
    * Compute a vector of boolean vectors indicating which data pages need to be decoded to
    * construct each input column based on the row mask, one vector per column
    *
-   * @tparam ColumnView Type of the row mask column view - cudf::mutable_column_view for filter
-   * columns and cudf::column_view for payload columns
-   *
-   * @param row_mask Boolean column indicating which rows need to be read after page-pruning
+   * @param row_mask Boolean column view indicating surviving rows
    * @param row_group_indices Input row groups indices
    * @param input_columns Input column information
-   * @param row_mask_offset Offset into the row mask column for the current pass
    * @param stream CUDA stream used for device memory operations and kernel launches
    *
    * @return Boolean vector indicating which data pages need to be decoded to produce
    *         the output table based on the input row mask across all input columns
    */
-  template <typename ColumnView>
   [[nodiscard]] thrust::host_vector<bool> compute_data_page_mask(
-    ColumnView const& row_mask,
+    cudf::column_view const& row_mask,
     std::span<std::vector<size_type> const> row_group_indices,
     std::span<input_column_info const> input_columns,
-    cudf::size_type row_mask_offset,
     cuda::stream_ref stream) const;
 };
 

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

@@ -9,6 +9,7 @@
 #include "hybrid_scan_helpers.hpp"
 #include "io/parquet/reader_impl_chunking_utils.cuh"
 #include "io/parquet/synthetic_column_helpers.hpp"
+#include "page_index_filter_utils.hpp"
 
 #include <cudf/copying.hpp>
 #include <cudf/detail/stream_compaction.hpp>
@@ -629,8 +630,8 @@ hybrid_scan_reader_impl::payload_pages_byte_ranges(
 
   // Compute the data page mask
   auto const mask_size = mask_offsets.back();
-  auto data_page_mask  = _extended_metadata->compute_data_page_mask(
-    row_mask, row_group_indices, _input_columns, 0, stream);
+  auto data_page_mask =
+    _extended_metadata->compute_data_page_mask(row_mask, row_group_indices, _input_columns, stream);
   CUDF_EXPECTS(data_page_mask.empty() or data_page_mask.size() == mask_size,
                "Computed data page mask does not match offset indexes");
 
@@ -729,13 +730,10 @@ table_with_metadata hybrid_scan_reader_impl::materialize_filter_columns(
     return read_chunk_internal(read_mode::READ_ALL, read_columns_mode::FILTER_COLUMNS, row_mask);
   }
 
-  auto data_page_mask = thrust::host_vector<bool>{};
-  if (mask_data_pages == use_data_page_mask::YES) {
-    data_page_mask = _extended_metadata->compute_data_page_mask(
-      row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
-  }
-
-  prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, data_page_mask);
+  auto const retention_mask = mask_data_pages == use_data_page_mask::YES
+                                ? std::optional<cudf::column_view>{cudf::column_view{row_mask}}
+                                : std::nullopt;
+  prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, retention_mask);
 
   return read_chunk_internal(read_mode::READ_ALL, read_columns_mode::FILTER_COLUMNS, row_mask);
 }
@@ -767,13 +765,10 @@ table_with_metadata hybrid_scan_reader_impl::materialize_payload_columns(
     return read_chunk_internal(read_mode::READ_ALL, read_columns_mode::PAYLOAD_COLUMNS, row_mask);
   }
 
-  auto data_page_mask = thrust::host_vector<bool>{};
-  if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) {
-    data_page_mask = _extended_metadata->compute_data_page_mask(
-      row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
-  }
-
-  prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, data_page_mask);
+  auto const retention_mask = mask_data_pages == use_data_page_mask::YES
+                                ? std::optional<cudf::column_view>{row_mask}
+                                : std::nullopt;
+  prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, retention_mask);
 
   return read_chunk_internal(read_mode::READ_ALL, read_columns_mode::PAYLOAD_COLUMNS, row_mask);
 }
@@ -838,13 +833,9 @@ void hybrid_scan_reader_impl::setup_chunking_for_filter_columns(
     return;
   }
 
-  auto data_page_mask = thrust::host_vector<bool>{};
-  if (mask_data_pages == use_data_page_mask::YES) {
-    data_page_mask = _extended_metadata->compute_data_page_mask(
-      row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
-  }
-
-  prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, data_page_mask);
+  auto const retention_mask =
+    mask_data_pages == use_data_page_mask::YES ? std::optional{row_mask} : std::nullopt;
+  prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, retention_mask);
 }
 
 table_with_metadata hybrid_scan_reader_impl::materialize_filter_columns_chunk(
@@ -898,13 +889,9 @@ void hybrid_scan_reader_impl::setup_chunking_for_payload_columns(
     return;
   }
 
-  auto data_page_mask = thrust::host_vector<bool>{};
-  if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) {
-    data_page_mask = _extended_metadata->compute_data_page_mask(
-      row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
-  }
-
-  prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, data_page_mask);
+  auto const retention_mask =
+    mask_data_pages == use_data_page_mask::YES ? std::optional{row_mask} : std::nullopt;
+  prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, retention_mask);
 }
 
 void hybrid_scan_reader_impl::setup_chunking_for_payload_columns(
@@ -1109,7 +1096,6 @@ bool hybrid_scan_reader_impl::has_next_table_chunk()
 
 void hybrid_scan_reader_impl::reset_internal_state()
 {
-  _row_mask_offset   = 0;
   _file_itm_data     = file_intermediate_data{};
   _file_preprocessed = false;
   _has_offset_index  = false;
@@ -1133,6 +1119,9 @@ void hybrid_scan_reader_impl::reset_internal_state()
   _output_chunk_read_limit = 0;
   _strings_to_categorical  = false;
   _reader_column_schema.reset();
+
+  _row_mask_offset = 0;
+
   _expr_conv = parquet_filter_normalizer{};
   _mr        = cudf::get_current_device_resource_ref();
 }
@@ -1188,7 +1177,7 @@ void hybrid_scan_reader_impl::prepare_data(
   read_mode mode,
   std::span<std::vector<size_type> const> row_group_indices,
   std::span<cudf::device_span<uint8_t const> const> column_chunk_data,
-  host_span<bool const> data_page_mask)
+  std::optional<cudf::column_view> row_mask)
 {
   // if we have not preprocessed at the whole-file level, do that now
   if (not _file_preprocessed) {
@@ -1199,10 +1188,17 @@ void hybrid_scan_reader_impl::prepare_data(
     prepare_row_groups(read_mode::READ_ALL, row_group_indices);
   }
 
+  // Compute data page mask from the row (retention) mask
+  auto data_page_mask = thrust::host_vector<bool>{};
+  if (_has_offset_index and row_mask.has_value() and not _sparse_page_io) {
+    data_page_mask = _extended_metadata->compute_data_page_mask(
+      row_mask.value(), row_group_indices, _input_columns, _stream);
+  }
+
   // handle any chunking work (ratcheting through the subpasses and chunks within
   // our current pass) if in bounds
   if (_file_itm_data._current_input_pass < _file_itm_data.num_passes()) {
-    handle_chunking(mode, column_chunk_data, data_page_mask);
+    handle_chunking(mode, column_chunk_data, data_page_mask, row_mask);
   }
 }
 
@@ -1354,9 +1350,9 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output(
   // Prepend the source and row index columns to filter columns only
   if (read_columns_mode == read_columns_mode::FILTER_COLUMNS) {
     if (_options.prepend_row_index_column) {
-      out_columns.emplace(
-        out_columns.begin(),
-        synthesize_row_index_column(_file_itm_data.row_groups, read_info, _stream, _mr));
+      out_columns.emplace(out_columns.begin(),
+                          parquet::detail::synthesize_row_index_column(
+                            _file_itm_data.row_groups, read_info, _stream, _mr));
       out_metadata.schema_info.emplace(out_metadata.schema_info.begin(),
                                        column_name_info{.name = "row_index", .is_nullable = false});
     }
@@ -1480,6 +1476,74 @@ void hybrid_scan_reader_impl::set_pass_page_mask(std::span<bool const> data_page
   mark_buffers_nullable_for_pruned_pages();
 }
 
+thrust::host_vector<bool> hybrid_scan_reader_impl::compute_data_page_mask_with_page_headers(
+  cudf::column_view const& row_mask)
+{
+  auto const& pass = *_pass_itm_data;
+
+  // Return an empty vector if all rows are required
+  if (are_all_rows_retained(row_mask, _stream)) { return thrust::host_vector<bool>{}; }
+
+  std::vector<cudf::size_type> page_row_offsets;
+  page_row_offsets.reserve(pass.pages.size() * 2);
+
+  // Maps each data page to its flat-page range; -1 keeps nested pages enabled.
+  std::vector<cudf::size_type> row_range_map;
+  row_range_map.reserve(pass.pages.size());
+
+  cudf::size_type previous_chunk_idx = -1;
+  auto max_page_size                 = cudf::size_type{0};
+
+  for (auto const& page : pass.pages) {
+    // Ignore dictionary pages altogether
+    if (page.flags & parquet::detail::PAGEINFO_FLAGS_DICTIONARY) { continue; }
+
+    auto const& chunk = pass.chunks[page.chunk_idx];
+
+    // Don't prune list column pages as rows may span page boundaries when offset index isn't
+    // present.
+    if (chunk.max_level[parquet::detail::level_type::REPETITION] > 0) {
+      row_range_map.push_back(-1);
+      continue;
+    }
+
+    auto const page_start = chunk.start_row + page.chunk_row;
+    auto const page_end   = page_start + page.num_rows;
+    max_page_size         = std::max<cudf::size_type>(max_page_size, page_end - page_start);
+
+    // Starting a new column chunk. Push page start row
+    if (page.chunk_idx != previous_chunk_idx) {
+      page_row_offsets.push_back(page_start);
+      previous_chunk_idx = page.chunk_idx;
+    }
+
+    // Push row range index and page end row
+    row_range_map.push_back(page_row_offsets.size() - 1);
+    page_row_offsets.push_back(page_end);
+  }
+
+  auto data_page_mask = thrust::host_vector<bool>{};
+
+  // Compute the row range mask
+  CUDF_EXPECTS(std::cmp_equal(row_mask.size(), pass.num_rows),
+               "Row mask must span across all rows in the pass");
+  auto const row_range_mask =
+    compute_row_range_selection_mask(row_mask, page_row_offsets, max_page_size, _stream);
+
+  if (row_range_mask.empty()) { return data_page_mask; }
+
+  CUDF_EXPECTS(row_range_mask.size() == page_row_offsets.size() - 1,
+               "Encountered invalid row range mask size");
+
+  data_page_mask.reserve(row_range_map.size());
+
+  // Scatter row range results while retaining list column pages.
+  for (auto const range_idx : row_range_map) {
+    data_page_mask.push_back(range_idx < 0 ? true : row_range_mask[range_idx]);
+  }
+  return data_page_mask;
+}
+
 void hybrid_scan_reader_impl::set_sparse_pass_page_mask(
   std::span<cudf::device_span<uint8_t const> const> page_data)
 {