Add Hybrid scan APIs for single-step table materialization (#20906)

This PR adds new hybrid scan parquet reader APIs for single-step table materialization

Authors:
  - Muhammad Haseeb (https://github.com/mhaseeb123)
  - https://github.com/apps/pre-commit-ci

Approvers:
  - David Wendt (https://github.com/davidwendt)
  - Bradley Dice (https://github.com/bdice)
  - Nghia Truong (https://github.com/ttnghia)
  - Vukasin Milovanovic (https://github.com/vuule)

URL: #20906

Author: mhaseeb123

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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -499,6 +499,31 @@ class hybrid_scan_reader {
     parquet_reader_options const& options,
     rmm::cuda_stream_view stream) const;
 
+  /**
+   * @brief Get byte ranges of column chunks of all (or selected) columns
+   *
+   * @param row_group_indices Input row groups indices
+   * @param options Parquet reader options
+   * @return Vector of byte ranges to column chunks of all (or selected) columns
+   */
+  [[nodiscard]] std::vector<byte_range_info> all_column_chunks_byte_ranges(
+    cudf::host_span<size_type const> row_group_indices,
+    parquet_reader_options const& options) const;
+
+  /**
+   * @brief Materializes all (or selected) columns and returns the final output table
+   *
+   * @param row_group_indices Input row groups indices
+   * @param column_chunk_buffers Device buffers containing column chunk data of all columns
+   * @param options Parquet reader options
+   * @param stream CUDA stream used for device memory operations and kernel launches
+   * @return Table of all materialized columns and metadata
+   */
+  [[nodiscard]] table_with_metadata materialize_all_columns(
+    cudf::host_span<size_type const> row_group_indices,
+    std::vector<rmm::device_buffer>&& column_chunk_buffers,
+    parquet_reader_options const& options,
+    rmm::cuda_stream_view stream) const;
   /**
    * @brief Setup chunking information for filter columns and preprocess the input data pages
    *

cpp/src/io/parquet/experimental/hybrid_scan.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -232,6 +232,33 @@ table_with_metadata hybrid_scan_reader::materialize_payload_columns(
                                             stream);
 }
 
+std::vector<byte_range_info> hybrid_scan_reader::all_column_chunks_byte_ranges(
+  cudf::host_span<size_type const> row_group_indices, parquet_reader_options const& options) const
+{
+  CUDF_FUNC_RANGE();
+
+  auto const input_row_group_indices =
+    std::vector<std::vector<size_type>>{{row_group_indices.begin(), row_group_indices.end()}};
+
+  return _impl->all_column_chunks_byte_ranges(input_row_group_indices, options).first;
+}
+
+table_with_metadata hybrid_scan_reader::materialize_all_columns(
+  cudf::host_span<size_type const> row_group_indices,
+  std::vector<rmm::device_buffer>&& column_chunk_buffers,
+  parquet_reader_options const& options,
+  rmm::cuda_stream_view stream) const
+{
+  CUDF_FUNC_RANGE();
+
+  // Temporary vector with row group indices from the first source
+  auto const input_row_group_indices =
+    std::vector<std::vector<size_type>>{{row_group_indices.begin(), row_group_indices.end()}};
+
+  return _impl->materialize_all_columns(
+    input_row_group_indices, std::move(column_chunk_buffers), options, stream);
+}
+
 void hybrid_scan_reader::setup_chunking_for_filter_columns(
   std::size_t chunk_read_limit,
   std::size_t pass_read_limit,

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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -105,9 +105,32 @@ void hybrid_scan_reader_impl::setup_page_index(
 void hybrid_scan_reader_impl::select_columns(read_columns_mode read_columns_mode,
                                              parquet_reader_options const& options)
 {
-  // Select only columns required by the filter
-  if (read_columns_mode == read_columns_mode::FILTER_COLUMNS) {
+  if (read_columns_mode == read_columns_mode::ALL_COLUMNS) {
+    if (_is_all_columns_selected) { return; }
+
+    // Select only columns required by the options and filter.
+    // Using as is from:
+    // https://github.com/rapidsai/cudf/blob/a8b25cd205dc5d04b9918dcb0b3abd6b8c4e4a74/cpp/src/io/parquet/reader_impl.cpp#L556-L569
+    std::optional<std::vector<std::string>> filter_only_columns_names;
+    if (options.get_filter().has_value() and options.get_columns().has_value()) {
+      filter_only_columns_names = cudf::io::parquet::detail::get_column_names_in_expression(
+        options.get_filter(), *(options.get_columns()));
+      _num_filter_only_columns = filter_only_columns_names->size();
+    }
+    std::tie(_input_columns, _output_buffers, _output_column_schemas) =
+      _metadata->select_columns(options.get_columns(),
+                                filter_only_columns_names,
+                                options.is_enabled_use_pandas_metadata(),
+                                _strings_to_categorical,
+                                options.is_enabled_ignore_missing_columns(),
+                                _options.timestamp_type.id());
+
+    _is_all_columns_selected     = true;
+    _is_filter_columns_selected  = false;
+    _is_payload_columns_selected = false;
+  } else if (read_columns_mode == read_columns_mode::FILTER_COLUMNS) {
     if (_is_filter_columns_selected) { return; }
+
     // list, struct, dictionary are not supported by AST filter yet.
     _filter_columns_names =
       names_from_expression(
@@ -124,6 +147,7 @@ void hybrid_scan_reader_impl::select_columns(read_columns_mode read_columns_mode
 
     _is_filter_columns_selected  = true;
     _is_payload_columns_selected = false;
+    _is_all_columns_selected     = false;
   } else {
     if (_is_payload_columns_selected) { return; }
 
@@ -137,6 +161,7 @@ void hybrid_scan_reader_impl::select_columns(read_columns_mode read_columns_mode
 
     _is_payload_columns_selected = true;
     _is_filter_columns_selected  = false;
+    _is_all_columns_selected     = false;
   }
 
   // Reset the rows processed so far
@@ -418,6 +443,17 @@ hybrid_scan_reader_impl::payload_column_chunks_byte_ranges(
   return get_input_column_chunk_byte_ranges(row_group_indices);
 }
 
+std::pair<std::vector<byte_range_info>, std::vector<cudf::size_type>>
+hybrid_scan_reader_impl::all_column_chunks_byte_ranges(
+  cudf::host_span<std::vector<size_type> const> row_group_indices,
+  parquet_reader_options const& options)
+{
+  CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered");
+
+  select_columns(read_columns_mode::ALL_COLUMNS, options);
+  return get_input_column_chunk_byte_ranges(row_group_indices);
+}
+
 table_with_metadata hybrid_scan_reader_impl::materialize_filter_columns(
   cudf::host_span<std::vector<size_type> const> row_group_indices,
   std::vector<rmm::device_buffer>&& column_chunk_buffers,
@@ -482,6 +518,29 @@ 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);
 }
 
+table_with_metadata hybrid_scan_reader_impl::materialize_all_columns(
+  cudf::host_span<std::vector<size_type> const> row_group_indices,
+  std::vector<rmm::device_buffer>&& column_chunk_buffers,
+  parquet_reader_options const& options,
+  rmm::cuda_stream_view stream)
+{
+  CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered");
+
+  reset_internal_state();
+
+  initialize_options(row_group_indices, options, stream);
+
+  select_columns(read_columns_mode::ALL_COLUMNS, options);
+
+  // Convert the input expression (must be done after column selection)
+  _expr_conv = build_converted_expression(options.get_filter());
+
+  prepare_data(read_mode::READ_ALL, row_group_indices, std::move(column_chunk_buffers), {});
+
+  // Use the main reader's function
+  return reader_impl::read_chunk_internal(read_mode::READ_ALL);
+}
+
 void hybrid_scan_reader_impl::setup_chunking_for_filter_columns(
   std::size_t chunk_read_limit,
   std::size_t pass_read_limit,
@@ -643,6 +702,8 @@ void hybrid_scan_reader_impl::initialize_options(
 named_to_reference_converter hybrid_scan_reader_impl::build_converted_expression(
   std::optional<std::reference_wrapper<const ast::expression>> filter)
 {
+  if (not filter.has_value()) { return named_to_reference_converter(std::nullopt, {}, {}); }
+
   table_metadata metadata;
   populate_metadata(metadata);
   auto expr_conv =
@@ -821,14 +882,16 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output(
   // Create a table from the output columns.
   auto read_table = std::make_unique<table>(std::move(out_columns));
 
+  CUDF_EXPECTS(row_mask.is_empty() or row_mask.type().id() == type_id::BOOL8,
+               "Input row mask must be empty or a boolean column");
+
   // If the input row mask is empty, return the table as is.
   if (row_mask.is_empty()) { return {std::move(read_table), std::move(out_metadata)}; }
 
-  // If reading filter columns, compute the predicate, apply it to the table, and update the input
-  // row mask to reflect the final surviving rows.
+  // For filter columns, apply the filter expression and update the input row mask
   if constexpr (std::is_same_v<RowMaskView, cudf::mutable_column_view>) {
     CUDF_EXPECTS(read_columns_mode == read_columns_mode::FILTER_COLUMNS, "Invalid read mode");
-    // Apply the row selection predicate on the read table to get the final row mask
+
     auto final_row_mask =
       cudf::detail::compute_column(*read_table,
                                    _expr_conv.get_converted_expr().value().get(),
@@ -850,11 +913,9 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output(
     // Return the final output table and metadata
     return {std::move(output_table), std::move(out_metadata)};
   }
-  // Otherwise, simply apply the input row mask to the table.
+  // For payload columns, simply apply the input row mask to the table.
   else {
     CUDF_EXPECTS(read_columns_mode == read_columns_mode::PAYLOAD_COLUMNS, "Invalid read mode");
-    CUDF_EXPECTS(row_mask.type().id() == type_id::BOOL8,
-                 "Predicate filter should return a boolean");
 
     auto const mask_offset = _rows_processed_so_far;
     _rows_processed_so_far += read_table->num_rows();

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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -180,6 +180,27 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
     parquet_reader_options const& options,
     rmm::cuda_stream_view stream);
 
+  /**
+   * @brief Fetches byte ranges for all (or selected) column chunks
+   *
+   * @param row_group_indices Input row groups indices
+   * @param options Parquet reader options
+   * @return Pair of a vector of byte ranges to column chunks of all (or selected) columns and a
+   * vector of their corresponding input source file indices
+   */
+  [[nodiscard]] std::pair<std::vector<byte_range_info>, std::vector<cudf::size_type>>
+  all_column_chunks_byte_ranges(cudf::host_span<std::vector<size_type> const> row_group_indices,
+                                parquet_reader_options const& options);
+
+  /**
+   * @copydoc cudf::io::experimental::hybrid_scan::materialize_all_columns
+   */
+  [[nodiscard]] table_with_metadata materialize_all_columns(
+    cudf::host_span<std::vector<size_type> const> row_group_indices,
+    std::vector<rmm::device_buffer>&& column_chunk_buffers,
+    parquet_reader_options const& options,
+    rmm::cuda_stream_view stream);
+
   /**
    * @copydoc cudf::io::experimental::hybrid_scan::setup_chunking_for_filter_columns
    */
@@ -242,9 +263,9 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
 
  private:
   /**
-   * @brief The enum indicating whether we are reading the filter columns or the payload columns
+   * @brief The enum indicating whether we are reading the filter, payload, or all columns
    */
-  enum class read_columns_mode { FILTER_COLUMNS, PAYLOAD_COLUMNS };
+  enum class read_columns_mode { FILTER_COLUMNS, PAYLOAD_COLUMNS, ALL_COLUMNS };
 
   /**
    * @brief Initialize the necessary options related internal variables for use later on
@@ -439,6 +460,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
 
   bool _is_filter_columns_selected{false};
   bool _is_payload_columns_selected{false};
+  bool _is_all_columns_selected{false};
 };
 
 }  // namespace cudf::io::parquet::experimental::detail

cpp/tests/io/experimental/hybrid_scan_common.cpp

@@ -1,4 +1,3 @@
-
 /*
  * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
@@ -200,7 +199,7 @@ std::tuple<std::unique_ptr<cudf::table>,
            cudf::io::table_metadata,
            cudf::io::table_metadata,
            std::unique_ptr<cudf::column>>
-hybrid_scan(std::vector<char>& buffer,
+hybrid_scan(cudf::host_span<uint8_t const> file_buffer_span,
             cudf::ast::operation const& filter_expression,
             cudf::size_type num_filter_columns,
             std::optional<std::vector<std::string>> const& payload_column_names,
@@ -215,10 +214,6 @@ hybrid_scan(std::vector<char>& buffer,
   // Set payload column names if provided
   if (payload_column_names.has_value()) { options.set_columns(payload_column_names.value()); }
 
-  // Input file buffer span
-  auto const file_buffer_span =
-    cudf::host_span<uint8_t const>(reinterpret_cast<uint8_t const*>(buffer.data()), buffer.size());
-
   auto [reader, filtered_row_group_indices, row_mask] =
     apply_parquet_filters(file_buffer_span, options, stream, mr);
 
@@ -271,7 +266,7 @@ std::tuple<std::unique_ptr<cudf::table>,
            cudf::io::table_metadata,
            cudf::io::table_metadata,
            std::unique_ptr<cudf::column>>
-chunked_hybrid_scan(std::vector<char> const& buffer,
+chunked_hybrid_scan(cudf::host_span<uint8_t const> file_buffer_span,
                     cudf::ast::operation const& filter_expression,
                     cudf::size_type num_filter_columns,
                     std::optional<std::vector<std::string>> const& payload_column_names,
@@ -286,10 +281,6 @@ chunked_hybrid_scan(std::vector<char> const& buffer,
   // Set payload column names if provided
   if (payload_column_names.has_value()) { options.set_columns(payload_column_names.value()); }
 
-  // Input file buffer span
-  auto const file_buffer_span =
-    cudf::host_span<uint8_t const>(reinterpret_cast<uint8_t const*>(buffer.data()), buffer.size());
-
   auto [reader, filtered_row_group_indices, row_mask] =
     apply_parquet_filters(file_buffer_span, options, stream, mr);
 
@@ -387,3 +378,34 @@ chunked_hybrid_scan(std::vector<char> const& buffer,
                     std::move(payload_metadata),
                     std::move(row_mask)};
 }
+
+cudf::io::table_with_metadata hybrid_scan_single_step(
+  cudf::host_span<uint8_t const> file_buffer_span,
+  std::optional<cudf::ast::operation> filter_expression,
+  std::optional<std::vector<std::string>> const& column_names,
+  rmm::cuda_stream_view stream,
+  rmm::device_async_resource_ref mr)
+{
+  // Create reader options with empty source info
+  cudf::io::parquet_reader_options options = cudf::io::parquet_reader_options::builder().build();
+
+  if (column_names.has_value()) { options.set_columns(column_names.value()); }
+  if (filter_expression.has_value()) { options.set_filter(filter_expression.value()); }
+
+  auto [reader, filtered_row_group_indices, _ /*row_mask*/] =
+    apply_parquet_filters(file_buffer_span, options, stream, mr);
+
+  auto current_row_group_indices = cudf::host_span<cudf::size_type>(filtered_row_group_indices);
+
+  // Get all column chunk byte ranges from the reader
+  auto const all_column_chunk_byte_ranges =
+    reader->all_column_chunks_byte_ranges(current_row_group_indices, options);
+
+  // Fetch column chunk device buffers from the input buffer
+  auto all_column_chunk_buffers =
+    fetch_byte_ranges(file_buffer_span, all_column_chunk_byte_ranges, stream, mr);
+
+  // Materialize the table with all columns
+  return reader->materialize_all_columns(
+    current_row_group_indices, std::move(all_column_chunk_buffers), options, stream);
+}