From 41d27dfb92d2c2cc75f04689bc599f5f6cfb0b1b Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 15 May 2026 20:47:29 +0000 Subject: [PATCH 01/18] Fix the parquet io utils --- .../io/parquet/io_utils/parquet_io_utils.cpp | 79 ++++++++++++------- 1 file changed, 50 insertions(+), 29 deletions(-) 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 9b6953b4bd1d..f70fd710a1c6 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,6 +22,7 @@ #include #include +#include /** * @file parquet_io_utils.cpp @@ -86,9 +88,13 @@ fetch_byte_ranges_to_device_async( return acc + range.size(); }); + // Coalesce adjacent byte ranges into a single IO call. Pure local computation, no shared state. std::vector io_offsets; std::vector io_sizes; std::vector destinations; + io_offsets.reserve(byte_ranges.size()); + io_sizes.reserve(byte_ranges.size()); + destinations.reserve(byte_ranges.size()); for (size_t chunk = 0; chunk < byte_ranges.size();) { auto const io_offset = static_cast(byte_ranges[chunk].offset()); @@ -110,12 +116,24 @@ fetch_byte_ranges_to_device_async( CUDF_EXPECTS(io_offsets.size() == io_sizes.size() and io_sizes.size() == destinations.size(), "Unexpected number of IO offsets, sizes, or destinations"); + using host_read_buffer = std::unique_ptr; + std::vector> device_read_tasks{}; - std::vector> host_read_tasks{}; - device_read_tasks.reserve(byte_ranges.size()); - host_read_tasks.reserve(byte_ranges.size()); + std::vector> host_read_tasks{}; + device_read_tasks.reserve(io_offsets.size()); + host_read_tasks.reserve(io_offsets.size()); + + // Vectors to store intermediate host read buffers and relevant pointers + std::vector copy_dsts{}; + std::vector copy_sizes{}; + copy_dsts.reserve(io_offsets.size()); + copy_sizes.reserve(io_offsets.size()); + + // Defer space reservation until we know if there are host reads + std::vector copy_srcs{}; + std::vector host_read_buffers{}; - // device_read_async is not guaranteed to follow stream-ordering (see datasource API docs). + // `device_read_async` is not guaranteed to follow stream-ordering (see datasource API docs) stream.synchronize(); { @@ -133,35 +151,38 @@ fetch_byte_ranges_to_device_async( device_read_tasks.emplace_back( datasource.device_read_async(io_offset, io_size, dest, stream)); } else { - // Read the column chunk data to the host buffer copy it to the device buffer - host_read_tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( - [&datasource, io_offset, io_size, dest, stream]() { - auto host_buffer = datasource.host_read(io_offset, io_size); - cudf::detail::cuda_memcpy_async( - cudf::device_span{dest, io_size}, - cudf::host_span{host_buffer->data(), io_size}, - stream); - return io_size; - })); + // Asynchronously read column chunk data to a host buffer + host_read_tasks.emplace_back(datasource.host_read_async(io_offset, io_size)); + copy_dsts.push_back(static_cast(dest)); + copy_sizes.push_back(io_size); } }); - } - auto sync_function = [](decltype(host_read_tasks) host_read_tasks, - decltype(device_read_tasks) device_read_tasks) { - for (auto& task : host_read_tasks) { - task.get(); - } - for (auto& task : device_read_tasks) { - task.get(); + // If there are host reads, schedule a batched memcpy to device + if (not host_read_tasks.empty()) { + host_read_buffers.reserve(host_read_tasks.size()); + copy_srcs.reserve(host_read_tasks.size()); + for (auto& task : host_read_tasks) { + host_read_buffers.emplace_back(task.get()); + copy_srcs.push_back(host_read_buffers.back().get()->data()); + } + + CUDF_CUDA_TRY(cudf::detail::memcpy_batch_async( + copy_dsts.data(), copy_srcs.data(), copy_sizes.data(), copy_dsts.size(), stream)); } - }; - return {std::move(column_chunk_buffers), - std::move(column_chunk_data), - std::async(std::launch::deferred, - sync_function, - std::move(host_read_tasks), - std::move(device_read_tasks))}; + + // Synchronize stream to ensure `memcpy_batch_async` completes before vectors are destroyed + stream.synchronize(); + + auto sync_function = [](decltype(device_read_tasks) device_read_tasks) { + for (auto& task : device_read_tasks) { + task.get(); + } + }; + return {std::move(column_chunk_buffers), + std::move(column_chunk_data), + std::async(std::launch::deferred, sync_function, std::move(device_read_tasks))}; + } } } // namespace cudf::io::parquet From dbbdea220dd17d1556b5109c9a1001d9ea8a75f7 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 15 May 2026 20:52:14 +0000 Subject: [PATCH 02/18] Remove unnecessary comments --- cpp/src/io/parquet/io_utils/parquet_io_utils.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 f70fd710a1c6..d67656abbc38 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -88,7 +88,6 @@ fetch_byte_ranges_to_device_async( return acc + range.size(); }); - // Coalesce adjacent byte ranges into a single IO call. Pure local computation, no shared state. std::vector io_offsets; std::vector io_sizes; std::vector destinations; @@ -118,6 +117,7 @@ fetch_byte_ranges_to_device_async( using host_read_buffer = std::unique_ptr; + // Vectors to hold futures from datasource std::vector> device_read_tasks{}; std::vector> host_read_tasks{}; device_read_tasks.reserve(io_offsets.size()); From 5a95c7eaf2ec1f8f788312cebeba0643b98e62f8 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 15 May 2026 22:38:27 +0000 Subject: [PATCH 03/18] Remove unnecessary sync --- cpp/src/io/parquet/io_utils/parquet_io_utils.cpp | 4 ---- 1 file changed, 4 deletions(-) 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 d67656abbc38..21a73b39d5e7 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -166,14 +166,10 @@ fetch_byte_ranges_to_device_async( host_read_buffers.emplace_back(task.get()); copy_srcs.push_back(host_read_buffers.back().get()->data()); } - CUDF_CUDA_TRY(cudf::detail::memcpy_batch_async( copy_dsts.data(), copy_srcs.data(), copy_sizes.data(), copy_dsts.size(), stream)); } - // Synchronize stream to ensure `memcpy_batch_async` completes before vectors are destroyed - stream.synchronize(); - auto sync_function = [](decltype(device_read_tasks) device_read_tasks) { for (auto& task : device_read_tasks) { task.get(); From 410d495d5ceb7019cfd954dafd901a43b6a5403a Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 15 May 2026 23:06:11 +0000 Subject: [PATCH 04/18] Minor updates --- cpp/src/io/parquet/io_utils/parquet_io_utils.cpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) 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 21a73b39d5e7..b83cff7c9066 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -129,10 +129,6 @@ fetch_byte_ranges_to_device_async( copy_dsts.reserve(io_offsets.size()); copy_sizes.reserve(io_offsets.size()); - // Defer space reservation until we know if there are host reads - std::vector copy_srcs{}; - std::vector host_read_buffers{}; - // `device_read_async` is not guaranteed to follow stream-ordering (see datasource API docs) stream.synchronize(); @@ -160,8 +156,11 @@ fetch_byte_ranges_to_device_async( // If there are host reads, schedule a batched memcpy to device if (not host_read_tasks.empty()) { + std::vector host_read_buffers{}; + std::vector copy_srcs{}; host_read_buffers.reserve(host_read_tasks.size()); copy_srcs.reserve(host_read_tasks.size()); + for (auto& task : host_read_tasks) { host_read_buffers.emplace_back(task.get()); copy_srcs.push_back(host_read_buffers.back().get()->data()); From 96e177849712fd6285baa9914a57da9e4e54fb09 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 15 May 2026 23:13:49 +0000 Subject: [PATCH 05/18] Minor update --- cpp/src/io/parquet/io_utils/parquet_io_utils.cpp | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) 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 b83cff7c9066..547fb5c26faa 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -129,6 +129,9 @@ fetch_byte_ranges_to_device_async( copy_dsts.reserve(io_offsets.size()); copy_sizes.reserve(io_offsets.size()); + // Vector to store intermediate host buffers + std::vector host_buffers{}; + // `device_read_async` is not guaranteed to follow stream-ordering (see datasource API docs) stream.synchronize(); @@ -156,19 +159,21 @@ fetch_byte_ranges_to_device_async( // If there are host reads, schedule a batched memcpy to device if (not host_read_tasks.empty()) { - std::vector host_read_buffers{}; std::vector copy_srcs{}; - host_read_buffers.reserve(host_read_tasks.size()); copy_srcs.reserve(host_read_tasks.size()); + host_buffers.reserve(host_read_tasks.size()); for (auto& task : host_read_tasks) { - host_read_buffers.emplace_back(task.get()); - copy_srcs.push_back(host_read_buffers.back().get()->data()); + host_buffers.emplace_back(task.get()); + copy_srcs.push_back(host_buffers.back().get()->data()); } CUDF_CUDA_TRY(cudf::detail::memcpy_batch_async( copy_dsts.data(), copy_srcs.data(), copy_sizes.data(), copy_dsts.size(), stream)); } + // Synchronize the stream so that the host buffers can be safely discarded + stream.synchronize(); + auto sync_function = [](decltype(device_read_tasks) device_read_tasks) { for (auto& task : device_read_tasks) { task.get(); From 7ae60be0436252c842327deed89b228c8fe346b6 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Mon, 18 May 2026 17:21:28 +0000 Subject: [PATCH 06/18] Minor improvement --- cpp/src/io/parquet/io_utils/parquet_io_utils.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 547fb5c26faa..2a3fad68d7c9 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -171,8 +171,9 @@ fetch_byte_ranges_to_device_async( copy_dsts.data(), copy_srcs.data(), copy_sizes.data(), copy_dsts.size(), stream)); } - // Synchronize the stream so that the host buffers can be safely discarded - stream.synchronize(); + // Synchronize the stream if `memcpy_batch_async` was scheduled to safely discard the host + // buffers + if (not host_buffers.empty()) { stream.synchronize(); } auto sync_function = [](decltype(device_read_tasks) device_read_tasks) { for (auto& task : device_read_tasks) { From f552803c5b60a54a607b5fa09137a125b0170cc1 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Mon, 18 May 2026 17:56:25 +0000 Subject: [PATCH 07/18] Optimize parquet io utils --- .../io/parquet/io_utils/parquet_io_utils.cpp | 77 +++++++++++-------- 1 file changed, 44 insertions(+), 33 deletions(-) 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 2a3fad68d7c9..eca20896bc59 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -123,22 +123,47 @@ fetch_byte_ranges_to_device_async( device_read_tasks.reserve(io_offsets.size()); host_read_tasks.reserve(io_offsets.size()); - // Vectors to store intermediate host read buffers and relevant pointers + // Vectors to store intermediate host buffers and relevant pointers + std::vector host_buffers{}; + std::vector copy_srcs{}; std::vector copy_dsts{}; std::vector copy_sizes{}; copy_dsts.reserve(io_offsets.size()); copy_sizes.reserve(io_offsets.size()); - // Vector to store intermediate host buffers - std::vector host_buffers{}; + auto iter = cuda::make_zip_iterator(io_offsets.begin(), io_sizes.begin(), destinations.begin()); + + // Schedule host reads in parallel + std::for_each(iter, iter + io_offsets.size(), [&](auto const& tuple) { + auto const io_offset = cuda::std::get<0>(tuple); + auto const io_size = cuda::std::get<1>(tuple); + auto const dest = cuda::std::get<2>(tuple); + + if (not datasource.supports_device_read() or not datasource.is_device_read_preferred(io_size)) { + // Asynchronously read column chunk data to a host buffer + host_read_tasks.emplace_back(datasource.host_read_async(io_offset, io_size)); + copy_dsts.push_back(static_cast(dest)); + copy_sizes.push_back(io_size); + } + }); + + // Complete host reads + if (not host_read_tasks.empty()) { + copy_srcs.reserve(host_read_tasks.size()); + host_buffers.reserve(host_read_tasks.size()); + + for (auto& task : host_read_tasks) { + host_buffers.emplace_back(task.get()); + copy_srcs.push_back(host_buffers.back().get()->data()); + } + } // `device_read_async` is not guaranteed to follow stream-ordering (see datasource API docs) stream.synchronize(); + // Ensure all device reads for this thread are scheduled together { - auto iter = cuda::make_zip_iterator(io_offsets.begin(), io_sizes.begin(), destinations.begin()); - - std::lock_guard lock(mutex); + std::scoped_lock lock(mutex); std::for_each(iter, iter + io_offsets.size(), [&](auto const& tuple) { auto const io_offset = cuda::std::get<0>(tuple); @@ -149,41 +174,27 @@ fetch_byte_ranges_to_device_async( if (datasource.supports_device_read() and datasource.is_device_read_preferred(io_size)) { device_read_tasks.emplace_back( datasource.device_read_async(io_offset, io_size, dest, stream)); - } else { - // Asynchronously read column chunk data to a host buffer - host_read_tasks.emplace_back(datasource.host_read_async(io_offset, io_size)); - copy_dsts.push_back(static_cast(dest)); - copy_sizes.push_back(io_size); } }); - // If there are host reads, schedule a batched memcpy to device - if (not host_read_tasks.empty()) { - std::vector copy_srcs{}; - copy_srcs.reserve(host_read_tasks.size()); - host_buffers.reserve(host_read_tasks.size()); - - for (auto& task : host_read_tasks) { - host_buffers.emplace_back(task.get()); - copy_srcs.push_back(host_buffers.back().get()->data()); - } + // Schedule a batched memcpy from host buffersto device + if (not host_buffers.empty()) { CUDF_CUDA_TRY(cudf::detail::memcpy_batch_async( copy_dsts.data(), copy_srcs.data(), copy_sizes.data(), copy_dsts.size(), stream)); } + } - // Synchronize the stream if `memcpy_batch_async` was scheduled to safely discard the host - // buffers - if (not host_buffers.empty()) { stream.synchronize(); } + // Synchronize stream if `memcpy_batch_async` was called to safely discard the host buffers + if (not host_buffers.empty()) { stream.synchronize(); } - auto sync_function = [](decltype(device_read_tasks) device_read_tasks) { - for (auto& task : device_read_tasks) { - task.get(); - } - }; - return {std::move(column_chunk_buffers), - std::move(column_chunk_data), - std::async(std::launch::deferred, sync_function, std::move(device_read_tasks))}; - } + auto sync_function = [](decltype(device_read_tasks) device_read_tasks) { + for (auto& task : device_read_tasks) { + task.get(); + } + }; + return {std::move(column_chunk_buffers), + std::move(column_chunk_data), + std::async(std::launch::deferred, sync_function, std::move(device_read_tasks))}; } } // namespace cudf::io::parquet From da4c404ea4e840e7813b3a0700f5a2c21a2d99f1 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Mon, 18 May 2026 20:44:55 +0000 Subject: [PATCH 08/18] Add and use multisource parquet io utils --- cpp/include/cudf/io/parquet_io_utils.hpp | 23 ++ .../io/parquet/io_utils/parquet_io_utils.cpp | 200 +++++++++++++----- cpp/src/io/parquet/reader_impl_preprocess.cu | 3 +- .../parquet/reader_impl_preprocess_utils.cu | 120 ++++------- .../parquet/reader_impl_preprocess_utils.cuh | 4 +- 5 files changed, 222 insertions(+), 128 deletions(-) diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index ef91c15d38bd..e99eb9c322c7 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 @@ -77,6 +78,28 @@ fetch_byte_ranges_to_device_async(cudf::io::datasource& datasource, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); +/** + * @brief Fetches lists of byte ranges from multiple datasources into device buffers + * + * @ingroup io_utils + * + * @param datasources Input datasources + * @param byte_ranges_per_source Vector of byte ranges to fetch, one per datasource + * @param stream CUDA stream + * @param mr Device memory resource + * + * @return A tuple containing a vector of device buffers, a vector of vectors of device spans (one + * per byte range per datasource), and a future to wait on the read tasks + */ +std::tuple, + std::vector>>, + std::future> +fetch_byte_ranges_to_device_async( + cudf::host_span const> datasources, + cudf::host_span const> byte_ranges_per_source, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + /** @} */ // end of group } // namespace io::parquet } // namespace CUDF_EXPORT cudf 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 eca20896bc59..31a0a3b2be5c 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,6 +22,8 @@ #include #include +#include +#include #include #include @@ -56,64 +59,104 @@ std::unique_ptr fetch_page_index_to_host( return datasource.host_read(page_index_bytes.offset(), page_index_bytes.size()); } +namespace detail { + +using device_spans_per_source_type = std::vector>; + std::tuple, - std::vector>, + std::vector, std::future> fetch_byte_ranges_to_device_async( - cudf::io::datasource& datasource, - cudf::host_span byte_ranges, + cudf::host_span const> datasources, + cudf::host_span const> + byte_ranges_per_source, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { static std::mutex mutex; - // Allocate device spans for each column chunk - std::vector> column_chunk_data{}; - column_chunk_data.reserve(byte_ranges.size()); + auto const num_sources = datasources.size(); - auto total_size = std::accumulate( - byte_ranges.begin(), byte_ranges.end(), std::size_t{0}, [&](auto acc, auto const& range) { - return acc + range.size(); - }); + CUDF_EXPECTS(num_sources == byte_ranges_per_source.size(), + "Encountered mismatch in number of datasources and the number of byte range spans"); - // Allocate single device buffer for all column chunks - std::vector column_chunk_buffers{}; - // Buffer needs to be padded. Required by `gpuDecodePageData`. - column_chunk_buffers.emplace_back( - cudf::util::round_up_safe(total_size, cudf::io::detail::BUFFER_PADDING_MULTIPLE), stream, mr); - auto buffer_data = static_cast(column_chunk_buffers.back().data()); - std::ignore = std::accumulate( - byte_ranges.begin(), byte_ranges.end(), std::size_t{0}, [&](auto acc, auto const& range) { - column_chunk_data.emplace_back(buffer_data + acc, static_cast(range.size())); - return acc + range.size(); - }); + // Total number of byte ranges across all sources + auto const total_byte_ranges = + std::accumulate(byte_ranges_per_source.begin(), + byte_ranges_per_source.end(), + std::size_t{0}, + [](auto acc, auto const& ranges) { return acc + ranges.size(); }); + // IO descriptors + std::vector io_source_indices; std::vector io_offsets; std::vector io_sizes; std::vector destinations; - io_offsets.reserve(byte_ranges.size()); - io_sizes.reserve(byte_ranges.size()); - destinations.reserve(byte_ranges.size()); - - for (size_t chunk = 0; chunk < byte_ranges.size();) { - auto const io_offset = static_cast(byte_ranges[chunk].offset()); - auto io_size = static_cast(byte_ranges[chunk].size()); - size_t next_chunk = chunk + 1; - while (next_chunk < byte_ranges.size()) { - size_t const next_offset = byte_ranges[next_chunk].offset(); - if (next_offset != io_offset + io_size) { break; } - io_size += byte_ranges[next_chunk].size(); - next_chunk++; - } - if (io_size != 0) { - io_offsets.push_back(io_offset); - io_sizes.push_back(io_size); - destinations.push_back(const_cast(column_chunk_data[chunk].data())); - } - chunk = next_chunk; - } - CUDF_EXPECTS(io_offsets.size() == io_sizes.size() and io_sizes.size() == destinations.size(), - "Unexpected number of IO offsets, sizes, or destinations"); + io_source_indices.reserve(total_byte_ranges); + io_offsets.reserve(total_byte_ranges); + io_sizes.reserve(total_byte_ranges); + destinations.reserve(total_byte_ranges); + + // Allocate one device buffer per byte ranges of a datasource + std::vector column_chunk_buffers{}; + column_chunk_buffers.reserve(num_sources); + + // Column chunk device spans, one per byte range per datasource + std::vector column_chunk_data_per_source(num_sources); + + std::for_each( + cuda::counting_iterator(0), + cuda::counting_iterator(num_sources), + [&](auto const source_idx) { + auto const& byte_ranges = byte_ranges_per_source[source_idx]; + + // Total buffer size required for column chunks of this source + auto const buffer_size = std::accumulate( + byte_ranges.begin(), byte_ranges.end(), std::size_t{0}, [](auto acc, auto const& range) { + return acc + range.size(); + }); + + // Buffer needs to be padded. Required by `gpuDecodePageData`. + column_chunk_buffers.emplace_back( + cudf::util::round_up_safe(buffer_size, cudf::io::detail::BUFFER_PADDING_MULTIPLE), + stream, + mr); + + auto buffer_data = static_cast(column_chunk_buffers.back().data()); + + // Build device spans for each byte range in this source + auto& column_chunk_data = column_chunk_data_per_source[source_idx]; + column_chunk_data.reserve(byte_ranges.size()); + std::ignore = std::accumulate( + byte_ranges.begin(), byte_ranges.end(), std::size_t{0}, [&](auto acc, auto const& range) { + column_chunk_data.emplace_back(buffer_data + acc, static_cast(range.size())); + return acc + range.size(); + }); + + // Coalesce contiguous byte ranges within this source into single IO request + for (size_t chunk = 0; chunk < byte_ranges.size();) { + auto const io_offset = static_cast(byte_ranges[chunk].offset()); + auto io_size = static_cast(byte_ranges[chunk].size()); + size_t next_chunk = chunk + 1; + while (next_chunk < byte_ranges.size()) { + size_t const next_offset = byte_ranges[next_chunk].offset(); + if (next_offset != io_offset + io_size) { break; } + io_size += byte_ranges[next_chunk].size(); + next_chunk++; + } + if (io_size != 0) { + io_source_indices.push_back(source_idx); + io_offsets.push_back(io_offset); + io_sizes.push_back(io_size); + destinations.push_back(const_cast(column_chunk_data[chunk].data())); + } + chunk = next_chunk; + } + }); + + CUDF_EXPECTS(io_offsets.size() == io_sizes.size() and io_sizes.size() == destinations.size() and + io_source_indices.size() == io_offsets.size(), + "Unexpected number of IO source indices, offsets, sizes, or destinations"); using host_read_buffer = std::unique_ptr; @@ -131,14 +174,17 @@ fetch_byte_ranges_to_device_async( copy_dsts.reserve(io_offsets.size()); copy_sizes.reserve(io_offsets.size()); - auto iter = cuda::make_zip_iterator(io_offsets.begin(), io_sizes.begin(), destinations.begin()); + auto iter = cuda::make_zip_iterator( + io_source_indices.begin(), io_offsets.begin(), io_sizes.begin(), destinations.begin()); // Schedule host reads in parallel std::for_each(iter, iter + io_offsets.size(), [&](auto const& tuple) { - auto const io_offset = cuda::std::get<0>(tuple); - auto const io_size = cuda::std::get<1>(tuple); - auto const dest = cuda::std::get<2>(tuple); + auto const src_idx = cuda::std::get<0>(tuple); + auto const io_offset = cuda::std::get<1>(tuple); + auto const io_size = cuda::std::get<2>(tuple); + auto const dest = cuda::std::get<3>(tuple); + auto& datasource = datasources[src_idx].get(); if (not datasource.supports_device_read() or not datasource.is_device_read_preferred(io_size)) { // Asynchronously read column chunk data to a host buffer host_read_tasks.emplace_back(datasource.host_read_async(io_offset, io_size)); @@ -166,10 +212,12 @@ fetch_byte_ranges_to_device_async( std::scoped_lock lock(mutex); std::for_each(iter, iter + io_offsets.size(), [&](auto const& tuple) { - auto const io_offset = cuda::std::get<0>(tuple); - auto const io_size = cuda::std::get<1>(tuple); - auto const dest = cuda::std::get<2>(tuple); + auto const src_idx = cuda::std::get<0>(tuple); + auto const io_offset = cuda::std::get<1>(tuple); + auto const io_size = cuda::std::get<2>(tuple); + auto const dest = cuda::std::get<3>(tuple); + auto& datasource = datasources[src_idx].get(); // Directly read the column chunk data to the device buffer if supported if (datasource.supports_device_read() and datasource.is_device_read_preferred(io_size)) { device_read_tasks.emplace_back( @@ -193,8 +241,56 @@ fetch_byte_ranges_to_device_async( } }; return {std::move(column_chunk_buffers), - std::move(column_chunk_data), + std::move(column_chunk_data_per_source), std::async(std::launch::deferred, sync_function, std::move(device_read_tasks))}; } +} // namespace detail + +std::tuple, + std::vector>, + std::future> +fetch_byte_ranges_to_device_async( + cudf::io::datasource& datasource, + cudf::host_span byte_ranges, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + 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, 1> byte_ranges_per_source{ + byte_ranges}; + + auto [buffers, fetched_byte_ranges, fut] = detail::fetch_byte_ranges_to_device_async( + {datasources.data(), datasources.size()}, + {byte_ranges_per_source.data(), byte_ranges_per_source.size()}, + stream, + mr); + + return {std::move(buffers), std::move(fetched_byte_ranges.front()), std::move(fut)}; +} + +std::tuple, + std::vector>>, + std::future> +fetch_byte_ranges_to_device_async( + cudf::host_span const> datasources, + cudf::host_span const> byte_ranges_per_source, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + CUDF_FUNC_RANGE(); + + // Convert input vectors into host spans for detail API + std::vector> source_byte_range_spans; + source_byte_range_spans.reserve(byte_ranges_per_source.size()); + for (auto const& ranges : byte_ranges_per_source) { + source_byte_range_spans.emplace_back(ranges); + } + return detail::fetch_byte_ranges_to_device_async( + datasources, {source_byte_range_spans.data(), source_byte_range_spans.size()}, stream, mr); +} + } // namespace cudf::io::parquet diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index 8ebb8879d7e6..31f25f49d0e8 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -539,7 +539,8 @@ std::pair> reader_impl::read_column_chunks() chunks.size(), column_chunk_offsets, chunk_source_map, - _stream)}; + _stream, + cudf::get_current_device_resource_ref())}; } void reader_impl::read_compressed_data() diff --git a/cpp/src/io/parquet/reader_impl_preprocess_utils.cu b/cpp/src/io/parquet/reader_impl_preprocess_utils.cu index 34507c129b89..a36f62f7af79 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess_utils.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess_utils.cu @@ -4,15 +4,15 @@ */ #include "error.hpp" -#include "io/comp/common.hpp" #include "reader_impl_preprocess_utils.cuh" #include #include #include -#include #include #include +#include +#include #include @@ -29,7 +29,9 @@ #include #include +#if defined(PREPROCESS_DEBUG) #include +#endif // PREPROCESS_DEBUG #include namespace cudf::io::parquet::detail { @@ -172,81 +174,51 @@ void generate_depth_remappings( size_t end_chunk, std::vector const& column_chunk_offsets, std::vector const& chunk_source_map, - rmm::cuda_stream_view stream) + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { - // Calculate total size per source and offset for each chunk - std::vector source_total_size(sources.size(), 0); - std::vector chunk_buffer_offset(end_chunk - begin_chunk); - - for (size_t chunk = begin_chunk; chunk < end_chunk; ++chunk) { - auto const source_idx = chunk_source_map[chunk]; - chunk_buffer_offset[chunk - begin_chunk] = source_total_size[source_idx]; - source_total_size[source_idx] += chunks[chunk].compressed_size; + // Construct per source byte ranges in chunk iteration order + std::vector> source_byte_ranges(sources.size()); + std::for_each(cuda::counting_iterator(begin_chunk), + cuda::counting_iterator(end_chunk), + [&](auto const chunk) { + auto const source_idx = chunk_source_map[chunk]; + source_byte_ranges[source_idx].emplace_back( + static_cast(column_chunk_offsets[chunk]), + static_cast(chunks[chunk].compressed_size)); + }); + + // Wrap datasource references + std::vector> datasource_refs; + datasource_refs.reserve(sources.size()); + for (auto const& source : sources) { + datasource_refs.emplace_back(std::ref(*source)); } - // Allocate one buffer per source - std::transform( - source_total_size.begin(), source_total_size.end(), page_data.begin(), [&](size_t total_size) { - return rmm::device_buffer( - cudf::util::round_up_safe(total_size, cudf::io::detail::BUFFER_PADDING_MULTIPLE), stream); - }); - // device_read_async is not guaranteed to follow stream-ordering (see datasource API docs). - stream.synchronize(); - - // Issue reads, coalescing adjacent chunks - std::vector> read_tasks; - for (size_t chunk = begin_chunk; chunk < end_chunk;) { - auto const source_idx = chunk_source_map[chunk]; - auto const io_offset = column_chunk_offsets[chunk]; - size_t io_size = chunks[chunk].compressed_size; - size_t const first_chunk = chunk; - size_t next_chunk = chunk + 1; - - while (next_chunk < end_chunk) { - if (chunk_source_map[next_chunk] != source_idx) { break; } - auto const next_offset = column_chunk_offsets[next_chunk]; - if (next_offset != io_offset + io_size) { break; } - io_size += chunks[next_chunk].compressed_size; - next_chunk++; - } - - if (io_size != 0) { - auto& source = sources[source_idx]; - auto* dest = static_cast(page_data[source_idx].data()) + - chunk_buffer_offset[first_chunk - begin_chunk]; - - if (source->is_device_read_preferred(io_size)) { - auto fut = source->device_read_async(io_offset, io_size, dest, stream); - read_tasks.emplace_back(std::move(fut)); - } else { - read_tasks.emplace_back(std::async( - std::launch::deferred, [source = std::ref(*source), io_offset, io_size, dest, stream]() { - auto const read_buffer = source.get().host_read(io_offset, io_size); - cudf::detail::cuda_memcpy_async( - cudf::device_span{static_cast(dest), io_size}, - cudf::host_span{read_buffer->data(), io_size}, - stream); - return io_size; - })); - } - - // Set compressed_data pointers for all coalesced chunks - auto* ptr = static_cast(dest); - for (size_t c = first_chunk; c < next_chunk; ++c) { - chunks[c].compressed_data = ptr; - ptr += chunks[c].compressed_size; - } - } - - chunk = next_chunk; - } - - auto sync_fn = [](decltype(read_tasks) read_tasks) { - for (auto& task : read_tasks) { - task.get(); - } - }; - return std::async(std::launch::deferred, sync_fn, std::move(read_tasks)); + // Read byte ranges into buffers + auto [buffers, data_per_source, read_task] = cudf::io::parquet::fetch_byte_ranges_to_device_async( + {datasource_refs.data(), datasource_refs.size()}, + {source_byte_ranges.data(), source_byte_ranges.size()}, + stream, + mr); + + // Extract data pointers from returned spans + size_t range_idx = 0; + std::for_each(cuda::counting_iterator(begin_chunk), + cuda::counting_iterator(end_chunk), + [&](auto const chunk) { + auto const src_idx = chunk_source_map[chunk]; + chunks[chunk].compressed_data = data_per_source[src_idx][range_idx++].data(); + if (range_idx == data_per_source[src_idx].size()) { range_idx = 0; } + }); + + // Transfer device buffer ownership to the page data span + CUDF_EXPECTS(page_data.size() == buffers.size(), + "Parquet reader encountered a mismatch in size of page data and read buffers"); + + std::move(buffers.begin(), buffers.end(), page_data.begin()); + + return std::move(read_task); } [[nodiscard]] size_t count_page_headers(cudf::detail::hostdevice_span chunks, diff --git a/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh b/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh index e333e060c633..3a0f54cd5cd6 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh +++ b/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh @@ -56,6 +56,7 @@ void generate_depth_remappings( * @param column_chunk_offsets File offset for all chunks * @param chunk_source_map Association between each column chunk and its source * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the returned device buffers * * @return A future object for reading synchronization */ @@ -67,7 +68,8 @@ void generate_depth_remappings( size_t end_chunk, std::vector const& column_chunk_offsets, std::vector const& chunk_source_map, - rmm::cuda_stream_view stream); + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); /** * @brief Return the number of total pages from the given column chunks. From f42c3e8525b5e281ba6db7f70917940a1f90e05e Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Wed, 20 May 2026 21:49:38 +0000 Subject: [PATCH 09/18] Extend footer and page index APIs --- cpp/include/cudf/io/parquet_io_utils.hpp | 24 +++ .../io/parquet/io_utils/parquet_io_utils.cpp | 164 +++++++++++++++--- 2 files changed, 163 insertions(+), 25 deletions(-) diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index e99eb9c322c7..1d536bbb548d 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -45,6 +45,17 @@ using cudf::io::text::byte_range_info; 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 + */ +std::vector> fetch_footer_to_host( + cudf::host_span const> datasources); + /** * @brief Fetches a host buffer of Parquet page index from the input data source * @@ -57,6 +68,19 @@ std::unique_ptr fetch_footer_to_host( 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 + */ +std::vector> fetch_page_index_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 31a0a3b2be5c..e129534fa2ac 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -34,32 +34,103 @@ namespace cudf::io::parquet { -std::unique_ptr fetch_footer_to_host(cudf::io::datasource& datasource) +namespace detail { + +auto constexpr parallel_threshold = 16; + +std::vector> fetch_footer_to_host( + cudf::host_span const> datasources) { - constexpr auto header_len = sizeof(file_header_s); - constexpr auto ender_len = sizeof(file_ender_s); - size_t const len = datasource.size(); - - 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(len > header_len + ender_len, "Incorrect data source"); - CUDF_EXPECTS(header->magic == detail::parquet_magic, "Corrupted header"); - 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 datasource.host_read(len - ender->footer_len - ender_len, ender->footer_len); + // 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(); + + 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(len > header_len + ender_len, "Incorrect data source"); + 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); + }; + + std::vector> footer_buffers; + footer_buffers.reserve(datasources.size()); + auto const num_sources = datasources.size(); + + if (num_sources < parallel_threshold) { + // Read footers sequentially to avoid task dispatch overhead + std::transform(datasources.begin(), + datasources.end(), + std::back_inserter(footer_buffers), + [&](auto const& datasource_ref) { return fetch_footer(datasource_ref.get()); }); + } else { + // Read footers in parallel + std::vector>> tasks; + tasks.reserve(datasources.size()); + for (auto const& datasource_ref : datasources) { + tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( + [&datasource = datasource_ref.get(), &fetch_footer]() { + return fetch_footer(datasource); + })); + } + std::transform(tasks.begin(), tasks.end(), std::back_inserter(footer_buffers), [](auto& task) { + return task.get(); + }); + } + return footer_buffers; } -std::unique_ptr fetch_page_index_to_host( - cudf::io::datasource& datasource, cudf::io::text::byte_range_info const page_index_bytes) +std::vector> fetch_page_index_to_host( + cudf::host_span const> datasources, + cudf::host_span page_index_bytes_per_source) { - return datasource.host_read(page_index_bytes.offset(), page_index_bytes.size()); -} + CUDF_EXPECTS(datasources.size() == page_index_bytes_per_source.size(), + "Encountered mismatch in number of datasources and page index byte ranges"); -namespace detail { + // 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) { + return datasource.host_read(page_index_bytes.offset(), page_index_bytes.size()); + }; + + std::vector> page_index_buffers; + page_index_buffers.reserve(datasources.size()); + + auto const num_sources = datasources.size(); + auto iter = cuda::make_zip_iterator(datasources.begin(), page_index_bytes_per_source.begin()); + + if (num_sources < parallel_threshold) { + // Read page indexes sequentially to avoid task dispatch overhead + std::transform( + iter, iter + datasources.size(), std::back_inserter(page_index_buffers), [&](auto const& t) { + return fetch_page_index(cuda::std::get<0>(t).get(), cuda::std::get<1>(t)); + }); + } else { + // Read page indexes in parallel + std::vector>> tasks; + tasks.reserve(datasources.size()); + std::for_each(iter, iter + datasources.size(), [&](auto const& tuple) { + auto const& datasource = cuda::std::get<0>(tuple); + auto const& page_index_bytes = cuda::std::get<1>(tuple); + tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( + [&datasource = datasource.get(), page_index_bytes, &fetch_page_index]() { + return fetch_page_index(datasource, page_index_bytes); + })); + }); + std::transform( + tasks.begin(), tasks.end(), std::back_inserter(page_index_buffers), [](auto& task) { + return task.get(); + }); + } + return page_index_buffers; +} using device_spans_per_source_type = std::vector>; @@ -247,6 +318,46 @@ fetch_byte_ranges_to_device_async( } // 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_footer_to_host({datasources.data(), datasources.size()}); + return std::move(footer_buffers.front()); +} + +std::vector> fetch_footer_to_host( + cudf::host_span const> datasources) +{ + CUDF_FUNC_RANGE(); + return detail::fetch_footer_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) +{ + 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_index_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_index_to_host( + cudf::host_span const> datasources, + cudf::host_span page_index_bytes_per_source) +{ + CUDF_FUNC_RANGE(); + return detail::fetch_page_index_to_host(datasources, page_index_bytes_per_source); +} + std::tuple, std::vector>, std::future> @@ -284,13 +395,16 @@ fetch_byte_ranges_to_device_async( CUDF_FUNC_RANGE(); // Convert input vectors into host spans for detail API - std::vector> source_byte_range_spans; - source_byte_range_spans.reserve(byte_ranges_per_source.size()); + std::vector> byte_range_spans_per_source; + byte_range_spans_per_source.reserve(byte_ranges_per_source.size()); for (auto const& ranges : byte_ranges_per_source) { - source_byte_range_spans.emplace_back(ranges); + byte_range_spans_per_source.emplace_back(ranges); } return detail::fetch_byte_ranges_to_device_async( - datasources, {source_byte_range_spans.data(), source_byte_range_spans.size()}, stream, mr); + datasources, + {byte_range_spans_per_source.data(), byte_range_spans_per_source.size()}, + stream, + mr); } } // namespace cudf::io::parquet From 7544248cea8ae5bb7f22b398207d8c673032f60a Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 29 May 2026 03:06:24 +0000 Subject: [PATCH 10/18] Revert unrelated changes --- cpp/src/io/parquet/reader_impl_preprocess.cu | 3 +- .../parquet/reader_impl_preprocess_utils.cu | 120 +++++++++++------- .../parquet/reader_impl_preprocess_utils.cuh | 4 +- 3 files changed, 76 insertions(+), 51 deletions(-) diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index 31f25f49d0e8..8ebb8879d7e6 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -539,8 +539,7 @@ std::pair> reader_impl::read_column_chunks() chunks.size(), column_chunk_offsets, chunk_source_map, - _stream, - cudf::get_current_device_resource_ref())}; + _stream)}; } void reader_impl::read_compressed_data() diff --git a/cpp/src/io/parquet/reader_impl_preprocess_utils.cu b/cpp/src/io/parquet/reader_impl_preprocess_utils.cu index a36f62f7af79..34507c129b89 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess_utils.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess_utils.cu @@ -4,15 +4,15 @@ */ #include "error.hpp" +#include "io/comp/common.hpp" #include "reader_impl_preprocess_utils.cuh" #include #include #include +#include #include #include -#include -#include #include @@ -29,9 +29,7 @@ #include #include -#if defined(PREPROCESS_DEBUG) #include -#endif // PREPROCESS_DEBUG #include namespace cudf::io::parquet::detail { @@ -174,51 +172,81 @@ void generate_depth_remappings( size_t end_chunk, std::vector const& column_chunk_offsets, std::vector const& chunk_source_map, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) + rmm::cuda_stream_view stream) { - // Construct per source byte ranges in chunk iteration order - std::vector> source_byte_ranges(sources.size()); - std::for_each(cuda::counting_iterator(begin_chunk), - cuda::counting_iterator(end_chunk), - [&](auto const chunk) { - auto const source_idx = chunk_source_map[chunk]; - source_byte_ranges[source_idx].emplace_back( - static_cast(column_chunk_offsets[chunk]), - static_cast(chunks[chunk].compressed_size)); - }); - - // Wrap datasource references - std::vector> datasource_refs; - datasource_refs.reserve(sources.size()); - for (auto const& source : sources) { - datasource_refs.emplace_back(std::ref(*source)); + // Calculate total size per source and offset for each chunk + std::vector source_total_size(sources.size(), 0); + std::vector chunk_buffer_offset(end_chunk - begin_chunk); + + for (size_t chunk = begin_chunk; chunk < end_chunk; ++chunk) { + auto const source_idx = chunk_source_map[chunk]; + chunk_buffer_offset[chunk - begin_chunk] = source_total_size[source_idx]; + source_total_size[source_idx] += chunks[chunk].compressed_size; } - // Read byte ranges into buffers - auto [buffers, data_per_source, read_task] = cudf::io::parquet::fetch_byte_ranges_to_device_async( - {datasource_refs.data(), datasource_refs.size()}, - {source_byte_ranges.data(), source_byte_ranges.size()}, - stream, - mr); - - // Extract data pointers from returned spans - size_t range_idx = 0; - std::for_each(cuda::counting_iterator(begin_chunk), - cuda::counting_iterator(end_chunk), - [&](auto const chunk) { - auto const src_idx = chunk_source_map[chunk]; - chunks[chunk].compressed_data = data_per_source[src_idx][range_idx++].data(); - if (range_idx == data_per_source[src_idx].size()) { range_idx = 0; } - }); - - // Transfer device buffer ownership to the page data span - CUDF_EXPECTS(page_data.size() == buffers.size(), - "Parquet reader encountered a mismatch in size of page data and read buffers"); - - std::move(buffers.begin(), buffers.end(), page_data.begin()); - - return std::move(read_task); + // Allocate one buffer per source + std::transform( + source_total_size.begin(), source_total_size.end(), page_data.begin(), [&](size_t total_size) { + return rmm::device_buffer( + cudf::util::round_up_safe(total_size, cudf::io::detail::BUFFER_PADDING_MULTIPLE), stream); + }); + // device_read_async is not guaranteed to follow stream-ordering (see datasource API docs). + stream.synchronize(); + + // Issue reads, coalescing adjacent chunks + std::vector> read_tasks; + for (size_t chunk = begin_chunk; chunk < end_chunk;) { + auto const source_idx = chunk_source_map[chunk]; + auto const io_offset = column_chunk_offsets[chunk]; + size_t io_size = chunks[chunk].compressed_size; + size_t const first_chunk = chunk; + size_t next_chunk = chunk + 1; + + while (next_chunk < end_chunk) { + if (chunk_source_map[next_chunk] != source_idx) { break; } + auto const next_offset = column_chunk_offsets[next_chunk]; + if (next_offset != io_offset + io_size) { break; } + io_size += chunks[next_chunk].compressed_size; + next_chunk++; + } + + if (io_size != 0) { + auto& source = sources[source_idx]; + auto* dest = static_cast(page_data[source_idx].data()) + + chunk_buffer_offset[first_chunk - begin_chunk]; + + if (source->is_device_read_preferred(io_size)) { + auto fut = source->device_read_async(io_offset, io_size, dest, stream); + read_tasks.emplace_back(std::move(fut)); + } else { + read_tasks.emplace_back(std::async( + std::launch::deferred, [source = std::ref(*source), io_offset, io_size, dest, stream]() { + auto const read_buffer = source.get().host_read(io_offset, io_size); + cudf::detail::cuda_memcpy_async( + cudf::device_span{static_cast(dest), io_size}, + cudf::host_span{read_buffer->data(), io_size}, + stream); + return io_size; + })); + } + + // Set compressed_data pointers for all coalesced chunks + auto* ptr = static_cast(dest); + for (size_t c = first_chunk; c < next_chunk; ++c) { + chunks[c].compressed_data = ptr; + ptr += chunks[c].compressed_size; + } + } + + chunk = next_chunk; + } + + auto sync_fn = [](decltype(read_tasks) read_tasks) { + for (auto& task : read_tasks) { + task.get(); + } + }; + return std::async(std::launch::deferred, sync_fn, std::move(read_tasks)); } [[nodiscard]] size_t count_page_headers(cudf::detail::hostdevice_span chunks, diff --git a/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh b/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh index 3a0f54cd5cd6..e333e060c633 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh +++ b/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh @@ -56,7 +56,6 @@ void generate_depth_remappings( * @param column_chunk_offsets File offset for all chunks * @param chunk_source_map Association between each column chunk and its source * @param stream CUDA stream used for device memory operations and kernel launches - * @param mr Device memory resource used to allocate the returned device buffers * * @return A future object for reading synchronization */ @@ -68,8 +67,7 @@ void generate_depth_remappings( size_t end_chunk, std::vector const& column_chunk_offsets, std::vector const& chunk_source_map, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr); + rmm::cuda_stream_view stream); /** * @brief Return the number of total pages from the given column chunks. From 1ddce7be1c782097efeaba3fde43776c44a33448 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 29 May 2026 03:18:58 +0000 Subject: [PATCH 11/18] Pluralize API names --- cpp/include/cudf/io/parquet_io_utils.hpp | 26 +- .../io/parquet/io_utils/parquet_io_utils.cpp | 288 ++++++------------ 2 files changed, 100 insertions(+), 214 deletions(-) diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index 1d536bbb548d..4e5b85d6d091 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -53,7 +53,7 @@ std::unique_ptr fetch_footer_to_host( * @param datasources Input data sources * @return Vector of host buffers containing footer bytes, one per datasource */ -std::vector> fetch_footer_to_host( +std::vector> fetch_footers_to_host( cudf::host_span const> datasources); /** @@ -77,7 +77,7 @@ std::unique_ptr fetch_page_index_to_host( * @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 */ -std::vector> fetch_page_index_to_host( +std::vector> fetch_page_indexes_to_host( cudf::host_span const> datasources, cudf::host_span page_index_bytes_per_source); @@ -102,28 +102,6 @@ fetch_byte_ranges_to_device_async(cudf::io::datasource& datasource, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); -/** - * @brief Fetches lists of byte ranges from multiple datasources into device buffers - * - * @ingroup io_utils - * - * @param datasources Input datasources - * @param byte_ranges_per_source Vector of byte ranges to fetch, one per datasource - * @param stream CUDA stream - * @param mr Device memory resource - * - * @return A tuple containing a vector of device buffers, a vector of vectors of device spans (one - * per byte range per datasource), and a future to wait on the read tasks - */ -std::tuple, - std::vector>>, - std::future> -fetch_byte_ranges_to_device_async( - cudf::host_span const> datasources, - cudf::host_span const> byte_ranges_per_source, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr); - /** @} */ // end of group } // namespace io::parquet } // namespace CUDF_EXPORT cudf 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 e129534fa2ac..fda05e876b99 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -38,7 +38,7 @@ namespace detail { auto constexpr parallel_threshold = 16; -std::vector> fetch_footer_to_host( +std::vector> fetch_footers_to_host( cudf::host_span const> datasources) { // Helper to fetch footer from a datasource @@ -87,7 +87,7 @@ std::vector> fetch_footer_to_host( return footer_buffers; } -std::vector> fetch_page_index_to_host( +std::vector> fetch_page_indexes_to_host( cudf::host_span const> datasources, cudf::host_span page_index_bytes_per_source) { @@ -132,102 +132,106 @@ std::vector> fetch_page_index_to_h return page_index_buffers; } -using device_spans_per_source_type = std::vector>; +} // 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) +{ + 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, - std::vector, + std::vector>, std::future> fetch_byte_ranges_to_device_async( - cudf::host_span const> datasources, - cudf::host_span const> - byte_ranges_per_source, + cudf::io::datasource& datasource, + cudf::host_span byte_ranges, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { static std::mutex mutex; - auto const num_sources = datasources.size(); + // Allocate device spans for each column chunk + std::vector> column_chunk_data{}; + column_chunk_data.reserve(byte_ranges.size()); - CUDF_EXPECTS(num_sources == byte_ranges_per_source.size(), - "Encountered mismatch in number of datasources and the number of byte range spans"); + auto total_size = std::accumulate( + byte_ranges.begin(), byte_ranges.end(), std::size_t{0}, [&](auto acc, auto const& range) { + return acc + range.size(); + }); - // Total number of byte ranges across all sources - auto const total_byte_ranges = - std::accumulate(byte_ranges_per_source.begin(), - byte_ranges_per_source.end(), - std::size_t{0}, - [](auto acc, auto const& ranges) { return acc + ranges.size(); }); + // Allocate single device buffer for all column chunks + std::vector column_chunk_buffers{}; + // Buffer needs to be padded. Required by `gpuDecodePageData`. + column_chunk_buffers.emplace_back( + cudf::util::round_up_safe(total_size, cudf::io::detail::BUFFER_PADDING_MULTIPLE), stream, mr); + auto buffer_data = static_cast(column_chunk_buffers.back().data()); + std::ignore = std::accumulate( + byte_ranges.begin(), byte_ranges.end(), std::size_t{0}, [&](auto acc, auto const& range) { + column_chunk_data.emplace_back(buffer_data + acc, static_cast(range.size())); + return acc + range.size(); + }); - // IO descriptors - std::vector io_source_indices; std::vector io_offsets; std::vector io_sizes; std::vector destinations; - io_source_indices.reserve(total_byte_ranges); - io_offsets.reserve(total_byte_ranges); - io_sizes.reserve(total_byte_ranges); - destinations.reserve(total_byte_ranges); - - // Allocate one device buffer per byte ranges of a datasource - std::vector column_chunk_buffers{}; - column_chunk_buffers.reserve(num_sources); - - // Column chunk device spans, one per byte range per datasource - std::vector column_chunk_data_per_source(num_sources); - - std::for_each( - cuda::counting_iterator(0), - cuda::counting_iterator(num_sources), - [&](auto const source_idx) { - auto const& byte_ranges = byte_ranges_per_source[source_idx]; - - // Total buffer size required for column chunks of this source - auto const buffer_size = std::accumulate( - byte_ranges.begin(), byte_ranges.end(), std::size_t{0}, [](auto acc, auto const& range) { - return acc + range.size(); - }); - - // Buffer needs to be padded. Required by `gpuDecodePageData`. - column_chunk_buffers.emplace_back( - cudf::util::round_up_safe(buffer_size, cudf::io::detail::BUFFER_PADDING_MULTIPLE), - stream, - mr); - - auto buffer_data = static_cast(column_chunk_buffers.back().data()); - - // Build device spans for each byte range in this source - auto& column_chunk_data = column_chunk_data_per_source[source_idx]; - column_chunk_data.reserve(byte_ranges.size()); - std::ignore = std::accumulate( - byte_ranges.begin(), byte_ranges.end(), std::size_t{0}, [&](auto acc, auto const& range) { - column_chunk_data.emplace_back(buffer_data + acc, static_cast(range.size())); - return acc + range.size(); - }); - - // Coalesce contiguous byte ranges within this source into single IO request - for (size_t chunk = 0; chunk < byte_ranges.size();) { - auto const io_offset = static_cast(byte_ranges[chunk].offset()); - auto io_size = static_cast(byte_ranges[chunk].size()); - size_t next_chunk = chunk + 1; - while (next_chunk < byte_ranges.size()) { - size_t const next_offset = byte_ranges[next_chunk].offset(); - if (next_offset != io_offset + io_size) { break; } - io_size += byte_ranges[next_chunk].size(); - next_chunk++; - } - if (io_size != 0) { - io_source_indices.push_back(source_idx); - io_offsets.push_back(io_offset); - io_sizes.push_back(io_size); - destinations.push_back(const_cast(column_chunk_data[chunk].data())); - } - chunk = next_chunk; - } - }); - - CUDF_EXPECTS(io_offsets.size() == io_sizes.size() and io_sizes.size() == destinations.size() and - io_source_indices.size() == io_offsets.size(), - "Unexpected number of IO source indices, offsets, sizes, or destinations"); + io_offsets.reserve(byte_ranges.size()); + io_sizes.reserve(byte_ranges.size()); + destinations.reserve(byte_ranges.size()); + + for (size_t chunk = 0; chunk < byte_ranges.size();) { + auto const io_offset = static_cast(byte_ranges[chunk].offset()); + auto io_size = static_cast(byte_ranges[chunk].size()); + size_t next_chunk = chunk + 1; + while (next_chunk < byte_ranges.size()) { + size_t const next_offset = byte_ranges[next_chunk].offset(); + if (next_offset != io_offset + io_size) { break; } + io_size += byte_ranges[next_chunk].size(); + next_chunk++; + } + if (io_size != 0) { + io_offsets.push_back(io_offset); + io_sizes.push_back(io_size); + destinations.push_back(const_cast(column_chunk_data[chunk].data())); + } + chunk = next_chunk; + } + CUDF_EXPECTS(io_offsets.size() == io_sizes.size() and io_sizes.size() == destinations.size(), + "Unexpected number of IO offsets, sizes, or destinations"); using host_read_buffer = std::unique_ptr; @@ -245,18 +249,15 @@ fetch_byte_ranges_to_device_async( copy_dsts.reserve(io_offsets.size()); copy_sizes.reserve(io_offsets.size()); - auto iter = cuda::make_zip_iterator( - io_source_indices.begin(), io_offsets.begin(), io_sizes.begin(), destinations.begin()); + auto iter = cuda::make_zip_iterator(io_offsets.begin(), io_sizes.begin(), destinations.begin()); // Schedule host reads in parallel std::for_each(iter, iter + io_offsets.size(), [&](auto const& tuple) { - auto const src_idx = cuda::std::get<0>(tuple); - auto const io_offset = cuda::std::get<1>(tuple); - auto const io_size = cuda::std::get<2>(tuple); - auto const dest = cuda::std::get<3>(tuple); + auto const io_offset = cuda::std::get<0>(tuple); + auto const io_size = cuda::std::get<1>(tuple); + auto const dest = cuda::std::get<2>(tuple); - auto& datasource = datasources[src_idx].get(); - if (not datasource.supports_device_read() or not datasource.is_device_read_preferred(io_size)) { + if (not datasource.is_device_read_preferred(io_size)) { // Asynchronously read column chunk data to a host buffer host_read_tasks.emplace_back(datasource.host_read_async(io_offset, io_size)); copy_dsts.push_back(static_cast(dest)); @@ -283,20 +284,18 @@ fetch_byte_ranges_to_device_async( std::scoped_lock lock(mutex); std::for_each(iter, iter + io_offsets.size(), [&](auto const& tuple) { - auto const src_idx = cuda::std::get<0>(tuple); - auto const io_offset = cuda::std::get<1>(tuple); - auto const io_size = cuda::std::get<2>(tuple); - auto const dest = cuda::std::get<3>(tuple); + auto const io_offset = cuda::std::get<0>(tuple); + auto const io_size = cuda::std::get<1>(tuple); + auto const dest = cuda::std::get<2>(tuple); - auto& datasource = datasources[src_idx].get(); // Directly read the column chunk data to the device buffer if supported - if (datasource.supports_device_read() and datasource.is_device_read_preferred(io_size)) { + if (datasource.is_device_read_preferred(io_size)) { device_read_tasks.emplace_back( datasource.device_read_async(io_offset, io_size, dest, stream)); } }); - // Schedule a batched memcpy from host buffersto device + // Schedule a batched memcpy from host buffers to device if (not host_buffers.empty()) { CUDF_CUDA_TRY(cudf::detail::memcpy_batch_async( copy_dsts.data(), copy_srcs.data(), copy_sizes.data(), copy_dsts.size(), stream)); @@ -312,99 +311,8 @@ fetch_byte_ranges_to_device_async( } }; return {std::move(column_chunk_buffers), - std::move(column_chunk_data_per_source), + std::move(column_chunk_data), std::async(std::launch::deferred, sync_function, std::move(device_read_tasks))}; } -} // 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_footer_to_host({datasources.data(), datasources.size()}); - return std::move(footer_buffers.front()); -} - -std::vector> fetch_footer_to_host( - cudf::host_span const> datasources) -{ - CUDF_FUNC_RANGE(); - return detail::fetch_footer_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) -{ - 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_index_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_index_to_host( - cudf::host_span const> datasources, - cudf::host_span page_index_bytes_per_source) -{ - CUDF_FUNC_RANGE(); - return detail::fetch_page_index_to_host(datasources, page_index_bytes_per_source); -} - -std::tuple, - std::vector>, - std::future> -fetch_byte_ranges_to_device_async( - cudf::io::datasource& datasource, - cudf::host_span byte_ranges, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) -{ - 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, 1> byte_ranges_per_source{ - byte_ranges}; - - auto [buffers, fetched_byte_ranges, fut] = detail::fetch_byte_ranges_to_device_async( - {datasources.data(), datasources.size()}, - {byte_ranges_per_source.data(), byte_ranges_per_source.size()}, - stream, - mr); - - return {std::move(buffers), std::move(fetched_byte_ranges.front()), std::move(fut)}; -} - -std::tuple, - std::vector>>, - std::future> -fetch_byte_ranges_to_device_async( - cudf::host_span const> datasources, - cudf::host_span const> byte_ranges_per_source, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) -{ - CUDF_FUNC_RANGE(); - - // Convert input vectors into host spans for detail API - std::vector> byte_range_spans_per_source; - byte_range_spans_per_source.reserve(byte_ranges_per_source.size()); - for (auto const& ranges : byte_ranges_per_source) { - byte_range_spans_per_source.emplace_back(ranges); - } - return detail::fetch_byte_ranges_to_device_async( - datasources, - {byte_range_spans_per_source.data(), byte_range_spans_per_source.size()}, - stream, - mr); -} - } // namespace cudf::io::parquet From 94df4703970e271a0c5c78b9870b478d69d1e8d3 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 29 May 2026 03:36:12 +0000 Subject: [PATCH 12/18] Add gtest --- .../experimental/hybrid_scan_filters_test.cpp | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp index b0cb3b3160e0..0e4ff41cea84 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,53 @@ TEST_F(HybridScanFiltersTest, Metadata) EXPECT_EQ(reader->total_rows_in_row_groups(input_row_group_indices), 2 * rows_per_row_group); } +TEST_F(HybridScanFiltersTest, MultiSourceFooterAndPageIndex) +{ + srand(0xdede); + using T = uint32_t; + + auto const file_buffer = std::get<1>(create_parquet_with_stats()); + + auto constexpr num_sources = 3; + auto const datasource = cudf::io::datasource::create(cudf::host_span( + reinterpret_cast(file_buffer.data()), file_buffer.size())); + std::vector> datasources(num_sources, + std::ref(*datasource)); + + // Fetch all footers at once + auto const footer_buffers = + cudf::io::parquet::fetch_footers_to_host({datasources.data(), datasources.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( + {datasources.data(), datasources.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); + auto const single_page_index = + cudf::io::parquet::fetch_page_index_to_host(*datasource, 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); + }); +} + TEST_F(HybridScanFiltersTest, ExternalMetadata) { srand(0xcaffe); From 27e01c4c49da6a255aaf1b090d25227b91d6926a Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 29 May 2026 03:55:10 +0000 Subject: [PATCH 13/18] Use dispatcher. Reduce line count --- .../io/parquet/io_utils/parquet_io_utils.cpp | 110 +++++++++--------- 1 file changed, 54 insertions(+), 56 deletions(-) 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 fda05e876b99..83d5f76997c2 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -26,6 +26,7 @@ #include #include #include +#include /** * @file parquet_io_utils.cpp @@ -36,8 +37,51 @@ namespace cudf::io::parquet { namespace detail { -auto constexpr parallel_threshold = 16; +auto constexpr parallel_threshold = 32; +/** + * @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) +{ + using result_type = std::invoke_result_t; + + 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; +} + +/** + * @copydoc cudf::io::parquet::fetch_footers_to_host + */ std::vector> fetch_footers_to_host( cudf::host_span const> datasources) { @@ -60,33 +104,14 @@ std::vector> fetch_footers_to_host return datasource.host_read(len - ender->footer_len - ender_len, ender->footer_len); }; - std::vector> footer_buffers; - footer_buffers.reserve(datasources.size()); - auto const num_sources = datasources.size(); - - if (num_sources < parallel_threshold) { - // Read footers sequentially to avoid task dispatch overhead - std::transform(datasources.begin(), - datasources.end(), - std::back_inserter(footer_buffers), - [&](auto const& datasource_ref) { return fetch_footer(datasource_ref.get()); }); - } else { - // Read footers in parallel - std::vector>> tasks; - tasks.reserve(datasources.size()); - for (auto const& datasource_ref : datasources) { - tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( - [&datasource = datasource_ref.get(), &fetch_footer]() { - return fetch_footer(datasource); - })); - } - std::transform(tasks.begin(), tasks.end(), std::back_inserter(footer_buffers), [](auto& task) { - return task.get(); - }); - } - return footer_buffers; + return dispatch_fetch_tasks(datasources.size(), [&](std::size_t source_idx) { + return fetch_footer(datasources[source_idx].get()); + }); } +/** + * @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) @@ -100,36 +125,9 @@ std::vector> fetch_page_indexes_to return datasource.host_read(page_index_bytes.offset(), page_index_bytes.size()); }; - std::vector> page_index_buffers; - page_index_buffers.reserve(datasources.size()); - - auto const num_sources = datasources.size(); - auto iter = cuda::make_zip_iterator(datasources.begin(), page_index_bytes_per_source.begin()); - - if (num_sources < parallel_threshold) { - // Read page indexes sequentially to avoid task dispatch overhead - std::transform( - iter, iter + datasources.size(), std::back_inserter(page_index_buffers), [&](auto const& t) { - return fetch_page_index(cuda::std::get<0>(t).get(), cuda::std::get<1>(t)); - }); - } else { - // Read page indexes in parallel - std::vector>> tasks; - tasks.reserve(datasources.size()); - std::for_each(iter, iter + datasources.size(), [&](auto const& tuple) { - auto const& datasource = cuda::std::get<0>(tuple); - auto const& page_index_bytes = cuda::std::get<1>(tuple); - tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( - [&datasource = datasource.get(), page_index_bytes, &fetch_page_index]() { - return fetch_page_index(datasource, page_index_bytes); - })); - }); - std::transform( - tasks.begin(), tasks.end(), std::back_inserter(page_index_buffers), [](auto& task) { - return task.get(); - }); - } - return page_index_buffers; + 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 From 91b281b72e3e53031585df95263b3aa00cb2e5a3 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 29 May 2026 04:00:10 +0000 Subject: [PATCH 14/18] Address a comment --- cpp/src/io/parquet/io_utils/parquet_io_utils.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 83d5f76997c2..c94671426265 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -37,8 +37,6 @@ namespace cudf::io::parquet { namespace detail { -auto constexpr parallel_threshold = 32; - /** * @brief Dispatches the fetch task for each source index and collects the results * @@ -54,6 +52,8 @@ auto dispatch_fetch_tasks(std::size_t num_sources, Task fetch_task) { using result_type = std::invoke_result_t; + auto constexpr parallel_threshold = 32; + std::vector results; results.reserve(num_sources); @@ -90,12 +90,12 @@ std::vector> fetch_footers_to_host 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(len > header_len + ender_len, "Incorrect data source"); 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), From 4ebce246488e369f719f90f5dfcdf005d4ace4c2 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 29 May 2026 04:01:33 +0000 Subject: [PATCH 15/18] Add nodiscard --- cpp/include/cudf/io/parquet_io_utils.hpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index 4e5b85d6d091..0af3d434ab14 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -42,7 +42,7 @@ 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); /** @@ -53,7 +53,7 @@ std::unique_ptr fetch_footer_to_host( * @param datasources Input data sources * @return Vector of host buffers containing footer bytes, one per datasource */ -std::vector> fetch_footers_to_host( +[[nodiscard]] std::vector> fetch_footers_to_host( cudf::host_span const> datasources); /** @@ -65,7 +65,7 @@ std::vector> fetch_footers_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); /** @@ -77,7 +77,7 @@ std::unique_ptr fetch_page_index_to_host( * @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 */ -std::vector> fetch_page_indexes_to_host( +[[nodiscard]] std::vector> fetch_page_indexes_to_host( cudf::host_span const> datasources, cudf::host_span page_index_bytes_per_source); From f38056035ae482dd14e362648c36e3e6c2c97720 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 29 May 2026 23:47:13 +0000 Subject: [PATCH 16/18] Improve test coverage --- .../experimental/hybrid_scan_filters_test.cpp | 81 ++++++++++--------- 1 file changed, 43 insertions(+), 38 deletions(-) diff --git a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp index 0e4ff41cea84..54f91f264499 100644 --- a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp @@ -161,51 +161,56 @@ TEST_F(HybridScanFiltersTest, Metadata) EXPECT_EQ(reader->total_rows_in_row_groups(input_row_group_indices), 2 * rows_per_row_group); } -TEST_F(HybridScanFiltersTest, MultiSourceFooterAndPageIndex) +TEST_F(HybridScanFiltersTest, MultiSourceMetadata) { srand(0xdede); using T = uint32_t; - auto const file_buffer = std::get<1>(create_parquet_with_stats()); + auto const test_multisource_metadata = [&](auto num_sources) { + auto const file_buffer = std::get<1>(create_parquet_with_stats()); - auto constexpr num_sources = 3; - auto const datasource = cudf::io::datasource::create(cudf::host_span( - reinterpret_cast(file_buffer.data()), file_buffer.size())); - std::vector> datasources(num_sources, - std::ref(*datasource)); + auto const datasource = cudf::io::datasource::create(cudf::host_span( + reinterpret_cast(file_buffer.data()), file_buffer.size())); + std::vector> datasources(num_sources, + std::ref(*datasource)); - // Fetch all footers at once - auto const footer_buffers = - cudf::io::parquet::fetch_footers_to_host({datasources.data(), datasources.size()}); - ASSERT_EQ(footer_buffers.size(), num_sources); + // Fetch all footers at once + auto const footer_buffers = + cudf::io::parquet::fetch_footers_to_host({datasources.data(), datasources.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( - {datasources.data(), datasources.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); - auto const single_page_index = - cudf::io::parquet::fetch_page_index_to_host(*datasource, 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); - }); + // 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( + {datasources.data(), datasources.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); + auto const single_page_index = + cudf::io::parquet::fetch_page_index_to_host(*datasource, 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); + }); + }; + + test_multisource_metadata(3); + test_multisource_metadata(32); } TEST_F(HybridScanFiltersTest, ExternalMetadata) From 311de40fcfb6e1d1d6a0c3eaf87ac82ce6892b20 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Mon, 1 Jun 2026 17:18:55 +0000 Subject: [PATCH 17/18] Address review comments --- .../io/parquet/io_utils/parquet_io_utils.cpp | 9 ++++++ .../experimental/hybrid_scan_filters_test.cpp | 31 ++++++++++++------- 2 files changed, 29 insertions(+), 11 deletions(-) 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 c94671426265..baccb96ab537 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -25,6 +25,7 @@ #include #include #include +#include #include #include @@ -122,6 +123,14 @@ std::vector> fetch_page_indexes_to // 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()); }; diff --git a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp index 54f91f264499..10960c185ef9 100644 --- a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp @@ -166,17 +166,24 @@ 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()); - auto const datasource = cudf::io::datasource::create(cudf::host_span( - reinterpret_cast(file_buffer.data()), file_buffer.size())); - std::vector> datasources(num_sources, - std::ref(*datasource)); + 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({datasources.data(), datasources.size()}); + 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 @@ -187,14 +194,14 @@ TEST_F(HybridScanFiltersTest, MultiSourceMetadata) 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( - {datasources.data(), datasources.size()}, + {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); - auto const single_page_index = - cudf::io::parquet::fetch_page_index_to_host(*datasource, reader->page_index_byte_range()); + 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) { @@ -209,8 +216,10 @@ TEST_F(HybridScanFiltersTest, MultiSourceMetadata) }); }; - test_multisource_metadata(3); - test_multisource_metadata(32); + auto num_sources = 4; + test_multisource_metadata(num_sources); + num_sources = 32; + test_multisource_metadata(num_sources); } TEST_F(HybridScanFiltersTest, ExternalMetadata) From bf5c5be74728723e500c5eecc95fcf5bd06867c6 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Mon, 1 Jun 2026 23:31:21 +0000 Subject: [PATCH 18/18] Address comments --- cpp/include/cudf/io/parquet_io_utils.hpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index 0af3d434ab14..a49e300bd345 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -52,6 +52,9 @@ using cudf::io::text::byte_range_info; * * @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); @@ -76,6 +79,10 @@ using cudf::io::text::byte_range_info; * @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,