Add multifile hybrid scan pass construction API

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

@@ -669,11 +669,11 @@ class hybrid_scan_reader {
    * size is estimated over all columns in each row group (not just the columns selected for
    * reading), for conservative estimates.
    *
-   * @throws cudf::logic_error if `row_group_indices` is empty
+   * @throws std::invalid_argument if no row group indices in the input
    *
    * @param row_group_indices Input row group indices
-   * @param pass_read_limit Limit on the amount of memory used for reading and decompressing row
-   * group data or `0` if there is no limit
+   * @param pass_read_limit Memory limit to read and decompress row group data, `0` if there is
+   * no limit (single pass)
    *
    * @return Vector of vectors of row group indices, one per constructed pass
    */

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

@@ -228,6 +228,27 @@ class hybrid_scan_multifile {
     rmm::cuda_stream_view stream,
     rmm::device_async_resource_ref mr) 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
+   *
+   * Note that the `pass_read_limit` is a hint, not an absolute limit - if a single row group
+   * cannot fit within the limit given, it will still constitute a pass. The compressed row group
+   * size is estimated over all columns in each row group (not just the columns selected for
+   * reading), for conservative estimates.
+   *
+   * @throws std::invalid_argument if no row group indices in the input
+   *
+   * @param row_group_indices 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)
+   *
+   * @return Vector of per-source row group indices, one per constructed pass
+   */
+  [[nodiscard]] std::vector<std::vector<std::vector<size_type>>> construct_row_group_passes(
+    cudf::host_span<std::vector<size_type> const> row_group_indices,
+    std::size_t pass_read_limit) const;
+
  private:
   std::unique_ptr<detail::hybrid_scan_reader_impl> _impl;
 };

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

@@ -370,7 +370,17 @@ table_with_metadata hybrid_scan_reader::materialize_all_columns_chunk() const
 std::vector<std::vector<cudf::size_type>> hybrid_scan_reader::construct_row_group_passes(
   cudf::host_span<cudf::size_type const> row_group_indices, std::size_t pass_read_limit) const
 {
-  return _impl->construct_row_group_passes(row_group_indices, pass_read_limit);
+  auto const total_row_groups = row_group_indices.size();
+
+  CUDF_EXPECTS(
+    total_row_groups > 0, "Empty input row group indices encountered", std::invalid_argument);
+
+  auto const input_row_group_indices =
+    std::vector<std::vector<size_type>>{{row_group_indices.begin(), row_group_indices.end()}};
+
+  return _impl
+    ->construct_row_group_passes(input_row_group_indices, total_row_groups, pass_read_limit)
+    .first;
 }
 
 bool hybrid_scan_reader::has_next_table_chunk() const { return _impl->has_next_table_chunk(); }

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

@@ -723,40 +723,55 @@ table_with_metadata hybrid_scan_reader_impl::materialize_all_columns_chunk()
   return result;
 }
 
