Enable hybrid scan page pruning for page level IO

cpp/CMakeLists.txt

@@ -818,6 +818,7 @@ add_library(
   src/io/parquet/writer_impl.cu
   src/io/parquet/writer_impl_helpers.cpp
   src/io/parquet/decode_fixed.cu
+  src/io/parquet/decode_pruned_pages.cu
   src/io/statistics/orc_column_statistics.cu
   src/io/statistics/parquet_column_statistics.cu
   src/io/text/byte_range_info.cpp

cpp/src/io/parquet/decode_pruned_pages.cu

@@ -0,0 +1,103 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "parquet_gpu.hpp"
+
+#include <cudf/detail/utilities/cuda.cuh>
+
+#include <cooperative_groups.h>
+#include <cuda/std/algorithm>
+
+namespace cudf::io::parquet::detail {
+
+namespace {
+
+auto constexpr block_size = 4 * cudf::detail::warp_size;
+
+/**
+ * @brief Initialize output entries for pruned string and list pages.
+ *
+ * String entries receive either the page's initial offset for small strings or a zero size for
+ * large strings. List offsets receive the start value of their child nesting level. These targeted
+ * writes avoid zero-initializing entire output buffers.
+ */
+CUDF_KERNEL void __launch_bounds__(block_size)
+  fill_pruned_offsets_kernel(device_span<PageInfo> pages,
+                             device_span<ColumnChunkDesc const> chunks,
+                             device_span<bool const> page_mask,
+                             size_t skip_rows,
+                             size_t num_rows)
+{
+  namespace cg = cooperative_groups;
+
+  auto const block    = cg::this_thread_block();
+  auto const page_idx = cg::this_grid().block_rank();
+  auto const t        = static_cast<size_type>(block.thread_rank());
+  if (page_mask[page_idx]) { return; }
+
+  auto const& page  = pages[page_idx];
+  auto const& chunk = chunks[page.chunk_idx];
+
+  if (chunk.column_data_base == nullptr) { return; }
+
+  auto const is_list_col = chunk.max_level[level_type::REPETITION] != 0;
+
+  // Write offsets for pruned non-list (flat) string columns.
+  // Mirrors `update_string_offsets_for_pruned_pages` in page_string_utils.cuh.
+  if (not is_list_col and is_string_col(chunk)) {
+    auto data = static_cast<size_type*>(chunk.column_data_base[chunk.max_nesting_depth - 1]);
+    if (data == nullptr) { return; }
+
+    auto const page_begin = chunk.start_row + page.chunk_row;
+    auto const page_end   = page_begin + page.num_rows;
+    auto const read_end   = skip_rows + num_rows;
+    auto const begin      = cuda::std::max(page_begin, skip_rows);
+    auto const end        = cuda::std::min(page_end, read_end);
+    // Write zeros for large strings and the page's initial offset otherwise.
+    auto const value =
+      chunk.is_large_string_col ? size_type{0} : static_cast<size_type>(page.str_offset);
+    for (auto row = begin + t; row < end; row += block.size()) {
+      data[row - skip_rows] = value;
+    }
+    return;
+  }
+
+  // Write offsets to list locations at each depth.
+  // Mirrors `update_list_offsets_for_pruned_pages` in page_decode.cuh.
+  if (is_list_col and page.nesting != nullptr and page.nesting_decode != nullptr) {
+    for (auto depth = 0; depth < chunk.max_nesting_depth - 1; depth++) {
+      auto offsets       = static_cast<size_type*>(chunk.column_data_base[depth]);
+      auto& nesting_info = page.nesting[depth];
+      // Pruned list pages retain rows through the first list level but contribute no child values.
+      // The preprocessing pass computes the output range and child start value at every list depth.
+      if (nesting_info.type != type_id::LIST or offsets == nullptr) { continue; }
+      // Emit an offset for the current nesting level equal to current length of the next nesting
+      // level
+      auto const output_begin = page.nesting_decode[depth].page_start_value;
+      auto const offset       = page.nesting_decode[depth + 1].page_start_value;
+      for (auto offset_idx = t; offset_idx < nesting_info.batch_size;
+           offset_idx += static_cast<size_type>(block.size())) {
+        offsets[output_begin + offset_idx] = offset;
+      }
+    }
+  }
+}
+
+}  // namespace
+
+void fill_pruned_offsets(cudf::device_span<PageInfo> pages,
+                         cudf::device_span<ColumnChunkDesc const> chunks,
+                         cudf::device_span<bool const> page_mask,
+                         size_t skip_rows,
+                         size_t num_rows,
+                         rmm::cuda_stream_view stream)
+{
+  CUDF_EXPECTS(pages.size() == page_mask.size(), "Page mask size does not match page count");
+  fill_pruned_offsets_kernel<<<pages.size(), block_size, 0, stream.value()>>>(
+    pages, chunks, page_mask, skip_rows, num_rows);
+  CUDF_CUDA_TRY(cudaGetLastError());
+}
+
+}  // namespace cudf::io::parquet::detail

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

@@ -57,7 +57,7 @@ void decode_dictionary_page_headers(cudf::detail::hostdevice_span<ColumnChunkDes
 
   parquet::kernel_error error_code(stream);
 
-  parquet::detail::decode_page_headers(chunks, chunk_page_info.begin(), error_code.data(), stream);
+  parquet::detail::decode_page_headers(chunks, chunk_page_info, error_code.data(), stream);
 
   if (auto const error = error_code.value_sync(stream); error != 0) {
     CUDF_FAIL("Parquet header parsing failed with code(s) " +

cpp/src/io/parquet/page_hdr.cu

@@ -520,6 +520,46 @@ void __forceinline__ __device__ zero_out_page_header_info(byte_stream_s* bs)
   bs->page.kernel_mask    = decode_kernel_mask::NONE;
 }
 
+/**
+ * @brief Decode a page header from an initialized byte stream.
+ *
+ * @param bs Byte stream
+ * @param chunk_idx Index of the chunk containing the page
+ * @param page Pointer to the page info to decode
+ * @param error_code Pointer to the error code for kernel failures
+ */
+__device__ void decode_page_header(byte_stream_s* bs,
+                                   cudf::size_type chunk_idx,
+                                   PageInfo* page,
+                                   kernel_error::pointer error_code)
+{
+  bs->page.chunk_idx      = chunk_idx;
+  bs->page.src_col_schema = bs->ck.src_col_schema;
+
+  if (not parse_valid_page_header(bs)) {
+    set_error(static_cast<kernel_error::value_type>(decode_error::INVALID_PAGE_HEADER), error_code);
+    return;
+  }
+  if (not is_supported_encoding(bs->page.encoding)) {
+    set_error(static_cast<kernel_error::value_type>(decode_error::UNSUPPORTED_ENCODING),
+              error_code);
+    return;
+  }
+
+  switch (bs->page_type) {
+    case PageType::DATA_PAGE: bs->page.num_rows = bs->page.num_input_values; break;
+    case PageType::DATA_PAGE_V2:
+      bs->page.flags |= PAGEINFO_FLAGS_V2;
+      bs->page.definition_level_encoding = Encoding::RLE;
+      bs->page.repetition_level_encoding = Encoding::RLE;
+      break;
+    case PageType::DICTIONARY_PAGE: bs->page.flags |= PAGEINFO_FLAGS_DICTIONARY; break;
+    default:
+      set_error(static_cast<kernel_error::value_type>(decode_error::INVALID_PAGE_TYPE), error_code);
+      return;
+  }
+}
+
 /**
  * @brief Kernel for outputting page headers from the specified column chunks
  *
@@ -530,7 +570,7 @@ void __forceinline__ __device__ zero_out_page_header_info(byte_stream_s* bs)
 CUDF_KERNEL
 void __launch_bounds__(decode_page_headers_block_size)
   decode_page_headers_kernel(device_span<ColumnChunkDesc const> chunks,
-                             chunk_page_info* chunk_pages,
+                             device_span<chunk_page_info> chunk_pages,
                              kernel_error::pointer error_code)
 {
   auto constexpr num_warps_per_block = decode_page_headers_block_size / cudf::detail::warp_size;
@@ -739,26 +779,22 @@ CUDF_KERNEL void __launch_bounds__(count_page_headers_block_size)
 /**
  * @brief Functor to decode page headers from specified page locations
  */
-struct decode_page_headers_with_pgidx_fn {
+struct decode_using_page_index_fn {
   cudf::device_span<ColumnChunkDesc const> colchunks;
   cudf::device_span<PageInfo> pages;
+  cudf::device_span<size_type const> chunk_page_offsets;
   uint8_t** page_locations;
-  size_type* chunk_page_offsets;
   kernel_error::pointer error_code;
 
   __device__ void operator()(size_type page_idx) const noexcept
   {
     auto const num_chunks = static_cast<cudf::size_type>(colchunks.size());
-
-    // Binary search the the column chunk index for this page
-    auto const chunk_idx = static_cast<cudf::size_type>(
+    auto const chunk_idx  = static_cast<cudf::size_type>(
       cuda::std::distance(
-        chunk_page_offsets,
+        chunk_page_offsets.begin(),
         thrust::upper_bound(
-          thrust::seq, chunk_page_offsets, chunk_page_offsets + num_chunks + 1, page_idx)) -
+          thrust::seq, chunk_page_offsets.begin(), chunk_page_offsets.end(), page_idx)) -
       1);
-
-    // Check if the chunk index is valid
     if (chunk_idx < 0 or chunk_idx >= num_chunks) {
       set_error(static_cast<kernel_error::value_type>(decode_error::DATA_STREAM_OVERRUN),
                 error_code);
@@ -778,48 +814,70 @@ struct decode_page_headers_with_pgidx_fn {
     // Clear page header info before writing known fields
     zero_out_page_header_info(&bs);
 
-    bs.page.chunk_idx      = chunk_idx;
-    bs.page.src_col_schema = bs.ck.src_col_schema;
+    decode_page_header(&bs, chunk_idx, &pages[page_idx], error_code);
+
+    bs.page.page_data   = const_cast<uint8_t*>(bs.cur);
+    bs.page.kernel_mask = kernel_mask_for_page(bs.page, bs.ck);
+
+    // Copy the page info to the output span
+    pages[page_idx] = bs.page;
+  }
+};
 
-    // bs.page.chunk_row not computed here and will be filled in later by
-    // `fill_in_page_info()`.
+/**
+ * @brief Functor to decode specified page headers from corresponding page data spans
+ */
+struct decode_from_page_data_fn {
+  cudf::device_span<ColumnChunkDesc const> colchunks;
+  cudf::device_span<PageInfo> pages;
+  cudf::device_span<cudf::device_span<uint8_t const> const> page_data;
+  cudf::device_span<size_type const> chunk_page_offsets;
+  kernel_error::pointer error_code;
 
-    // Parsed page must be valid and not empty
-    if (not parse_valid_page_header(&bs)) {
-      set_error(static_cast<kernel_error::value_type>(decode_error::INVALID_PAGE_HEADER),
+  __device__ void operator()(size_type page_idx) const noexcept
+  {
+    auto const num_chunks = static_cast<cudf::size_type>(colchunks.size());
+    auto const chunk_idx  = static_cast<cudf::size_type>(
+      cuda::std::distance(
+        chunk_page_offsets.begin(),
+        thrust::upper_bound(
+          thrust::seq, chunk_page_offsets.begin(), chunk_page_offsets.end(), page_idx)) -
+      1);
+    if (chunk_idx < 0 or chunk_idx >= num_chunks) {
+      set_error(static_cast<kernel_error::value_type>(decode_error::DATA_STREAM_OVERRUN),
                 error_code);
       return;
     }
-    if (not is_supported_encoding(bs.page.encoding)) {
-      set_error(static_cast<kernel_error::value_type>(decode_error::UNSUPPORTED_ENCODING),
-                error_code);
+
+    byte_stream_s bs{};
+    bs.ck = colchunks[chunk_idx];
+    zero_out_page_header_info(&bs);
+
+    auto const page_span = page_data[page_idx];
+    if (page_span.empty()) {
+      // Initialize the logical page descriptor. Page-index fields are populated by
+      // fill_in_page_info().
+      bs.page.chunk_idx      = chunk_idx;
+      bs.page.src_col_schema = bs.ck.src_col_schema;
+      pages[page_idx]        = bs.page;
       return;
     }
-    switch (bs.page_type) {
-      case PageType::DATA_PAGE:
-        // this computation is only valid for flat schemas. for nested schemas,
-        // they will be recomputed in the preprocess step by examining repetition and
-        // definition levels
-        bs.page.num_rows = bs.page.num_input_values;
-        break;
-      case PageType::DATA_PAGE_V2:
-        bs.page.flags |= PAGEINFO_FLAGS_V2;
-        // V2 only uses RLE, so it was removed from the header
-        bs.page.definition_level_encoding = Encoding::RLE;
-        bs.page.repetition_level_encoding = Encoding::RLE;
-        break;
-      case PageType::DICTIONARY_PAGE: bs.page.flags |= PAGEINFO_FLAGS_DICTIONARY; break;
-      default:
-        set_error(static_cast<kernel_error::value_type>(decode_error::INVALID_PAGE_TYPE),
-                  error_code);
-        return;
+
+    bs.base = bs.cur = page_span.data();
+    bs.end           = page_span.data() + page_span.size();
+
+    decode_page_header(&bs, chunk_idx, &pages[page_idx], error_code);
+
+    if (bs.page.compressed_page_size < 0 or
+        static_cast<size_t>(bs.end - bs.cur) != static_cast<size_t>(bs.page.compressed_page_size)) {
+      set_error(static_cast<kernel_error::value_type>(decode_error::DATA_STREAM_OVERRUN),
+                error_code);
+      return;
     }
 
     bs.page.page_data   = const_cast<uint8_t*>(bs.cur);
     bs.page.kernel_mask = kernel_mask_for_page(bs.page, bs.ck);
-
-    // Copy over the page info from byte stream
-    pages[page_idx] = bs.page;
+    pages[page_idx]     = bs.page;
   }
 };
 
@@ -917,12 +975,14 @@ void count_page_headers(cudf::detail::hostdevice_span<ColumnChunkDesc> chunks,
 }
 
 void decode_page_headers(cudf::device_span<ColumnChunkDesc const> chunks,
-                         chunk_page_info* chunk_pages,
+                         cudf::device_span<chunk_page_info> chunk_pages,
                          kernel_error::pointer error_code,
                          rmm::cuda_stream_view stream)
 {
   static_assert(decode_page_headers_block_size % cudf::detail::warp_size == 0,
                 "Block size for decode page headers kernel must be a multiple of warp size");
+  CUDF_EXPECTS(chunk_pages.size() == chunks.size(),
+               "Chunk page info must contain one entry per chunk");
 
   auto const num_chunks              = static_cast<cudf::size_type>(chunks.size());
   auto constexpr num_warps_per_block = decode_page_headers_block_size / cudf::detail::warp_size;
@@ -937,21 +997,45 @@ void decode_page_headers(cudf::device_span<ColumnChunkDesc const> chunks,
   CUDF_CUDA_TRY(cudaGetLastError());
 }
 
-void decode_page_headers_with_pgidx(cudf::device_span<ColumnChunkDesc const> chunks,
-                                    cudf::device_span<PageInfo> pages,
-                                    uint8_t** page_locations,
-                                    size_type* chunk_page_offsets,
-                                    kernel_error::pointer error_code,
-                                    rmm::cuda_stream_view stream)
+void decode_page_headers_using_page_index(cudf::device_span<ColumnChunkDesc const> chunks,
+                                          cudf::device_span<PageInfo> pages,
+                                          uint8_t** page_locations,
+                                          cudf::device_span<size_type const> chunk_page_offsets,
+                                          kernel_error::pointer error_code,
+                                          rmm::cuda_stream_view stream)
+{
+  CUDF_EXPECTS(chunk_page_offsets.size() == chunks.size() + 1,
+               "Chunk page offsets must cover all chunks");
+  thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+                   cuda::counting_iterator<cudf::size_type>{0},
+                   cuda::counting_iterator{static_cast<cudf::size_type>(pages.size())},
+                   decode_using_page_index_fn{.colchunks          = chunks,
+                                              .pages              = pages,
+                                              .chunk_page_offsets = chunk_page_offsets,
+                                              .page_locations     = page_locations,
+                                              .error_code         = error_code});
+}
+
+void decode_page_headers_from_page_data(
+  cudf::device_span<ColumnChunkDesc const> chunks,
+  cudf::device_span<PageInfo> pages,
+  cudf::device_span<cudf::device_span<uint8_t const> const> page_data,
+  cudf::device_span<size_type const> chunk_page_offsets,
+  kernel_error::pointer error_code,
+  rmm::cuda_stream_view stream)
 {
+  CUDF_EXPECTS(page_data.size() == pages.size(),
+               "Page span count must match the number of logical pages");
+  CUDF_EXPECTS(chunk_page_offsets.size() == chunks.size() + 1,
+               "Chunk page offsets must cover all chunks");
   thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
                    cuda::counting_iterator<cudf::size_type>{0},
                    cuda::counting_iterator{static_cast<cudf::size_type>(pages.size())},
-                   decode_page_headers_with_pgidx_fn{.colchunks          = chunks,
-                                                     .pages              = pages,
-                                                     .page_locations     = page_locations,
-                                                     .chunk_page_offsets = chunk_page_offsets,
-                                                     .error_code         = error_code});
+                   decode_from_page_data_fn{.colchunks          = chunks,
+                                            .pages              = pages,
+                                            .page_data          = page_data,
+                                            .chunk_page_offsets = chunk_page_offsets,
+                                            .error_code         = error_code});
 }
 
 void build_string_dictionary_index(ColumnChunkDesc* chunks,