From 41d27dfb92d2c2cc75f04689bc599f5f6cfb0b1b Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 15 May 2026 20:47:29 +0000 Subject: [PATCH 01/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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/13] 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 683acc91c8ba1459831e7d14fdc18a8cdd8091d0 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Wed, 20 May 2026 21:50:23 +0000 Subject: [PATCH 09/13] Minor variable name improvement --- cpp/src/io/parquet/io_utils/parquet_io_utils.cpp | 11 +++++++---- 1 file changed, 7 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 31a0a3b2be5c..0ee09efae577 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -284,13 +284,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 a7928f1109ebd7620219ca6500623b433e75fd01 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 22 May 2026 22:55:59 +0000 Subject: [PATCH 10/13] Update --- .../io/parquet/io_utils/parquet_io_utils.cpp | 46 +++++++++++-------- 1 file changed, 28 insertions(+), 18 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 2b0218bc7724..98cd06dae6f0 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -73,7 +73,8 @@ fetch_byte_ranges_to_device_async( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - static std::mutex mutex; + static std::mutex host_read_mutex; + static std::mutex device_read_mutex; auto const num_sources = datasources.size(); @@ -177,21 +178,29 @@ fetch_byte_ranges_to_device_async( 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 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)); - copy_dsts.push_back(static_cast(dest)); - copy_sizes.push_back(io_size); - } - }); + // Schedule host reads holding the `host_read_mutex` so that all reads for a caller thread + // are scheduled without interleaving with reads from other threads yielding better pipelining + { + std::scoped_lock lock(host_read_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& datasource = datasources[src_idx].get(); + if (not datasource.is_device_read_preferred(io_size)) { + // Asynchronously read column chunk data to a host buffer + host_read_tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( + [&datasource, io_offset, io_size]() -> host_read_buffer { + return datasource.host_read(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()) { @@ -207,9 +216,10 @@ fetch_byte_ranges_to_device_async( // `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 + // Schedule device reads holding the `device_read_mutex` so that all reads for a caller thread + // are scheduled without interleaving with reads from other threads yielding better pipelining { - std::scoped_lock lock(mutex); + std::scoped_lock lock(device_read_mutex); std::for_each(iter, iter + io_offsets.size(), [&](auto const& tuple) { auto const src_idx = cuda::std::get<0>(tuple); From ca2e3f863709e8bdbe51d2e1ba9bdd3f160cd986 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Tue, 26 May 2026 15:48:46 -0700 Subject: [PATCH 11/13] Update cpp/src/io/parquet/io_utils/parquet_io_utils.cpp Co-authored-by: Bradley Dice --- 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 98cd06dae6f0..42843bc96101 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -235,7 +235,7 @@ fetch_byte_ranges_to_device_async( } }); - // 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)); From 867341ef60de51819637a1310e114ba6d053941d Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Mon, 1 Jun 2026 16:20:22 -0700 Subject: [PATCH 12/13] Update cpp/src/io/parquet/io_utils/parquet_io_utils.cpp Co-authored-by: Yunsong Wang <12716979+PointKernel@users.noreply.github.com> --- 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 46ed795c54f2..bcb494f3a1b2 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -104,7 +104,7 @@ fetch_byte_ranges_to_device_async( std::vector column_chunk_data_per_source(num_sources); std::for_each( - cuda::counting_iterator(0), + cuda::counting_iterator(0), cuda::counting_iterator(num_sources), [&](auto const source_idx) { auto const& byte_ranges = byte_ranges_per_source[source_idx]; From 3a56e67406eabb9c58e94bc7fa1cc6eba5b39dc8 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Tue, 2 Jun 2026 01:09:28 +0000 Subject: [PATCH 13/13] Address comments --- .../io/parquet/io_utils/parquet_io_utils.cpp | 106 +++++++++--------- 1 file changed, 51 insertions(+), 55 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 ed62fd927719..7170d818cd87 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -36,7 +36,7 @@ namespace cudf::io::parquet { -namespace detail { +namespace { /** * @brief Dispatches the fetch task for each source index and collects the results @@ -83,7 +83,7 @@ auto dispatch_fetch_tasks(std::size_t num_sources, Task fetch_task) /** * @copydoc cudf::io::parquet::fetch_footers_to_host */ -std::vector> fetch_footers_to_host( +std::vector> fetch_footers_to_host_impl( cudf::host_span const> datasources) { // Helper to fetch footer from a datasource @@ -97,8 +97,8 @@ std::vector> fetch_footers_to_host 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(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"); @@ -113,7 +113,7 @@ std::vector> fetch_footers_to_host /** * @copydoc cudf::io::parquet::fetch_page_indexes_to_host */ -std::vector> fetch_page_indexes_to_host( +std::vector> fetch_page_indexes_to_host_impl( cudf::host_span const> datasources, cudf::host_span page_index_bytes_per_source) { @@ -139,56 +139,12 @@ std::vector> fetch_page_indexes_to }); } -} // 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); -} - -namespace detail { - using device_spans_per_source_type = std::vector>; std::tuple, std::vector, std::future> -fetch_byte_ranges_to_device_async( +fetch_byte_ranges_to_device_async_impl( cudf::host_span const> datasources, cudf::host_span const> byte_ranges_per_source, @@ -377,7 +333,47 @@ fetch_byte_ranges_to_device_async( std::async(std::launch::deferred, sync_function, std::move(device_read_tasks))}; } -} // namespace detail +} // namespace + +std::unique_ptr fetch_footer_to_host(cudf::io::datasource& datasource) +{ + CUDF_FUNC_RANGE(); + + // Wrap the input into an array and delegate to the multi-source implementation + std::array, 1> datasources{std::ref(datasource)}; + auto footer_buffers = fetch_footers_to_host_impl({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 fetch_footers_to_host_impl(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 multi-source implementation + std::array, 1> datasources{std::ref(datasource)}; + std::array page_index_bytes_per_source{page_index_bytes}; + + auto page_index_buffers = fetch_page_indexes_to_host_impl( + {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 fetch_page_indexes_to_host_impl(datasources, page_index_bytes_per_source); +} std::tuple, std::vector>, @@ -390,12 +386,12 @@ fetch_byte_ranges_to_device_async( { CUDF_FUNC_RANGE(); - // Wrap the inputs into arrays and delegate to the detail multi-source API + // Wrap the inputs into arrays and delegate to the multi-source implementation 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( + auto [buffers, fetched_byte_ranges, fut] = fetch_byte_ranges_to_device_async_impl( {datasources.data(), datasources.size()}, {byte_ranges_per_source.data(), byte_ranges_per_source.size()}, stream, @@ -415,13 +411,13 @@ fetch_byte_ranges_to_device_async( { CUDF_FUNC_RANGE(); - // Convert input vectors into host spans for detail API + // Convert input vectors into host spans for the implementation 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( + return fetch_byte_ranges_to_device_async_impl( datasources, {byte_range_spans_per_source.data(), byte_range_spans_per_source.size()}, stream,