diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index a49e300bd345..97e2fe97b5db 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -109,6 +109,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 baccb96ab537..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,106 +139,103 @@ 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); -} +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, +fetch_byte_ranges_to_device_async_impl( + 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; + static std::mutex host_read_mutex; + static std::mutex device_read_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; @@ -256,21 +253,32 @@ 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); + // 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); - 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)); - copy_sizes.push_back(io_size); - } - }); + 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()) { @@ -286,15 +294,18 @@ 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 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.is_device_read_preferred(io_size)) { device_read_tasks.emplace_back( @@ -318,8 +329,99 @@ 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 + +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>, + 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 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] = fetch_byte_ranges_to_device_async_impl( + {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 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 fetch_byte_ranges_to_device_async_impl( + datasources, + {byte_range_spans_per_source.data(), byte_range_spans_per_source.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.