Add Hybrid scan page pruning when offset index is absent

cpp/examples/hybrid_scan_io/hybrid_scan_composer.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -426,22 +426,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,
-  rmm::cuda_stream_view 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>(
@@ -460,6 +444,14 @@ template std::unique_ptr<cudf::table> hybrid_scan<true, true>(
   rmm::cuda_stream_view,
   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,
+  rmm::cuda_stream_view,
+  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/src/io/parquet/experimental/hybrid_scan_chunking.cu

@@ -126,7 +126,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);
+      // Without an offset index, pages can only be pruned once their headers have been decoded, so
+      // compute the data page mask here rather than using the input one
+      auto const data_page_mask_pghdr = [&]() {
+        if (not _has_offset_index and not _row_mask.is_empty()) {
+          return compute_data_page_mask_with_page_headers();
+        }
+        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_impl.cpp

@@ -8,6 +8,7 @@
 #include "cudf/io/text/byte_range_info.hpp"
 #include "hybrid_scan_helpers.hpp"
 #include "io/parquet/reader_impl_chunking_utils.cuh"
+#include "page_index_filter_utils.hpp"
 
 #include <cudf/copying.hpp>
 #include <cudf/detail/stream_compaction.hpp>
@@ -696,6 +697,7 @@ table_with_metadata hybrid_scan_reader_impl::materialize_filter_columns(
 
   auto data_page_mask = thrust::host_vector<bool>{};
   if (mask_data_pages == use_data_page_mask::YES) {
+    _row_mask      = row_mask;
     data_page_mask = _extended_metadata->compute_data_page_mask(
       row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
   }
@@ -734,6 +736,7 @@ table_with_metadata hybrid_scan_reader_impl::materialize_payload_columns(
 
   auto data_page_mask = thrust::host_vector<bool>{};
   if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) {
+    _row_mask      = row_mask;
     data_page_mask = _extended_metadata->compute_data_page_mask(
       row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
   }
@@ -805,6 +808,7 @@ void hybrid_scan_reader_impl::setup_chunking_for_filter_columns(
 
   auto data_page_mask = thrust::host_vector<bool>{};
   if (mask_data_pages == use_data_page_mask::YES) {
+    _row_mask      = row_mask;
     data_page_mask = _extended_metadata->compute_data_page_mask(
       row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
   }
@@ -865,6 +869,7 @@ void hybrid_scan_reader_impl::setup_chunking_for_payload_columns(
 
   auto data_page_mask = thrust::host_vector<bool>{};
   if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) {
+    _row_mask      = row_mask;
     data_page_mask = _extended_metadata->compute_data_page_mask(
       row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
   }
@@ -1074,7 +1079,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;
@@ -1098,6 +1102,10 @@ void hybrid_scan_reader_impl::reset_internal_state()
   _output_chunk_read_limit = 0;
   _strings_to_categorical  = false;
   _reader_column_schema.reset();
+
+  _row_mask        = column_view{};
+  _row_mask_offset = 0;
+
   _expr_conv = parquet_filter_normalizer{};
   _mr        = cudf::get_current_device_resource_ref();
 }
@@ -1441,6 +1449,70 @@ void hybrid_scan_reader_impl::set_pass_page_mask(std::span<bool const> data_page
                "Encountered mismatch in number of pass pages and page mask size");
 }
 
+thrust::host_vector<bool> hybrid_scan_reader_impl::compute_data_page_mask_with_page_headers()
+{
+  auto& pass = *_pass_itm_data;
+  pass.pages.device_to_host_async(_stream);
+  _stream.sync();
+
+  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 (previous_chunk_idx == -1 or 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
+  auto const row_range_mask = compute_row_range_selection_mask(
+    _row_mask, _row_mask_offset, pass.num_rows, 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)
 {

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

@@ -394,6 +394,11 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
    */
   void set_sparse_pass_page_mask(std::span<cudf::device_span<uint8_t const> const> page_data);
 
+  /**
+   * @brief Compute a data page mask from the decoded page headers.
+   */
+  [[nodiscard]] thrust::host_vector<bool> compute_data_page_mask_with_page_headers();
+
   /**
    * @brief Select the columns to be read based on the read mode
    *
@@ -616,6 +621,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
 
   std::optional<std::vector<std::string>> _filter_columns_names;
 
+  cudf::column_view _row_mask{};
   cudf::size_type _row_mask_offset{0};
   bool _output_chunk_produced{false};
 

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

@@ -978,6 +978,68 @@ std::unique_ptr<cudf::column> aggregate_reader_metadata::build_row_mask_with_pag
     page_stats_table, stats_expr.get_stats_expr().get(), stream, mr);
 }
 
+thrust::host_vector<bool> compute_row_range_selection_mask(
+  cudf::column_view const& row_mask,
+  cudf::size_type row_mask_offset,
+  cudf::size_type total_rows,
+  std::span<cudf::size_type const> page_row_offsets,
+  cudf::size_type max_page_size,
+  cuda::stream_ref stream)
+{
+  // Need at least two offsets (or one range) to search the Fenwick tree
+  if (page_row_offsets.size() < 2) return thrust::host_vector<bool>{};
+
+  // Return early if all rows are needed.
+  if (cudf::detail::all_of(row_mask.begin<bool>() + row_mask_offset,
+                           row_mask.begin<bool>() + row_mask_offset + total_rows,
+                           cuda::std::identity{},
+                           stream)) {
+    return thrust::host_vector<bool>{};
+  }
+
+  auto const mr                 = cudf::get_current_device_resource_ref();
+  auto const tree_level_offsets = compute_fenwick_tree_level_offsets(total_rows, max_page_size);
+  auto const num_levels         = static_cast<cudf::size_type>(tree_level_offsets.size());
+  auto tree_levels_data         = rmm::device_uvector<bool>(tree_level_offsets.back(), stream, mr);
+  auto host_tree_level_ptrs     = cudf::detail::make_pinned_vector_async<bool*>(num_levels, stream);
+  host_tree_level_ptrs[0]       = const_cast<bool*>(row_mask.begin<bool>()) + row_mask_offset;
+  std::for_each(cuda::counting_iterator<cudf::size_type>{1},
+                cuda::counting_iterator<cudf::size_type>{num_levels},
+                [&](auto const level_idx) {
+                  host_tree_level_ptrs[level_idx] =
+                    tree_levels_data.data() + tree_level_offsets[level_idx - 1];
+                });
+  auto tree_level_ptrs = cudf::detail::make_device_uvector_async(host_tree_level_ptrs, stream, mr);
+
+  auto prev_level_size = total_rows;
+  std::for_each(
+    cuda::counting_iterator<cudf::size_type>{0},
+    cuda::counting_iterator<cudf::size_type>{num_levels - 1},
+    [&](auto const prev_level) {
+      auto const current_level_size = cudf::util::div_rounding_up_safe(prev_level_size, 2);
+      thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+                       cuda::counting_iterator<cudf::size_type>{0},
+                       cuda::counting_iterator{current_level_size},
+                       build_fenwick_tree_level_functor{
+                         tree_level_ptrs.data(), prev_level, prev_level_size, current_level_size});
+      prev_level_size = current_level_size;
+    });
+
+  auto const num_ranges    = static_cast<cudf::size_type>(page_row_offsets.size() - 1);
+  auto device_results      = rmm::device_uvector<bool>(num_ranges, stream, mr);
+  auto pinned_page_offsets = cudf::detail::make_pinned_vector(page_row_offsets, stream);
+  auto page_offsets = cudf::detail::make_device_uvector_async(pinned_page_offsets, stream, mr);
+  thrust::transform(
+    rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+    cuda::counting_iterator<cudf::size_type>{0},
+    cuda::counting_iterator{num_ranges},
+    device_results.begin(),
+    search_fenwick_tree_functor{tree_level_ptrs.data(), page_offsets.data(), num_ranges});
+  auto results = cudf::detail::make_pinned_vector_async(device_results, stream);
+  stream.sync();
+  return thrust::host_vector<bool>(results.begin(), results.end());
+}
+
 template <typename ColumnView>
 thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(
   ColumnView const& row_mask,
@@ -1115,79 +1177,24 @@ thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(
                  "Row mask must not contain nulls for payload columns");
   }
 
-  auto const mr = cudf::get_current_device_resource_ref();
-
-  // Compute fenwick tree level offsets and total size (level 1 and higher)
-  auto const tree_level_offsets = compute_fenwick_tree_level_offsets(total_rows, max_page_size);
-  auto const num_levels         = static_cast<cudf::size_type>(tree_level_offsets.size());
-  // Buffer to store Fenwick tree levels (level 1 and higher) data
-  auto tree_levels_data = rmm::device_uvector<bool>(tree_level_offsets.back(), stream, mr);
-
-  // Pointers to each Fenwick tree level data
-  auto host_tree_level_ptrs = cudf::detail::make_pinned_vector_async<bool*>(num_levels, stream);
-  // Zeroth level is just the row mask itself
-  host_tree_level_ptrs[0] = const_cast<bool*>(row_mask.template begin<bool>()) + row_mask_offset;
-  std::for_each(cuda::counting_iterator<cudf::size_type>{1},
-                cuda::counting_iterator{num_levels},
-                [&](auto const level_idx) {
-                  host_tree_level_ptrs[level_idx] =
-                    tree_levels_data.data() + tree_level_offsets[level_idx - 1];
-                });
-
-  auto fenwick_tree_level_ptrs =
-    cudf::detail::make_device_uvector_async(host_tree_level_ptrs, stream, mr);
+  auto data_page_mask = thrust::host_vector<bool>{};
 
-  // Build Fenwick tree levels (zeroth level is just the row mask itself)
-  auto prev_level_size = static_cast<cudf::size_type>(total_rows);
-  std::for_each(
-    cuda::counting_iterator<cudf::size_type>{0},
-    cuda::counting_iterator{num_levels - 1},
-    [&](auto const prev_level) {
-      auto const current_level_size = cudf::util::div_rounding_up_safe(prev_level_size, 2);
-      thrust::for_each(
-        rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
-        cuda::counting_iterator<cudf::size_type>{0},
-        cuda::counting_iterator{current_level_size},
-        build_fenwick_tree_level_functor{
-          fenwick_tree_level_ptrs.data(), prev_level, prev_level_size, current_level_size});
-      prev_level_size = current_level_size;
-    });
-
-  //  Search the Fenwick tree to see if there's a surviving row in each page's row range
-  auto const num_ranges = static_cast<cudf::size_type>(page_row_offsets.size() - 1);
-  rmm::device_uvector<bool> device_data_page_mask(num_ranges, stream, mr);
-  // Use a pinned bounce buffer to avoid pageable h2d copy
-  auto pinned_page_offsets = cudf::detail::make_pinned_vector(
-    cudf::host_span<cudf::size_type const>{page_row_offsets}, stream);
-  auto page_offsets = cudf::detail::make_device_uvector_async(pinned_page_offsets, stream, mr);
-  thrust::transform(
-    rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
-    cuda::counting_iterator<cudf::size_type>{0},
-    cuda::counting_iterator{num_ranges},
-    device_data_page_mask.begin(),
-    search_fenwick_tree_functor{fenwick_tree_level_ptrs.data(), page_offsets.data(), num_ranges});
-
-  //  Copy over search results to host
-  auto host_results      = cudf::detail::make_pinned_vector_async(device_data_page_mask, stream);
-  auto const total_pages = pinned_page_offsets.size() - num_columns;
-  auto data_page_mask    = thrust::host_vector<bool>{};
-  data_page_mask.reserve(total_pages);
-  auto host_results_iter = host_results.begin();
-  stream.sync();
+  auto const row_range_mask = compute_row_range_selection_mask(
+    row_mask, row_mask_offset, total_rows, page_row_offsets, max_page_size, stream);
+  if (row_range_mask.empty()) { return data_page_mask; }
 
+  data_page_mask.reserve(page_row_offsets.size() - num_columns);
   // Discard results for invalid ranges. i.e. ranges starting at the last page of a column and
   // ending at the first page of the next column
-  auto num_pages_inserted = 0;
   std::for_each(cuda::counting_iterator<std::size_t>{0},
                 cuda::counting_iterator{num_columns},
                 [&](auto col_idx) {
                   auto const col_num_pages =
                     col_page_offsets[col_idx + 1] - col_page_offsets[col_idx] - 1;
-                  data_page_mask.insert(data_page_mask.begin() + num_pages_inserted,
-                                        host_results_iter,
-                                        host_results_iter + col_num_pages);
-                  host_results_iter += col_num_pages + 1;
-                  num_pages_inserted += col_num_pages;
+                  auto const first_page_range = col_page_offsets[col_idx];
+                  data_page_mask.insert(data_page_mask.end(),
+                                        row_range_mask.begin() + first_page_range,
+                                        row_range_mask.begin() + first_page_range + col_num_pages);
                 });
   return data_page_mask;
 }