Parquet utilities to fetch footer and page index buffers from multiple sources

cpp/include/cudf/io/parquet_io_utils.hpp

@@ -12,6 +12,7 @@
 #include <rmm/device_buffer.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <functional>
 #include <future>
 #include <tuple>
 #include <vector>
@@ -41,9 +42,23 @@ using cudf::io::text::byte_range_info;
  * @param datasource Input data source
  * @return Host buffer containing footer bytes
  */
-std::unique_ptr<cudf::io::datasource::buffer> fetch_footer_to_host(
+[[nodiscard]] std::unique_ptr<cudf::io::datasource::buffer> fetch_footer_to_host(
   cudf::io::datasource& datasource);
 
+/**
+ * @brief Fetches host buffers of Parquet footer bytes from multiple input data sources
+ *
+ * @ingroup io_utils
+ *
+ * @param datasources Input data sources
+ * @return Vector of host buffers containing footer bytes, one per datasource
+ *
+ * @throw cudf::logic_error if any datasource contains a corrupted Parquet magic number, header or
+ * footer, or has an invalid footer length.
+ */
+[[nodiscard]] std::vector<std::unique_ptr<cudf::io::datasource::buffer>> fetch_footers_to_host(
+  cudf::host_span<std::reference_wrapper<cudf::io::datasource> const> datasources);
+
 /**
  * @brief Fetches a host buffer of Parquet page index from the input data source
  *
@@ -53,9 +68,26 @@ std::unique_ptr<cudf::io::datasource::buffer> fetch_footer_to_host(
  * @param page_index_bytes Byte range of page index
  * @return Host buffer containing page index bytes
  */
-std::unique_ptr<cudf::io::datasource::buffer> fetch_page_index_to_host(
+[[nodiscard]] std::unique_ptr<cudf::io::datasource::buffer> fetch_page_index_to_host(
   cudf::io::datasource& datasource, byte_range_info const page_index_bytes);
 
+/**
+ * @brief Fetches host buffers of Parquet page index bytes from multiple input data sources
+ *
+ * @ingroup io_utils
+ *
+ * @param datasources Input datasources
+ * @param page_index_bytes_per_source Byte ranges of page index, one per datasource
+ * @return Vector of host buffers containing page index bytes, one per datasource
+ *
+ * @throw cudf::logic_error if the number of datasources does not match the number of page index
+ * byte ranges
+ * @throw std::out_of_range if any page index byte range is out of range for its datasource
+ */
+[[nodiscard]] std::vector<std::unique_ptr<cudf::io::datasource::buffer>> fetch_page_indexes_to_host(
+  cudf::host_span<std::reference_wrapper<cudf::io::datasource> const> datasources,
+  cudf::host_span<byte_range_info const> page_index_bytes_per_source);
+
 /**
  * @brief Fetches a list of byte ranges from a datasource into device buffers
  *

cpp/src/io/parquet/io_utils/parquet_io_utils.cpp

@@ -6,6 +6,7 @@
 #include "io/comp/common.hpp"
 #include "io/parquet/parquet_common.hpp"
 
+#include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/utilities/cuda_memcpy.hpp>
 #include <cudf/detail/utilities/host_worker_pool.hpp>
 #include <cudf/detail/utilities/integer_utils.hpp>
@@ -21,8 +22,12 @@
 #include <cuda/iterator>
 #include <cuda/std/tuple>
 
+#include <functional>
+#include <mutex>
 #include <numeric>
+#include <stdexcept>
 #include <tuple>
+#include <type_traits>
 
 /**
  * @file parquet_io_utils.cpp
@@ -31,27 +36,149 @@
 
 namespace cudf::io::parquet {
 
-std::unique_ptr<cudf::io::datasource::buffer> fetch_footer_to_host(cudf::io::datasource& datasource)
+namespace detail {
+
+/**
+ * @brief Dispatches the fetch task for each source index and collects the results
+ *
+ * Dispatches sequentially or using host worker pool depending on the number of sources.
+ *
+ * @tparam Task Callable invocable as `fetch_task(std::size_t source_idx)`
+ * @param num_sources Number of sources to process
+ * @param fetch_task Task to run for each source index
+ * @return Vector of results, one per source, in source order
+ */
+template <typename Task>
+auto dispatch_fetch_tasks(std::size_t num_sources, Task fetch_task)
 {
-  constexpr auto header_len = sizeof(file_header_s);
-  constexpr auto ender_len  = sizeof(file_ender_s);
-  size_t const len          = datasource.size();
+  using result_type = std::invoke_result_t<Task, std::size_t>;
+
+  auto constexpr parallel_threshold = 32;
+
+  std::vector<result_type> results;
+  results.reserve(num_sources);
+
+  if (num_sources < parallel_threshold) {
+    // Run sequentially to avoid task dispatch overhead
+    std::for_each(cuda::counting_iterator<std::size_t>(0),
+                  cuda::counting_iterator<std::size_t>(num_sources),
+                  [&](std::size_t source_idx) { results.emplace_back(fetch_task(source_idx)); });
+  } else {
+    // Dispatch the tasks to the host worker pool
+    std::vector<std::future<result_type>> tasks;
+    tasks.reserve(num_sources);
+    std::for_each(cuda::counting_iterator<std::size_t>(0),
+                  cuda::counting_iterator<std::size_t>(num_sources),
+                  [&](std::size_t source_idx) {
+                    tasks.emplace_back(cudf::detail::host_worker_pool().submit_task(
+                      [&fetch_task, source_idx]() { return fetch_task(source_idx); }));
+                  });
+    std::transform(tasks.begin(), tasks.end(), std::back_inserter(results), [](auto& task) {
+      return task.get();
+    });
+  }
+  return results;
+}
 
-  CUDF_EXPECTS(len > header_len + ender_len, "Incorrect data source");
+/**
+ * @copydoc cudf::io::parquet::fetch_footers_to_host
+ */
+std::vector<std::unique_ptr<cudf::io::datasource::buffer>> fetch_footers_to_host(
+  cudf::host_span<std::reference_wrapper<cudf::io::datasource> const> datasources)
+{
+  // Helper to fetch footer from a datasource
+  auto const fetch_footer = [](cudf::io::datasource& datasource) {
+    constexpr auto header_len = sizeof(file_header_s);
+    constexpr auto ender_len  = sizeof(file_ender_s);
+    size_t const len          = datasource.size();
+    CUDF_EXPECTS(len > header_len + ender_len, "Incorrect data source");
+
+    auto header_buffer = datasource.host_read(0, header_len);
+    auto const header  = reinterpret_cast<file_header_s const*>(header_buffer->data());
+    auto ender_buffer  = datasource.host_read(len - ender_len, ender_len);
+    auto const ender   = reinterpret_cast<file_ender_s const*>(ender_buffer->data());
+    CUDF_EXPECTS(header->magic == parquet_magic, "Corrupted header");
+    CUDF_EXPECTS(ender->magic == parquet_magic, "Corrupted footer");
+    CUDF_EXPECTS(ender->footer_len != 0 && ender->footer_len <= (len - header_len - ender_len),
+                 "Incorrect footer length");
+
+    return datasource.host_read(len - ender->footer_len - ender_len, ender->footer_len);
+  };
 
-  auto ender_buffer = datasource.host_read(len - ender_len, ender_len);
-  auto const ender  = reinterpret_cast<file_ender_s const*>(ender_buffer->data());
-  CUDF_EXPECTS(ender->magic == detail::parquet_magic, "Corrupted footer");
-  CUDF_EXPECTS(ender->footer_len != 0 && ender->footer_len <= (len - header_len - ender_len),
-               "Incorrect footer length");
+  return dispatch_fetch_tasks(datasources.size(), [&](std::size_t source_idx) {
+    return fetch_footer(datasources[source_idx].get());
+  });
+}
 
-  return datasource.host_read(len - ender->footer_len - ender_len, ender->footer_len);
+/**
+ * @copydoc cudf::io::parquet::fetch_page_indexes_to_host
+ */
+std::vector<std::unique_ptr<cudf::io::datasource::buffer>> fetch_page_indexes_to_host(
+  cudf::host_span<std::reference_wrapper<cudf::io::datasource> const> datasources,
+  cudf::host_span<cudf::io::text::byte_range_info const> page_index_bytes_per_source)
+{
+  CUDF_EXPECTS(datasources.size() == page_index_bytes_per_source.size(),
+               "Encountered mismatch in number of datasources and page index byte ranges");
+
+  // Helper to fetch page index bytes from a datasource
+  auto const fetch_page_index = [](cudf::io::datasource& datasource,
+                                   cudf::io::text::byte_range_info const& page_index_bytes) {
+    CUDF_EXPECTS(
+      page_index_bytes.offset() >= 0 and
+        std::cmp_less_equal(page_index_bytes.offset() + page_index_bytes.size(), datasource.size()),
+      std::format("Invalid page index byte range: offset={}, size={}, datasource_size={}",
+                  page_index_bytes.offset(),
+                  page_index_bytes.size(),
+                  datasource.size()),
+      std::out_of_range);
+    return datasource.host_read(page_index_bytes.offset(), page_index_bytes.size());
+  };
+
+  return dispatch_fetch_tasks(datasources.size(), [&](std::size_t source_idx) {
+    return fetch_page_index(datasources[source_idx].get(), page_index_bytes_per_source[source_idx]);
+  });
+}
+
+}  // namespace detail
+
+std::unique_ptr<cudf::io::datasource::buffer> fetch_footer_to_host(cudf::io::datasource& datasource)
+{
+  CUDF_FUNC_RANGE();
+
+  // Wrap the input into an array and delegate to the detail multi-source API
+  std::array<std::reference_wrapper<cudf::io::datasource>, 1> datasources{std::ref(datasource)};
+  auto footer_buffers = detail::fetch_footers_to_host({datasources.data(), datasources.size()});
+  return std::move(footer_buffers.front());
+}
+
+std::vector<std::unique_ptr<cudf::io::datasource::buffer>> fetch_footers_to_host(
+  cudf::host_span<std::reference_wrapper<cudf::io::datasource> const> datasources)
+{
+  CUDF_FUNC_RANGE();
+  return detail::fetch_footers_to_host(datasources);
 }
 
 std::unique_ptr<cudf::io::datasource::buffer> fetch_page_index_to_host(
   cudf::io::datasource& datasource, cudf::io::text::byte_range_info const page_index_bytes)
 {
-  return datasource.host_read(page_index_bytes.offset(), page_index_bytes.size());
+  CUDF_FUNC_RANGE();
+
+  // Wrap the inputs into arrays and delegate to the detail multi-source API
+  std::array<std::reference_wrapper<cudf::io::datasource>, 1> datasources{std::ref(datasource)};
+  std::array<cudf::io::text::byte_range_info, 1> page_index_bytes_per_source{page_index_bytes};
+
+  auto page_index_buffers = detail::fetch_page_indexes_to_host(
+    {datasources.data(), datasources.size()},
+    {page_index_bytes_per_source.data(), page_index_bytes_per_source.size()});
+  return std::move(page_index_buffers.front());
+}
+
+std::vector<std::unique_ptr<cudf::io::datasource::buffer>> fetch_page_indexes_to_host(
+  cudf::host_span<std::reference_wrapper<cudf::io::datasource> const> datasources,
+  cudf::host_span<cudf::io::text::byte_range_info const> page_index_bytes_per_source)
+{
+  CUDF_FUNC_RANGE();
+  return detail::fetch_page_indexes_to_host(datasources, page_index_bytes_per_source);
 }
 
 std::tuple<std::vector<rmm::device_buffer>,

cpp/tests/io/experimental/hybrid_scan_filters_test.cpp

@@ -19,8 +19,12 @@
 
 #include <src/io/parquet/parquet_gpu.hpp>
 
+#include <cstring>
 #include <filesystem>
 #include <fstream>
+#include <functional>
+#include <memory>
+#include <vector>
 
 namespace {
 
@@ -157,6 +161,67 @@ TEST_F(HybridScanFiltersTest, Metadata)
   EXPECT_EQ(reader->total_rows_in_row_groups(input_row_group_indices), 2 * rows_per_row_group);
 }
 
+TEST_F(HybridScanFiltersTest, MultiSourceMetadata)
+{
+  srand(0xdede);
+  using T = uint32_t;
+
+  // Helper to test multi-source metadata
+  auto const test_multisource_metadata = [&](auto num_sources) {
+    auto const file_buffer = std::get<1>(create_parquet_with_stats<T, 1>());
+
+    std::vector<std::unique_ptr<cudf::io::datasource>> datasources(num_sources);
+    std::vector<std::reference_wrapper<cudf::io::datasource>> datasource_refs{};
+    std::transform(datasources.begin(),
+                   datasources.end(),
+                   std::back_inserter(datasource_refs),
+                   [&](auto& datasource) {
+                     datasource = cudf::io::datasource::create(cudf::host_span<std::byte const>(
+                       reinterpret_cast<std::byte const*>(file_buffer.data()), file_buffer.size()));
+                     return std::ref(*datasource);
+                   });
+
+    // Fetch all footers at once
+    auto const footer_buffers =
+      cudf::io::parquet::fetch_footers_to_host({datasource_refs.data(), datasource_refs.size()});
+    ASSERT_EQ(footer_buffers.size(), num_sources);
+
+    // Fetch all page indexes at once
+    auto const reader = std::make_unique<cudf::io::parquet::experimental::hybrid_scan_reader>(
+      cudf::host_span<uint8_t const>{static_cast<uint8_t const*>(footer_buffers.front()->data()),
+                                     footer_buffers.front()->size()},
+      cudf::io::parquet_reader_options::builder().build());
+    std::vector<cudf::io::parquet::byte_range_info> page_index_byte_ranges(
+      num_sources, reader->page_index_byte_range());
+    auto const page_index_buffers = cudf::io::parquet::fetch_page_indexes_to_host(
+      {datasource_refs.data(), datasource_refs.size()},
+      {page_index_byte_ranges.data(), page_index_byte_ranges.size()});
+    ASSERT_EQ(page_index_buffers.size(), num_sources);
+
+    // Footer and page index from multi-source and single-source APIs should match
+    auto const single_footer     = cudf::io::parquet::fetch_footer_to_host(datasource_refs.front());
+    auto const single_page_index = cudf::io::parquet::fetch_page_index_to_host(
+      datasource_refs.front(), reader->page_index_byte_range());
+
+    auto const iter = cuda::make_zip_iterator(footer_buffers.begin(), page_index_buffers.begin());
+    std::for_each(iter, iter + num_sources, [&](auto const& pair) {
+      auto const& [footer_buffer, page_index_buffer] = pair;
+      ASSERT_EQ(footer_buffer->size(), single_footer->size());
+      EXPECT_EQ(std::memcmp(footer_buffer->data(), single_footer->data(), single_footer->size()),
+                0);
+      ASSERT_EQ(page_index_buffer->size(), single_page_index->size());
+      EXPECT_EQ(std::memcmp(
+                  page_index_buffer->data(), single_page_index->data(), single_page_index->size()),
+                0);
+    });
+  };
+
+  auto num_sources = 4;
+  test_multisource_metadata(num_sources);
+  num_sources = 32;
+  test_multisource_metadata(num_sources);
+}
+
 TEST_F(HybridScanFiltersTest, ExternalMetadata)
 {
   srand(0xcaffe);