-std::vector<std::vector<cudf::size_type>> hybrid_scan_reader_impl::construct_row_group_passes(
-  cudf::host_span<cudf::size_type const> row_group_indices, std::size_t pass_read_limit) const
+std::pair<std::vector<std::vector<cudf::size_type>>, std::vector<cudf::size_type>>
+hybrid_scan_reader_impl::construct_row_group_passes(
+  cudf::host_span<std::vector<size_type> const> row_group_indices,
+  std::size_t total_row_groups,
+  std::size_t pass_read_limit) const
 {
-  CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered");
+  CUDF_EXPECTS(
+    total_row_groups > 0, "Empty input row group indices encountered", std::invalid_argument);
+
+  CUDF_EXPECTS(row_group_indices.size() == _extended_metadata->get_num_sources(),
+               "Mismatch in the number of row group indices vectors and the number of input "
+               "datasources",
+               std::invalid_argument);
 
-  // If pass_read_limit is 0 or there is only one row group, return all in a single pass
-  if (pass_read_limit == 0 or row_group_indices.size() == 1) {
-    return {{row_group_indices.begin(), row_group_indices.end()}};
+  if (pass_read_limit == 0) {
+    return {
+      std::vector<std::vector<cudf::size_type>>{row_group_indices.begin(), row_group_indices.end()},
+      std::vector<cudf::size_type>{}};
   }
 
-  // TODO(mh): Need to handle multiple sources in the future
-  auto constexpr source_index = 0;
+  CUDF_EXPECTS(
+    pass_read_limit > 0, "Pass read limit must be greater than 0", std::invalid_argument);
 
-  // Construct row group information
   auto row_groups_info = std::vector<row_group_info>{};
-  row_groups_info.reserve(row_group_indices.size());
+  row_groups_info.reserve(total_row_groups);
   size_t start_row = 0;
-  std::transform(row_group_indices.begin(),
-                 row_group_indices.end(),
-                 std::back_inserter(row_groups_info),
-                 [&](auto const& rg_index) {
-                   auto const& row_group =
-                     _extended_metadata->get_row_group(rg_index, source_index);
-                   auto const [compressed_size, total_size, num_rows, max_leaf_values] =
-                     _extended_metadata->get_row_group_properties(row_group);
-                   auto rg_info = row_group_info{.index               = rg_index,
-                                                 .start_row           = start_row,
-                                                 .unadjusted_num_rows = num_rows,
-                                                 .source_index        = source_index,
-                                                 .compressed_size     = compressed_size,
-                                                 .max_leaf_values     = max_leaf_values};
-                   start_row += num_rows;
-                   return rg_info;
-                 });
+  std::for_each(cuda::counting_iterator<cudf::size_type>(0),
+                cuda::counting_iterator<cudf::size_type>(row_group_indices.size()),
+                [&](auto const source_index) {
+                  auto const& src_row_groups = row_group_indices[source_index];
+                  std::transform(
+                    src_row_groups.begin(),
+                    src_row_groups.end(),
+                    std::back_inserter(row_groups_info),
+                    [&](auto const rg_index) {
+                      auto const& row_group =
+                        _extended_metadata->get_row_group(rg_index, source_index);
+                      auto const [compressed_size, total_size, num_rows, max_leaf_values] =
+                        _extended_metadata->get_row_group_properties(row_group);
+                      auto rg_info = row_group_info{.index               = rg_index,
+                                                    .start_row           = start_row,
+                                                    .unadjusted_num_rows = num_rows,
+                                                    .source_index        = source_index,
+                                                    .compressed_size     = compressed_size,
+                                                    .max_leaf_values     = max_leaf_values};
+                      start_row += num_rows;
+                      return rg_info;
+                    });
+                });
 
   auto const comp_read_limit = static_cast<std::size_t>(
     pass_read_limit * cudf::io::parquet::detail::input_limit_compression_reserve);
@@ -768,15 +783,27 @@ std::vector<std::vector<cudf::size_type>> hybrid_scan_reader_impl::construct_row
   auto const& offsets = pass_data.pass_row_group_offsets;
   auto passes         = std::vector<std::vector<cudf::size_type>>{};
   passes.reserve(offsets.size() - 1);
+  auto row_group_source_map       = std::vector<cudf::size_type>{};
+  auto const has_multiple_sources = row_group_indices.size() > 1;
+  if (has_multiple_sources) { row_group_source_map.reserve(row_groups_info.size()); }
   std::transform(offsets.begin(),
                  offsets.end() - 1,
                  offsets.begin() + 1,
                  std::back_inserter(passes),
                  [&](auto const start, auto const end) {
-                   return std::vector<cudf::size_type>{row_group_indices.begin() + start,
-                                                       row_group_indices.begin() + end};
+                   auto pass = std::vector<cudf::size_type>{};
+                   pass.reserve(end - start);
+                   std::for_each(row_groups_info.begin() + start,
+                                 row_groups_info.begin() + end,
+                                 [&](auto const& rg_info) {
+                                   pass.emplace_back(rg_info.index);
+                                   if (has_multiple_sources) {
+                                     row_group_source_map.emplace_back(rg_info.source_index);
+                                   }
+                                 });
+                   return pass;
                  });
-  return passes;
+  return {std::move(passes), std::move(row_group_source_map)};
 }
 
 bool hybrid_scan_reader_impl::has_next_table_chunk()

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

@@ -25,6 +25,7 @@
 
 #include <memory>
 #include <optional>
+#include <utility>
 #include <vector>
 
 namespace cudf::io::parquet::experimental::detail {
@@ -270,10 +271,23 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
   [[nodiscard]] table_with_metadata materialize_all_columns_chunk();
 
   /**
-   * @copydoc cudf::io::experimental::hybrid_scan_reader::construct_row_group_passes
+   * @brief Partition per-source row groups into read passes
+   *
+   * @throws std::invalid_argument if @p row_group_indices.size() is all empty or not equal to the
+   * number of input datasources
+   *
+   * @param row_group_indices Input row group indices, one per source
+   * @param total_row_groups Total number of row groups across all sources
+   * @param pass_read_limit Memory limit to read and decompress row
+   * group data
+   *
+   * @return Pair of a vector of flattened row group passes and a source index map. The source index
+   * map is empty for single source input
    */
-  [[nodiscard]] std::vector<std::vector<cudf::size_type>> construct_row_group_passes(
-    cudf::host_span<cudf::size_type const> row_group_indices, std::size_t pass_read_limit) const;
+  [[nodiscard]] std::pair<std::vector<std::vector<cudf::size_type>>, std::vector<cudf::size_type>>
+  construct_row_group_passes(cudf::host_span<std::vector<size_type> const> row_group_indices,
+                             std::size_t total_row_groups,
+                             std::size_t pass_read_limit) const;
 
   /**
    * @copydoc cudf::io::experimental::hybrid_scan::has_next_table_chunk

cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp

@@ -7,6 +7,9 @@
 
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/io/experimental/hybrid_scan_multifile.hpp>
+#include <cudf/utilities/error.hpp>
+
+#include <numeric>
 
 namespace cudf::io::parquet::experimental {
 
@@ -122,4 +125,44 @@ table_with_metadata hybrid_scan_multifile::materialize_all_columns(
   return _impl->materialize_all_columns(row_group_indices, column_chunk_data, options, stream, mr);
 }
 
+std::vector<std::vector<std::vector<size_type>>> hybrid_scan_multifile::construct_row_group_passes(
+  cudf::host_span<std::vector<size_type> const> row_group_indices,
+  std::size_t pass_read_limit) const
+{
+  auto const total_row_groups =
+    std::accumulate(row_group_indices.begin(),
+                    row_group_indices.end(),
+                    std::size_t{0},
+                    [](auto sum, auto const& rgs) { return sum + rgs.size(); });
+  CUDF_EXPECTS(
+    total_row_groups > 0, "Empty input row group indices encountered", std::invalid_argument);
+
+  auto [passes, source_map] =
+    _impl->construct_row_group_passes(row_group_indices, total_row_groups, pass_read_limit);
+
+  if (pass_read_limit == 0) { return {passes}; }
+
+  auto source_passes = std::vector<std::vector<std::vector<size_type>>>{};
+  source_passes.reserve(passes.size());
+
+  if (row_group_indices.size() == 1) {
+    for (auto& pass : passes) {
+      source_passes.emplace_back();
+      source_passes.back().push_back(std::move(pass));
+    }
+    return source_passes;
+  }
+
+  auto source_map_it = source_map.begin();
+  for (auto const& pass : passes) {
+    auto source_pass = std::vector<std::vector<size_type>>(row_group_indices.size());
+    for (auto const row_group_index : pass) {
+      source_pass[*source_map_it++].push_back(row_group_index);
+    }
+    source_passes.push_back(std::move(source_pass));
+  }
+
+  return source_passes;
+}
+
 }  // namespace cudf::io::parquet::experimental

cpp/src/io/parquet/reader_impl_helpers.hpp

@@ -411,6 +411,13 @@ class aggregate_reader_metadata {
    */
   [[nodiscard]] auto get_num_row_groups() const { return num_row_groups; }
 
+  /**
+   * @brief Get total number of sources
+   *
+   * @return Total number of sources
+   */
+  [[nodiscard]] auto get_num_sources() const { return per_file_metadata.size(); }
+
   /**
    * @brief Get the number of row groups per file
    *