diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index ef91c15d38bd..a49e300bd345 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -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 fetch_footer_to_host( +[[nodiscard]] std::unique_ptr 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> fetch_footers_to_host( + cudf::host_span const> datasources); + /** * @brief Fetches a host buffer of Parquet page index from the input data source * @@ -53,9 +68,26 @@ std::unique_ptr fetch_footer_to_host( * @param page_index_bytes Byte range of page index * @return Host buffer containing page index bytes */ -std::unique_ptr fetch_page_index_to_host( +[[nodiscard]] std::unique_ptr 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> fetch_page_indexes_to_host( + cudf::host_span const> datasources, + cudf::host_span page_index_bytes_per_source); + /** * @brief Fetches a list of byte ranges from a datasource into device buffers * diff --git a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp index 6f3316581365..baccb96ab537 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/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 #include #include #include @@ -21,8 +22,12 @@ #include #include +#include +#include #include +#include #include +#include /** * @file parquet_io_utils.cpp @@ -31,27 +36,149 @@ namespace cudf::io::parquet { -std::unique_ptr 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 +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; + + auto constexpr parallel_threshold = 32; + + std::vector results; + results.reserve(num_sources); + + if (num_sources < parallel_threshold) { + // Run sequentially to avoid task dispatch overhead + std::for_each(cuda::counting_iterator(0), + cuda::counting_iterator(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> tasks; + tasks.reserve(num_sources); + std::for_each(cuda::counting_iterator(0), + cuda::counting_iterator(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> fetch_footers_to_host( + cudf::host_span 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(header_buffer->data()); + auto ender_buffer = datasource.host_read(len - ender_len, ender_len); + auto const ender = reinterpret_cast(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(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> fetch_page_indexes_to_host( + cudf::host_span const> datasources, + cudf::host_span 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 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, 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> fetch_footers_to_host( + cudf::host_span const> datasources) +{ + CUDF_FUNC_RANGE(); + return detail::fetch_footers_to_host(datasources); } std::unique_ptr 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, 1> datasources{std::ref(datasource)}; + std::array 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> fetch_page_indexes_to_host( + cudf::host_span const> datasources, + cudf::host_span page_index_bytes_per_source) +{ + CUDF_FUNC_RANGE(); + return detail::fetch_page_indexes_to_host(datasources, page_index_bytes_per_source); } std::tuple, diff --git a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp index b0cb3b3160e0..10960c185ef9 100644 --- a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp @@ -19,8 +19,12 @@ #include +#include #include #include +#include +#include +#include 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()); + + std::vector> datasources(num_sources); + std::vector> datasource_refs{}; + std::transform(datasources.begin(), + datasources.end(), + std::back_inserter(datasource_refs), + [&](auto& datasource) { + datasource = cudf::io::datasource::create(cudf::host_span( + reinterpret_cast(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::host_span{static_cast(footer_buffers.front()->data()), + footer_buffers.front()->size()}, + cudf::io::parquet_reader_options::builder().build()); + std::vector 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);