diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index 2a511b7a9e39..ca10c215dd5b 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -17,6 +17,7 @@ #include #include #include +#include #include /** @@ -151,6 +152,48 @@ fetch_byte_ranges_to_device_async( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); +/** + * @brief Fetches Parquet bloom filter bitsets from a datasource into device buffers + * + * @ingroup io_utils + * + * @param datasource Input datasource + * @param bloom_filter_byte_ranges Byte ranges of complete bloom filters to fetch, must span a + * complete bloom filter + * @param stream CUDA stream + * @param mr Device memory resource used to allocate the returned device buffers + * + * @return A pair containing buffers that own the fetched bitsets and one device span per input byte + * range + */ +std::pair, std::vector>> +fetch_bloom_filters_to_device(cudf::io::datasource& datasource, + cudf::host_span bloom_filter_byte_ranges, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/** + * @brief Fetches Parquet bloom filter bitsets from multiple datasources into device buffers + * + * @ingroup io_utils + * + * @param datasources Input datasources + * @param bloom_filter_byte_ranges_per_source Byte ranges of complete bloom filters to fetch, one + * vector per datasource. Each byte range must span a complete bloom filter. + * @param stream CUDA stream + * @param mr Device memory resource used to allocate the returned device buffers + * + * @return A pair containing buffers that own the fetched bitsets and per-source device spans, with + * one inner vector per datasource + */ +std::pair, + std::vector>>> +fetch_bloom_filters_to_device( + cudf::host_span const> datasources, + cudf::host_span const> bloom_filter_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/bloom_filter_reader.cu b/cpp/src/io/parquet/bloom_filter_reader.cu index e0591f2d2115..eb7dfdc90b6b 100644 --- a/cpp/src/io/parquet/bloom_filter_reader.cu +++ b/cpp/src/io/parquet/bloom_filter_reader.cu @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -29,9 +30,12 @@ #include #include +#include #include #include #include +#include +#include namespace cudf::io::parquet::detail { namespace { @@ -299,219 +303,110 @@ class bloom_filter_expression_converter : public equality_literals_collector { std::unique_ptr _always_true; }; -/** - * @brief Reads bloom filter data to device. - * - * @param sources Dataset sources - * @param num_chunks Number of total column chunks to read - * @param bloom_filter_data Device buffers to hold bloom filter bitsets for each chunk - * @param bloom_filter_offsets Bloom filter offsets for all chunks - * @param bloom_filter_sizes Bloom filter sizes 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 aligned_mr Aligned device memory resource to allocate bloom filter buffers - */ -void read_bloom_filter_data(host_span const> sources, - std::size_t num_chunks, - cudf::host_span bloom_filter_data, - cudf::host_span> bloom_filter_offsets, - cudf::host_span> bloom_filter_sizes, - std::vector const& chunk_source_map, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref aligned_mr) -{ - // Using `arrow_filter_policy` with a temporary `cuda::std::byte` key type to extract bloom - // filter properties - using policy_type = arrow_filter_policy; - auto constexpr filter_block_alignment = - alignof(cuco::bloom_filter_ref, - cuco::thread_scope_thread, - policy_type>::filter_block_type); - auto constexpr words_per_block = policy_type::words_per_block; - - // Read tasks for bloom filter data - std::vector> read_tasks; - - // Read bloom filters for all column chunks - std::for_each( - cuda::counting_iterator{0}, - cuda::counting_iterator{num_chunks}, - [&](auto const chunk) { - // If bloom filter offset absent, fill in an empty buffer and skip ahead - if (not bloom_filter_offsets[chunk].has_value()) { - bloom_filter_data[chunk] = {}; - return; - } - // Read bloom filter iff present - auto const bloom_filter_offset = bloom_filter_offsets[chunk].value(); - - // If Bloom filter size (header + bitset) is available, just read the entire thing. - // Else just read 256 bytes which will contain the entire header and may contain the - // entire bitset as well. - auto constexpr bloom_filter_size_guess = 256; - auto const initial_read_size = - static_cast(bloom_filter_sizes[chunk].value_or(bloom_filter_size_guess)); - - // Read an initial buffer from source - auto& source = sources[chunk_source_map[chunk]]; - auto buffer = source->host_read(bloom_filter_offset, initial_read_size); - - // Deserialize the Bloom filter header from the buffer. - BloomFilterHeader header; - CompactProtocolReader cp{buffer->data(), buffer->size()}; - cp.read(&header); - - // Check if the bloom filter header is valid. - auto const is_header_valid = - (header.num_bytes % words_per_block) == 0 and - header.compression.compression == BloomFilterCompression::UNCOMPRESSED and - header.algorithm.algorithm == BloomFilterAlgorithm::SPLIT_BLOCK and - header.hash.hash == BloomFilterHash::XXHASH; - - // Do not read if the bloom filter is invalid - if (not is_header_valid) { - bloom_filter_data[chunk] = {}; - CUDF_LOG_WARN("Encountered an invalid bloom filter header. Skipping"); - return; - } - - // Bloom filter header size - auto const bloom_filter_header_size = static_cast(cp.bytecount()); - auto const bitset_size = static_cast(header.num_bytes); - - // Check if we already read in the filter bitset in the initial read. - if (initial_read_size >= bloom_filter_header_size + bitset_size) { - bloom_filter_data[chunk] = rmm::device_buffer{ - buffer->data() + bloom_filter_header_size, bitset_size, stream, aligned_mr}; - // The allocated bloom filter buffer must be aligned - CUDF_EXPECTS(reinterpret_cast(bloom_filter_data[chunk].data()) % - filter_block_alignment == - 0, - "Encountered misaligned bloom filter block"); - } - // Read the bitset from datasource. - else { - auto const bitset_offset = bloom_filter_offset + bloom_filter_header_size; - // Directly read to device if preferred - if (source->is_device_read_preferred(bitset_size)) { - bloom_filter_data[chunk] = rmm::device_buffer{bitset_size, stream, aligned_mr}; - // The allocated bloom filter buffer must be aligned - CUDF_EXPECTS(reinterpret_cast(bloom_filter_data[chunk].data()) % - filter_block_alignment == - 0, - "Encountered misaligned bloom filter block"); - auto future_read_size = - source->device_read_async(bitset_offset, - bitset_size, - static_cast(bloom_filter_data[chunk].data()), - stream); - - read_tasks.emplace_back(std::move(future_read_size)); - } else { - buffer = source->host_read(bitset_offset, bitset_size); - bloom_filter_data[chunk] = - rmm::device_buffer{buffer->data(), buffer->size(), stream, aligned_mr}; - // The allocated bloom filter buffer must be aligned - CUDF_EXPECTS(reinterpret_cast(bloom_filter_data[chunk].data()) % - filter_block_alignment == - 0, - "Encountered misaligned bloom filter block"); - } - } - }); - - // Read task sync function - for (auto& task : read_tasks) { - task.get(); - } -} - } // namespace -std::size_t aggregate_reader_metadata::get_bloom_filter_alignment() const +std::optional> parse_bloom_filter_header( + host_span bytes) { - // Required alignment: - // https://github.com/NVIDIA/cuCollections/blob/deab5799f3e4226cb8a49acf2199c03b14941ee4/include/cuco/detail/bloom_filter/bloom_filter_impl.cuh#L55-L67 - using policy_type = arrow_filter_policy; - auto constexpr alignment = alignof(cuco::bloom_filter_ref, - cuco::thread_scope_thread, - policy_type>::filter_block_type); - static_assert((alignment & (alignment - 1)) == 0, "Alignment must be a power of 2"); - return std::max(alignment, rmm::CUDA_ALLOCATION_ALIGNMENT); + using policy_type = arrow_filter_policy; + using word_type = typename policy_type::word_type; + auto constexpr bytes_per_block = sizeof(word_type) * policy_type::words_per_block; + + // Deserialize the bloom filter header from the front of the buffer + BloomFilterHeader header; + CompactProtocolReader cp{bytes.data(), bytes.size()}; + cp.read(&header); + + // Check if the bloom filter header is valid + auto const is_header_valid = + (header.num_bytes % bytes_per_block) == 0 and + header.compression.compression == BloomFilterCompression::UNCOMPRESSED and + header.algorithm.algorithm == BloomFilterAlgorithm::SPLIT_BLOCK and + header.hash.hash == BloomFilterHash::XXHASH; + if (not is_header_valid) { return std::nullopt; } + + return std::pair{static_cast(cp.bytecount()), + static_cast(header.num_bytes)}; } -std::vector aggregate_reader_metadata::read_bloom_filters( +std::pair, std::vector>> +aggregate_reader_metadata::read_bloom_filters( host_span const> sources, host_span const> row_group_indices, host_span column_schemas, size_type total_row_groups, rmm::cuda_stream_view stream, - rmm::device_async_resource_ref aligned_mr) const + rmm::device_async_resource_ref mr) const { // Descriptors for all the chunks that make up the selected columns auto const num_input_columns = column_schemas.size(); auto const num_chunks = total_row_groups * num_input_columns; - // Association between each column chunk and its source - std::vector chunk_source_map(num_chunks); - - // Keep track of column chunk file offsets - std::vector> bloom_filter_offsets(num_chunks); - std::vector> bloom_filter_sizes(num_chunks); - - // Gather all bloom filter offsets and sizes. - size_type chunk_count = 0; - // Flag to check if we have at least one valid bloom filter offset auto have_bloom_filters = false; - + // Speculatively read when a bloom filter's length is absent, enough to cover the header (and + // often the whole bitset). + auto constexpr speculative_read_size = int64_t{256}; + // Build complete bloom filter byte ranges (header + bitset) for every column chunk + std::vector> bloom_filter_byte_ranges_per_source( + row_group_indices.size()); // For all data sources - std::for_each(cuda::counting_iterator{0}, - cuda::counting_iterator{row_group_indices.size()}, - [&](auto const src_index) { - // Get all row group indices in the data source - auto const& rg_indices = row_group_indices[src_index]; - // For all row groups - std::for_each(rg_indices.cbegin(), rg_indices.cend(), [&](auto const rg_index) { - // For all column chunks - std::for_each( - column_schemas.begin(), column_schemas.end(), [&](auto const schema_idx) { - auto& col_meta = get_column_metadata(rg_index, src_index, schema_idx); - - // Get bloom filter offsets and sizes - bloom_filter_offsets[chunk_count] = col_meta.bloom_filter_offset; - bloom_filter_sizes[chunk_count] = col_meta.bloom_filter_length; - - // Set `have_bloom_filters` if `bloom_filter_offset` is valid - if (col_meta.bloom_filter_offset.has_value()) { have_bloom_filters = true; } - - // Map each column chunk to its source index - chunk_source_map[chunk_count] = src_index; - chunk_count++; - }); - }); - }); + std::for_each( + cuda::counting_iterator{0}, + cuda::counting_iterator{row_group_indices.size()}, + [&](auto const src_index) { + auto const& rg_indices = row_group_indices[src_index]; + auto& source_ranges = bloom_filter_byte_ranges_per_source[src_index]; + auto const source_size = static_cast(sources[src_index]->size()); + source_ranges.reserve(rg_indices.size() * num_input_columns); + // For all row groups in the source + std::for_each(rg_indices.cbegin(), rg_indices.cend(), [&](auto const rg_index) { + // For all column chunks in the row group + std::for_each(column_schemas.begin(), column_schemas.end(), [&](auto const schema_idx) { + auto const& col_meta = get_column_metadata(rg_index, src_index, schema_idx); + if (col_meta.bloom_filter_offset.has_value()) { + have_bloom_filters = true; + auto const offset = col_meta.bloom_filter_offset.value(); + CUDF_EXPECTS(offset >= 0 and offset < source_size, + "Bloom filter offset is out of datasource bounds"); + // Length absent: speculatively read enough to recover the header, clamped at EOF + auto const length = col_meta.bloom_filter_length.has_value() + ? static_cast(col_meta.bloom_filter_length.value()) + : std::min(speculative_read_size, source_size - offset); + CUDF_EXPECTS(length >= 0 and offset + length <= source_size, + "Bloom filter length is out of datasource bounds"); + source_ranges.push_back({offset, length}); + } else { + source_ranges.push_back({0, 0}); + } + }); + }); + }); // Exit early if we don't have any bloom filters if (not have_bloom_filters) { return {}; } - // Vector to hold bloom filter data - std::vector bloom_filter_data(num_chunks); - - // Read bloom filter data - read_bloom_filter_data(sources, - num_chunks, - bloom_filter_data, - bloom_filter_offsets, - bloom_filter_sizes, - chunk_source_map, - stream, - aligned_mr); - - // Return bloom filter data - return bloom_filter_data; + // Fetch the header-stripped, 32-byte-aligned bloom filter bitsets to device + std::vector> datasource_refs; + datasource_refs.reserve(sources.size()); + std::transform( + sources.begin(), sources.end(), std::back_inserter(datasource_refs), [](auto const& source) { + return std::ref(*source); + }); + + auto [bloom_filter_buffers, bitset_spans_per_source] = + fetch_bloom_filters_to_device(datasource_refs, bloom_filter_byte_ranges_per_source, stream, mr); + + // Flatten the per-source bitset spans into per-chunk order + std::vector> bloom_filter_data; + bloom_filter_data.reserve(num_chunks); + auto flat_bitset_spans = bitset_spans_per_source | std::views::join; + std::transform(flat_bitset_spans.begin(), + flat_bitset_spans.end(), + std::back_inserter(bloom_filter_data), + [](auto const& span) { return cuda::std::as_bytes(span); }); + + return {std::move(bloom_filter_buffers), std::move(bloom_filter_data)}; } std::optional>> aggregate_reader_metadata::apply_bloom_filters( 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 9740ae71a435..b6f0413c57c1 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -5,17 +5,20 @@ #include "io/comp/common.hpp" #include "io/parquet/parquet_common.hpp" +#include "io/parquet/reader_impl_helpers.hpp" #include #include #include #include #include +#include #include #include #include #include #include +#include #include #include @@ -25,16 +28,21 @@ #include #include +#include #include #include #include +#include +#include #include #include +#include #include #include #include #include #include +#include #include /** @@ -46,43 +54,66 @@ namespace cudf::io::parquet { namespace { +using device_spans_per_source_type = std::vector>; +using host_read_buffer = std::unique_ptr; + +/** + * @brief Serializes host-read submission batches to avoid cross-thread request interleaving and + * completion stalls + */ +std::mutex& host_read_mutex() +{ + static std::mutex mutex; + return mutex; +} + +/** + * @brief Serializes device reads and host-to-device copies to avoid cross-thread request + * interleaving and completion stalls + */ +std::mutex& device_read_mutex() +{ + static std::mutex mutex; + return mutex; +} + /** - * @brief Dispatches the fetch task for each source index and collects the results + * @brief Dispatches each indexed task and collects the results * - * Dispatches sequentially or using host worker pool depending on the number of sources. + * Dispatches sequentially or using host worker pool depending on the number of tasks. * - * @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 + * @tparam Task Callable invocable as `task(std::size_t task_idx)` + * @param num_tasks Number of tasks to dispatch + * @param task Task to run for each task index + * @return Vector of results, one per task, in task-index order */ template -auto dispatch_fetch_tasks(std::size_t num_sources, Task fetch_task) +auto dispatch_tasks(std::size_t num_tasks, Task task) { using result_type = std::invoke_result_t; auto constexpr parallel_threshold = 32; std::vector results; - results.reserve(num_sources); + results.reserve(num_tasks); - if (num_sources < parallel_threshold) { + if (num_tasks < 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)); }); + cuda::counting_iterator(num_tasks), + [&](std::size_t task_idx) { results.emplace_back(task(task_idx)); }); } else { // Dispatch the tasks to the host worker pool - std::vector> tasks; - tasks.reserve(num_sources); + std::vector> futures; + futures.reserve(num_tasks); 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); })); + cuda::counting_iterator(num_tasks), + [&](std::size_t task_idx) { + futures.emplace_back(cudf::detail::host_worker_pool().submit_task( + [&task, task_idx]() { return task(task_idx); })); }); - std::transform(tasks.begin(), tasks.end(), std::back_inserter(results), [](auto& task) { - return task.get(); + std::transform(futures.begin(), futures.end(), std::back_inserter(results), [](auto& fut) { + return fut.get(); }); } return results; @@ -148,7 +179,7 @@ std::vector> fetch_footers_to_host return cudf::io::datasource::buffer::create(std::move(footer_bytes)); }; - return dispatch_fetch_tasks(datasources.size(), [&](std::size_t source_idx) { + return dispatch_tasks(datasources.size(), [&](std::size_t source_idx) { return fetch_footer(datasources[source_idx].get()); }); } @@ -177,12 +208,74 @@ std::vector> fetch_page_indexes_to return datasource.host_read(page_index_bytes.offset(), page_index_bytes.size()); }; - return dispatch_fetch_tasks(datasources.size(), [&](std::size_t source_idx) { + return dispatch_tasks(datasources.size(), [&](std::size_t source_idx) { return fetch_page_index(datasources[source_idx].get(), page_index_bytes_per_source[source_idx]); }); } -using device_spans_per_source_type = std::vector>; +/** + * @brief Reads the given byte ranges into a caller-provided host buffer. + * + * Holds a mutex while scheduling so each thread's host reads are submitted contiguously. The ranges + * are packed consecutively into `dst` in iterator order; zero-size ranges are skipped. + * + * @param datasources Input datasources + * @param source_indices Iterator over the datasource index for each byte range + * @param offsets Iterator over the datasource offset for each byte range + * @param sizes Iterator over the size of each byte range + * @param count Number of byte ranges + * @param dst Host buffer that receives the read data + */ +template +void read_ranges_to_host( + cudf::host_span const> datasources, + SourceIndexIterator source_indices, + OffsetIterator offsets, + SizeIterator sizes, + std::size_t count, + cudf::host_span dst) +{ + std::vector> host_read_tasks; + std::vector expected_sizes; + host_read_tasks.reserve(count); + expected_sizes.reserve(count); + + auto iter = cuda::make_zip_iterator(source_indices, offsets, sizes); + std::size_t dst_byte_offset = 0; + + // 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 + count, [&](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 dst_offset = dst_byte_offset; + dst_byte_offset += io_size; + + if (io_size == 0) { return; } + + auto& datasource = datasources[src_idx].get(); + auto* const dst_ptr = dst.data() + dst_offset; + expected_sizes.push_back(io_size); + host_read_tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( + [&datasource, io_offset, io_size, dst_ptr]() -> std::size_t { + return datasource.host_read(io_offset, io_size, dst_ptr); + })); + }); + } + CUDF_EXPECTS(dst_byte_offset == dst.size(), "Unexpected destination host buffer size"); + + // Complete the reads; every range must be read in full + std::for_each(cuda::counting_iterator(0), + cuda::counting_iterator(host_read_tasks.size()), + [&](std::size_t i) { + CUDF_EXPECTS(host_read_tasks[i].get() == expected_sizes[i], + "Failed to read complete byte range to host"); + }); +} std::tuple, std::vector, @@ -194,9 +287,6 @@ fetch_byte_ranges_to_device_async_impl( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - static std::mutex host_read_mutex; - static std::mutex device_read_mutex; - auto const num_sources = datasources.size(); CUDF_EXPECTS(num_sources == byte_ranges_per_source.size(), @@ -280,8 +370,6 @@ fetch_byte_ranges_to_device_async_impl( io_source_indices.size() == io_offsets.size(), "Unexpected number of IO source indices, offsets, sizes, or destinations"); - using host_read_buffer = std::unique_ptr; - // Vectors to hold futures from datasource std::vector> device_read_tasks{}; std::vector> host_read_tasks{}; @@ -302,7 +390,7 @@ fetch_byte_ranges_to_device_async_impl( // 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::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); @@ -340,7 +428,7 @@ fetch_byte_ranges_to_device_async_impl( // 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(device_read_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); @@ -376,6 +464,187 @@ fetch_byte_ranges_to_device_async_impl( std::async(std::launch::deferred, sync_function, std::move(device_read_tasks))}; } +std::pair, std::vector> +fetch_bloom_filters_to_device_impl( + cudf::host_span const> datasources, + cudf::host_span const> + bloom_filter_byte_ranges_per_source, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto const num_sources = datasources.size(); + CUDF_EXPECTS(num_sources == bloom_filter_byte_ranges_per_source.size(), + "Encountered mismatch in number of datasources and bloom filter byte range spans"); + + // Each bitset must align to 32-byte boundaries, as required by cuco's Arrow bloom filter bitsets. + // TODO(NVIDIA/cuCollections#829): replace with a cuco-provided block-size / alignment accessor. + auto constexpr bloom_filter_block_bytes = std::size_t{32}; + + auto const total_filters = + std::accumulate(bloom_filter_byte_ranges_per_source.begin(), + bloom_filter_byte_ranges_per_source.end(), + std::size_t{0}, + [](auto acc, auto const& bloom_ranges) { return acc + bloom_ranges.size(); }); + + std::vector bitset_spans_per_source(num_sources); + + // Phase 1: Initial read. Cover the complete bloom filter or enough bytes to parse the header + std::vector initial_source_indices(total_filters); + std::vector initial_offsets(total_filters); + std::vector initial_sizes(total_filters); + std::vector initial_dst_offsets(total_filters); + std::size_t total_initial_read_size = 0; + + { + std::size_t filter_idx = 0; + std::for_each(cuda::counting_iterator(0), + cuda::counting_iterator(num_sources), + [&](auto const source_idx) { + auto const& bloom_ranges = bloom_filter_byte_ranges_per_source[source_idx]; + bitset_spans_per_source[source_idx].resize(bloom_ranges.size()); + std::for_each(bloom_ranges.begin(), bloom_ranges.end(), [&](auto const& range) { + initial_source_indices[filter_idx] = source_idx; + initial_offsets[filter_idx] = static_cast(range.offset()); + initial_sizes[filter_idx] = static_cast(range.size()); + initial_dst_offsets[filter_idx] = total_initial_read_size; + total_initial_read_size += initial_sizes[filter_idx]; + ++filter_idx; + }); + }); + CUDF_EXPECTS(filter_idx == total_filters, "Unexpected number of bloom filter byte ranges"); + } + + // Read every initial bloom filter bytes into one host buffer + auto initial_buffer = cudf::detail::make_host_vector(total_initial_read_size, stream); + read_ranges_to_host(datasources, + initial_source_indices.cbegin(), + initial_offsets.cbegin(), + initial_sizes.cbegin(), + total_filters, + initial_buffer); + + // Phase 2: Parse headers, organize bitset slots, and record deferred bitset reads + std::vector copy_srcs; + std::vector copy_sizes; + copy_srcs.reserve(total_filters); + copy_sizes.reserve(total_filters); + std::size_t total_device_size = 0; + + std::vector deferred_filter_indices; + std::vector deferred_offsets; + std::size_t total_deferred_size = 0; + + std::for_each( + cuda::counting_iterator(0), + cuda::counting_iterator(total_filters), + [&](std::size_t filter_idx) { + auto const push_empty_filter = [&]() { + copy_srcs.push_back(nullptr); + copy_sizes.push_back(0); + }; + + // Absent filter: no bloom filter to read + if (initial_sizes[filter_idx] == 0) { + push_empty_filter(); + return; + } + + auto const* const filter_addr = initial_buffer.data() + initial_dst_offsets[filter_idx]; + auto const header_info = + detail::parse_bloom_filter_header({filter_addr, initial_sizes[filter_idx]}); + if (not header_info.has_value()) { + CUDF_LOG_WARN("Encountered an invalid bloom filter header. Skipping"); + push_empty_filter(); + return; + } + auto const [header_bytes, bitset_bytes] = header_info.value(); + if (bitset_bytes % bloom_filter_block_bytes != 0) { + CUDF_LOG_WARN(std::format( + "Encountered a bloom filter bitset size that is not a multiple of {} bytes. Skipping", + bloom_filter_block_bytes)); + push_empty_filter(); + return; + } + + auto const header_size = static_cast(header_bytes); + auto const bitset_size = static_cast(bitset_bytes); + copy_sizes.push_back(bitset_size); + total_device_size += bitset_size; + + if (initial_sizes[filter_idx] >= header_size + bitset_size) { + // Whole bitset already in the host buffer: point at it, stripping the header + copy_srcs.push_back(filter_addr + header_size); + } else { + // Whole bitset not in the host buffer: defer the read + copy_srcs.push_back(nullptr); + + deferred_filter_indices.push_back(filter_idx); + deferred_offsets.push_back(initial_offsets[filter_idx] + header_size); + total_deferred_size += bitset_size; + } + }); + + // Phase 3: Resolve deferred reads, then batch copy all bitsets to the device + auto deferred_buffer = cudf::detail::make_host_vector(total_deferred_size, stream); + { + auto deferred_source_indices = + cuda::permutation_iterator{initial_source_indices.cbegin(), deferred_filter_indices.cbegin()}; + auto deferred_sizes = + cuda::permutation_iterator{copy_sizes.cbegin(), deferred_filter_indices.cbegin()}; + read_ranges_to_host(datasources, + deferred_source_indices, + deferred_offsets.cbegin(), + deferred_sizes, + deferred_filter_indices.size(), + deferred_buffer); + std::size_t deferred_dst_offset = 0; + std::for_each( + deferred_filter_indices.begin(), deferred_filter_indices.end(), [&](auto const filter_idx) { + copy_srcs[filter_idx] = deferred_buffer.data() + deferred_dst_offset; + deferred_dst_offset += copy_sizes[filter_idx]; + }); + } + + // Add the buffer base to every output span and copy destination. + rmm::device_buffer bitset_buffer(total_device_size, bloom_filter_block_bytes, stream, mr); + std::vector copy_dsts(total_filters); + auto* const device_base = static_cast(bitset_buffer.data()); + std::size_t device_offset = 0; + if (device_base != nullptr) { + std::for_each(cuda::counting_iterator(0), + cuda::counting_iterator(total_filters), + [&](std::size_t filter_idx) { + copy_dsts[filter_idx] = device_base + device_offset; + device_offset += copy_sizes[filter_idx]; + }); + } + CUDF_EXPECTS(device_offset == total_device_size, "Unexpected bloom filter device buffer size"); + + // Populate the nested per-source spans through a flattened view + auto flat_output_spans = bitset_spans_per_source | std::views::join; + std::transform(copy_dsts.begin(), + copy_dsts.end(), + copy_sizes.begin(), + flat_output_spans.begin(), + [](auto const dst, auto const size) { + return cudf::device_span{static_cast(dst), size}; + }); + + // One batched copy (entries with a null source or zero size are ignored by the batch API) + if (total_device_size != 0) { + { + std::scoped_lock lock(device_read_mutex()); + CUDF_CUDA_TRY(cudf::detail::memcpy_batch_async( + copy_dsts.data(), copy_srcs.data(), copy_sizes.data(), total_filters, stream)); + } + stream.synchronize(); + } + + std::vector bitset_buffers; + bitset_buffers.push_back(std::move(bitset_buffer)); + return {std::move(bitset_buffers), std::move(bitset_spans_per_source)}; +} + } // namespace [[nodiscard]] std::size_t metadata_size_hint() @@ -471,4 +740,52 @@ fetch_byte_ranges_to_device_async( mr); } +std::pair, std::vector>> +fetch_bloom_filters_to_device( + cudf::io::datasource& datasource, + cudf::host_span bloom_filter_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> + bloom_filter_byte_ranges_per_source{bloom_filter_byte_ranges}; + + auto [buffers, fetched_byte_ranges] = fetch_bloom_filters_to_device_impl( + {datasources.data(), datasources.size()}, + {bloom_filter_byte_ranges_per_source.data(), bloom_filter_byte_ranges_per_source.size()}, + stream, + mr); + + return {std::move(buffers), std::move(fetched_byte_ranges.front())}; +} + +std::pair, + std::vector>>> +fetch_bloom_filters_to_device( + cudf::host_span const> datasources, + cudf::host_span const> + bloom_filter_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> + bloom_filter_byte_range_spans_per_source; + bloom_filter_byte_range_spans_per_source.reserve(bloom_filter_byte_ranges_per_source.size()); + for (auto const& ranges : bloom_filter_byte_ranges_per_source) { + bloom_filter_byte_range_spans_per_source.emplace_back(ranges); + } + return fetch_bloom_filters_to_device_impl(datasources, + {bloom_filter_byte_range_spans_per_source.data(), + bloom_filter_byte_range_spans_per_source.size()}, + stream, + mr); +} + } // namespace cudf::io::parquet diff --git a/cpp/src/io/parquet/predicate_pushdown.cpp b/cpp/src/io/parquet/predicate_pushdown.cpp index 35b4b407e2ad..a762d16d575c 100644 --- a/cpp/src/io/parquet/predicate_pushdown.cpp +++ b/cpp/src/io/parquet/predicate_pushdown.cpp @@ -18,8 +18,6 @@ #include #include -#include - #include #include @@ -319,36 +317,22 @@ aggregate_reader_metadata::filter_row_groups( {std::make_optional(num_stats_filtered_row_groups), std::nullopt}}; } - // Aligned resource adaptor to allocate bloom filter buffers with - auto aligned_mr = rmm::mr::aligned_resource_adaptor(cudf::get_current_device_resource_ref(), - get_bloom_filter_alignment()); - // Read a vector of bloom filter bitset device buffers for all columns with equality // predicate(s) across all row groups - auto bloom_filter_buffers = read_bloom_filters(sources, - bloom_filter_input_row_groups, - equality_col_schemas, - num_stats_filtered_row_groups, - stream, - aligned_mr); - - // No bloom filter buffers, return early - if (bloom_filter_buffers.empty()) { + auto const [bloom_filter_buffers, bloom_filter_data] = + read_bloom_filters(sources, + bloom_filter_input_row_groups, + equality_col_schemas, + num_stats_filtered_row_groups, + stream, + cudf::get_current_device_resource_ref()); + + // No bloom filters, return early + if (bloom_filter_data.empty()) { return {stats_filtered_row_groups, {std::make_optional(num_stats_filtered_row_groups), std::nullopt}}; } - // Create spans from bloom filter buffers - std::vector> bloom_filter_data; - bloom_filter_data.reserve(bloom_filter_buffers.size()); - std::transform(bloom_filter_buffers.begin(), - bloom_filter_buffers.end(), - std::back_inserter(bloom_filter_data), - [](auto& buffer) { - return cudf::device_span( - static_cast(buffer.data()), buffer.size()); - }); - // Apply bloom filtering on the output row groups from stats filter auto const bloom_filtered_row_groups = apply_bloom_filters(bloom_filter_data, bloom_filter_input_row_groups, diff --git a/cpp/src/io/parquet/reader_impl_helpers.hpp b/cpp/src/io/parquet/reader_impl_helpers.hpp index 40c70babf2b8..c9f70d3b49f1 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.hpp +++ b/cpp/src/io/parquet/reader_impl_helpers.hpp @@ -15,10 +15,12 @@ #include #include +#include #include #include #include #include +#include #include namespace cudf::io::parquet::detail { @@ -185,6 +187,17 @@ struct column_selection_options { bool case_sensitive_names = true; }; +/** + * @brief Parses and validates a Parquet `BloomFilterHeader` from the front of `bytes` + * + * @param bytes Host bytes starting at the beginning of a bloom filter (header followed by bitset) + * + * @return A pair of the bloom filter header size and the bitset size in bytes, or `std::nullopt` + * if the header is missing or unsupported + */ +[[nodiscard]] std::optional> parse_bloom_filter_header( + host_span bytes); + class aggregate_reader_metadata { protected: std::vector per_file_metadata; @@ -256,11 +269,6 @@ class aggregate_reader_metadata { */ void column_info_for_row_group(row_group_info& rg_info, size_t chunk_start_row) const; - /** - * @brief Returns the required alignment for bloom filter buffers - */ - [[nodiscard]] size_t get_bloom_filter_alignment() const; - /** * @brief Reads bloom filter bitsets for the specified columns from the given lists of row * groups. @@ -270,18 +278,19 @@ class aggregate_reader_metadata { * @param column_schemas Schema indices of columns whose bloom filters will be read * @param num_row_groups Number of row groups in the file * @param stream CUDA stream used for device memory operations and kernel launches - * @param aligned_mr Aligned device memory resource to allocate bloom filter buffers - * - * @return A flattened list of bloom filter bitset device buffers for each predicate column across - * row group - */ - [[nodiscard]] std::vector read_bloom_filters( - host_span const> sources, - host_span const> row_group_indices, - host_span column_schemas, - size_type num_row_groups, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref aligned_mr) const; + * @param mr Device memory resource used to allocate bloom filter buffers + * + * @return A pair of the device buffers backing the bloom filter bitsets and a flattened, + * per-chunk list of bitset device spans (empty spans for chunks without a bloom filter) + */ + [[nodiscard]] std::pair, + std::vector>> + read_bloom_filters(host_span const> sources, + host_span const> row_group_indices, + host_span column_schemas, + size_type num_row_groups, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; /** * @brief Collects Parquet types for the columns with the specified schema indices