Add hybrid scan multifile chunked materializers

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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -513,7 +513,7 @@ class hybrid_scan_reader {
    *
    * @param row_group_indices Input row groups indices
    * @param column_chunk_data Device spans of column chunk data of payload columns
-   * @param row_mask Boolean column indicating which rows need to be read. All rows read if empty
+   * @param row_mask Boolean column indicating which rows need to be read
    * @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
@@ -564,7 +564,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. All rows read if empty
+   * @param row_mask Boolean column indicating which rows need to be read
    * @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
@@ -601,7 +601,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. All rows read if empty
+   * @param row_mask Boolean column indicating which rows need to be read
    * @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 payload columns
    * @param options Parquet reader options
@@ -623,7 +623,7 @@ class hybrid_scan_reader {
    * @brief Materializes a chunk of payload columns and applies the corresponding range of input row
    * mask to the output table chunk
    *
-   * @param row_mask Boolean column indicating which rows need to be read. All rows read if empty
+   * @param row_mask Boolean column indicating which rows need to be read
    *
    * @return Table chunk of materialized filter columns and metadata
    */

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

@@ -100,15 +100,15 @@ class hybrid_scan_multifile {
    * @brief Get all available per-source row group indices from the parquet files
    *
    * @param options Parquet reader options
-   * @return Vector of row group indices, one inner vector per source
+   * @return Vector of vectors of row group indices, one per source
    */
   [[nodiscard]] std::vector<std::vector<size_type>> all_row_groups(
     parquet_reader_options const& options) const;
 
   /**
    * @brief Get the total number of top-level rows in the per-source row groups
    *
-   * @param row_group_indices Input per-source row group indices (one inner vector per source)
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @return Total number of top-level rows across all sources
    */
   [[nodiscard]] size_type total_rows_in_row_groups(
@@ -130,9 +130,9 @@ class hybrid_scan_multifile {
    * Filters the row groups such that only the row groups that start within the byte range are
    * selected. Note that the last selected row group may end beyond the byte range.
    *
-   * @param row_group_indices Input row group indices, one per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param options Parquet reader options
-   * @return Filtered per-source row group indices (one inner vector per source)
+   * @return Vector of vectors of filtered row group indices, one per source
    */
   [[nodiscard]] std::vector<std::vector<size_type>> filter_row_groups_with_byte_range(
     cudf::host_span<std::vector<size_type> const> row_group_indices,
@@ -141,10 +141,10 @@ class hybrid_scan_multifile {
   /**
    * @brief Filter the input row groups using column chunk statistics
    *
-   * @param row_group_indices Input row group indices, one per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param options Parquet reader options
    * @param stream CUDA stream used for device memory operations and kernel launches
-   * @return Filtered row group indices, one per source
+   * @return Vector of vectors of filtered row group indices, one per source
    */
   [[nodiscard]] std::vector<std::vector<size_type>> filter_row_groups_with_stats(
     cudf::host_span<std::vector<size_type> const> row_group_indices,
@@ -158,7 +158,7 @@ class hybrid_scan_multifile {
    * @note Device buffers for bloom filter byte ranges must be allocated using a 32 byte
    *       aligned memory resource
    *
-   * @param row_group_indices Input row group indices, one per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param options Parquet reader options
    * @return Pair of vectors of byte ranges of column chunk with bloom filters and dictionary
    *         pages subject to filter predicate
@@ -171,7 +171,7 @@ class hybrid_scan_multifile {
    * @brief Builds a boolean survival column of size equal to the total number of rows in the row
    * groups containing all `true` values
    *
-   * @param row_group_indices Input per-source row group indices (one inner vector per source)
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @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 An all-true boolean (survival) column spanning all selected rows across all sources
@@ -185,7 +185,7 @@ class hybrid_scan_multifile {
    * @brief Builds a boolean column indicating surviving rows using page-level statistics in the
    * page index
    *
-   * @param row_group_indices Input per-source row group indices (one inner vector per source)
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param options Parquet reader options
    * @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
@@ -206,7 +206,7 @@ class hybrid_scan_multifile {
    * map has one source index per byte range and can be used to regroup byte ranges by datasource
    * before fetching.
    *
-   * @param row_group_indices Input row group indices, one inner vector per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param options Parquet reader options
    * @return Pair of flattened byte ranges to column chunks of filter columns and their
    *         corresponding source indices
@@ -219,7 +219,7 @@ class hybrid_scan_multifile {
    * @brief Materializes filter columns and updates the input row mask to only the rows that exist
    * in the output table
    *
-   * @param row_group_indices Input row group indices, one inner vector per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @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
@@ -247,7 +247,7 @@ class hybrid_scan_multifile {
    * map has one source index per byte range and can be used to regroup byte ranges by datasource
    * before fetching.
    *
-   * @param row_group_indices Input row group indices, one inner vector per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param options Parquet reader options
    * @return Pair of flattened byte ranges to column chunks of payload columns and their
    *         corresponding source indices
@@ -259,11 +259,11 @@ class hybrid_scan_multifile {
   /**
    * @brief Materialize payload columns and applies the row mask to the output table
    *
-   * @param row_group_indices Input row group indices, one inner vector per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param column_chunk_data Flattened device spans of payload column chunk data returned in the
    *        same order as `payload_column_chunks_byte_ranges`
    * @param row_mask Boolean column spanning all selected rows across all sources and indicating
-   *        which rows need to be read. All rows read if empty
+   *        which rows need to be read
    * @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
@@ -282,7 +282,7 @@ class hybrid_scan_multifile {
   /**
    * @brief Get byte ranges of column chunks of all (or selected) columns
    *
-   * @param row_group_indices Input row group indices, one inner vector per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param options Parquet reader options
    * @return Pair of flattened byte ranges to column chunks of all (or selected) columns and their
    *         corresponding source indices
@@ -294,7 +294,7 @@ class hybrid_scan_multifile {
   /**
    * @brief Materializes all (or selected) columns and returns the final output table
    *
-   * @param row_group_indices Input row group indices, one inner vector per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param column_chunk_data Flattened device spans of column chunk data returned in the same order
    *        as `all_column_chunks_byte_ranges`
    * @param options Parquet reader options
@@ -309,6 +309,119 @@ class hybrid_scan_multifile {
     rmm::cuda_stream_view stream,
     rmm::device_async_resource_ref mr) const;
 
+  /**
+   * @brief Setup chunking information for filter columns and preprocess the input data pages
+   *
+   * @param chunk_read_limit Limit on total number of bytes to be returned per table chunk. `0` if
+   * there is no limit
+   * @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 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`
+   * @param options Parquet reader options
+   * @param stream CUDA stream used for device memory operations and kernel launches
+   * @param mr Device memory resource used to allocate the device memory for the output table chunks
+   */
+  void setup_chunking_for_filter_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,
+    use_data_page_mask mask_data_pages,
+    cudf::host_span<cudf::device_span<uint8_t const> const> column_chunk_data,
+    parquet_reader_options const& options,
+    rmm::cuda_stream_view stream,
+    rmm::device_async_resource_ref mr) const;
+
+  /**
+   * @brief Materializes a chunk of filter columns and updates the corresponding range of input row
+   * mask to only the rows that exist in the output table
+   *
+   * @param[in,out] row_mask Mutable boolean column spanning all selected rows across all sources
+   * and indicating surviving rows from page pruning. The row mask size must equal the total
+   * number of rows in the input row groups, and is empty only when there are no such rows
+   * (yielding an empty output table)
+   *
+   * @return Table chunk of materialized filter columns and metadata
+   */
+  [[nodiscard]] table_with_metadata materialize_filter_columns_chunk(
+    cudf::mutable_column_view& row_mask) const;
+
+  /**
+   * @brief Setup chunking information for payload columns and preprocess the input data pages
+   *
+   * @param chunk_read_limit Limit on total number of bytes to be returned per table chunk. `0` if
+   * there is no limit
+   * @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 mask_data_pages Whether to build and use a data page mask using the row mask
+   * @param column_chunk_data Flattened device spans of payload column chunk data returned in the
+   * same order as `payload_column_chunks_byte_ranges`
+   * @param options Parquet reader options
+   * @param stream CUDA stream used for device memory operations and kernel launches
+   * @param mr Device memory resource used to allocate the device memory for the 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,
+    use_data_page_mask mask_data_pages,
+    cudf::host_span<cudf::device_span<uint8_t const> const> column_chunk_data,
+    parquet_reader_options const& options,
+    rmm::cuda_stream_view 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
+   *
+   * @param row_mask Boolean column spanning all selected rows across all sources and indicating
+   * which rows need to be read
+   *
+   * @return Table chunk of materialized payload columns and metadata
+   */
+  [[nodiscard]] table_with_metadata materialize_payload_columns_chunk(
+    cudf::column_view const& row_mask) const;
+
+  /**
+   * @brief Setup chunking information for all (or selected) columns and preprocess the input data
+   * pages
+   *
+   * @param chunk_read_limit Limit on total number of bytes to be returned per table chunk. `0` if
+   * there is no limit
+   * @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 column_chunk_data Flattened device spans of column chunk data returned in the same order
+   * as `all_column_chunks_byte_ranges`
+   * @param options Parquet reader options
+   * @param stream CUDA stream used for device memory operations and kernel launches
+   * @param mr Device memory resource used to allocate the device memory for the output table chunks
+   */
+  void setup_chunking_for_all_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::host_span<cudf::device_span<uint8_t const> const> column_chunk_data,
+    parquet_reader_options const& options,
+    rmm::cuda_stream_view stream,
+    rmm::device_async_resource_ref mr) const;
+
+  /**
+   * @brief Materializes a chunk of all (or selected) columns and returns the output table chunk
+   *
+   * @return Table chunk of materialized all (or selected) columns and metadata
+   */
+  [[nodiscard]] table_with_metadata materialize_all_columns_chunk() const;
+
   /**
    * @brief Partition row groups into passes such that the amount of GPU memory required to read,
    * decompress and decode a pass is bounded by the specified limit
@@ -320,7 +433,7 @@ class hybrid_scan_multifile {
    *
    * @throws std::invalid_argument if no row group indices in the input
    *
-   * @param row_group_indices Input row group indices, one per source
+   * @param row_group_indices Span of vectors of input row group indices, one per source
    * @param pass_read_limit Memory limit to read and decompress row group data, `0` if there is
    * no limit (single pass)
    *
@@ -330,6 +443,13 @@ class hybrid_scan_multifile {
     cudf::host_span<std::vector<size_type> const> row_group_indices,
     std::size_t pass_read_limit) const;
 
+  /**
+   * @brief Check if there is any parquet data left to read for the current chunked setup
+   *
+   * @return Boolean indicating if there is any data left to read
+   */
+  [[nodiscard]] bool has_next_table_chunk() const;
+
  private:
   std::unique_ptr<detail::hybrid_scan_reader_impl> _impl;
 };