Add page-level I/O and materialization in Hybrid Scan

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

@@ -307,6 +307,27 @@ class hybrid_scan_multifile {
   payload_column_chunks_byte_ranges(cudf::host_span<std::vector<size_type> const> row_group_indices,
                                     parquet_reader_options const& options) const;
 
+  /**
+   * @brief Get byte ranges of pages of payload columns
+   *
+   * Byte ranges are flattened in source, row group, column chunk, and page order. The returned
+   * source map has one source index per byte range. Dictionary pages precede data pages within each
+   * column chunk. Pruned pages are represented by empty byte ranges.
+   *
+   * @throws cudf::logic_error if any selected column chunk does not have a valid offset index
+   *
+   * @param row_group_indices Input row group indices, one vector per source
+   * @param row_mask Boolean mask spanning the selected row groups
+   * @param options Parquet reader options
+   * @param stream CUDA stream used to compute the page mask
+   * @return Pair of flattened payload page byte ranges and their corresponding source indices
+   */
+  [[nodiscard]] std::pair<std::vector<byte_range_info>, std::vector<size_type>>
+  payload_pages_byte_ranges(cudf::host_span<std::vector<size_type> const> row_group_indices,
+                            cudf::column_view const& row_mask,
+                            parquet_reader_options const& options,
+                            cuda::stream_ref stream) const;
+
   /**
    * @brief Materialize payload columns and applies the row mask to the output table
    *
@@ -430,6 +451,33 @@ class hybrid_scan_multifile {
     cuda::stream_ref stream,
     rmm::device_async_resource_ref mr) const;
 
+  /**
+   * @brief Setup chunking information for payload columns and preprocess the input data pages
+   *
+   * Input page data spans (including empty ones) must have the same shape as the byte ranges
+   * returned by `payload_pages_byte_ranges`. The data page mask is inferred from the data spans.
+   *
+   * @param chunk_read_limit Maximum bytes returned per output table chunk, or zero
+   * @param pass_read_limit Maximum read/decompression memory, or zero
+   * @param row_group_indices Input row group indices, one vector per source
+   * @param row_mask Boolean column spanning all selected rows across all sources and indicating
+   * which rows need to be read
+   * @param page_data Flattened device spans of payload page data in the same order as the byte
+   * ranges from `payload_pages_byte_ranges`
+   * @param options Parquet reader options
+   * @param stream CUDA stream used for preprocessing
+   * @param mr Device memory resource used for output table chunks
+   */
+  void setup_chunking_for_payload_columns(
+    std::size_t chunk_read_limit,
+    std::size_t pass_read_limit,
+    cudf::host_span<std::vector<size_type> const> row_group_indices,
+    cudf::column_view const& row_mask,
+    cudf::host_span<cudf::device_span<uint8_t const> const> page_data,
+    parquet_reader_options const& options,
+    cuda::stream_ref stream,
+    rmm::device_async_resource_ref mr) const;
+
   /**
    * @brief Materializes a chunk of payload columns and applies the corresponding range of input row
    * mask to the output table chunk

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

@@ -38,10 +38,7 @@ void hybrid_scan_reader_impl::handle_chunking(
   // 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);
-
-    // Must be called as soon as we create the pass
-    set_pass_page_mask(data_page_mask);
+    setup_next_pass(column_chunk_data, data_page_mask);
   }
 
   auto& pass = *_pass_itm_data;
@@ -78,7 +75,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<cudf::device_span<uint8_t const> const> column_chunk_data,
+  std::span<bool const> data_page_mask)
 {
   auto const num_passes = _file_itm_data.num_passes();
   CUDF_EXPECTS(num_passes == 1,
@@ -120,7 +118,16 @@ void hybrid_scan_reader_impl::setup_next_pass(
     pass.num_rows  = _file_itm_data.global_num_rows;
 
     // Setup page information for the chunk (which we can access without decompressing)
-    setup_compressed_data(column_chunk_data);
+    if (_sparse_page_io) {
+      CUDF_EXPECTS(data_page_mask.empty(),
+                   "Encountered a non-empty input data page mask in sparse I/O path.",
+                   std::invalid_argument);
+      setup_sparse_compressed_data(column_chunk_data);
+      set_sparse_pass_page_mask(column_chunk_data);
+    } else {
+      setup_compressed_data(column_chunk_data);
+      set_pass_page_mask(data_page_mask);
+    }
 
     // detect malformed columns.
     // - we have seen some cases in the wild where we have a row group containing N
@@ -148,12 +155,20 @@ void hybrid_scan_reader_impl::setup_next_pass(
     // store off how much memory we've used so far. This includes the compressed page data and the
     // decompressed dictionary data. we will subtract this from the available total memory for the
     // subpasses
-    auto chunk_iter =
-      cuda::transform_iterator(pass.chunks.d_begin(), parquet::detail::get_chunk_compressed_size{});
-    pass.base_mem_size =
-      decomp_dict_data_size +
-      cudf::detail::reduce(
+    auto const compressed_data_size = [&] {
+      // In Sparse I/O case, compressed chunk size is the sum of its page data span sizes
+      if (_sparse_page_io) {
+        return std::accumulate(column_chunk_data.begin(),
+                               column_chunk_data.end(),
+                               std::size_t{0},
+                               [](auto size, auto const& page) { return size + page.size(); });
+      }
+      auto chunk_iter = cuda::transform_iterator(pass.chunks.d_begin(),
+                                                 parquet::detail::get_chunk_compressed_size{});
+      return cudf::detail::reduce(
         chunk_iter, chunk_iter + pass.chunks.size(), size_t{0}, cuda::std::plus<size_t>{}, _stream);
+    }();
+    pass.base_mem_size = decomp_dict_data_size + compressed_data_size;
 
     // if we are doing subpass reading, generate more accurate num_row estimates for list columns.
     // this helps us to generate more accurate subpass splits.

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

@@ -47,19 +47,6 @@ 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
    *
@@ -91,6 +78,19 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base {
     cuda::stream_ref stream) const;
 
  public:
+  /**
+   * @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 Constructor for aggregate_reader_metadata
    *