diff --git a/cpp/include/cudf/io/experimental/hybrid_scan.hpp b/cpp/include/cudf/io/experimental/hybrid_scan.hpp index bab4bf3fb93e..86c23db4ca36 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -513,7 +513,7 @@ class hybrid_scan_reader { * * @param row_group_indices Input row groups indices * @param column_chunk_data Device spans of column chunk data of payload columns - * @param row_mask Boolean column indicating which rows need to be read. All rows read if empty + * @param row_mask Boolean column indicating which rows need to be read * @param mask_data_pages Whether to build and use a data page mask using the row mask * @param options Parquet reader options * @param stream CUDA stream used for device memory operations and kernel launches @@ -564,7 +564,7 @@ class hybrid_scan_reader { * @param pass_read_limit Limit on the memory used for reading and decompressing data. `0` if * there is no limit * @param row_group_indices Input row groups indices - * @param row_mask Boolean column indicating which rows need to be read. All rows read if empty + * @param row_mask Boolean column indicating which rows need to be read * @param mask_data_pages Whether to build and use a data page mask using the row mask * @param column_chunk_data Device spans of column chunk data of filter columns * @param options Parquet reader options @@ -601,7 +601,7 @@ class hybrid_scan_reader { * @param pass_read_limit Limit on the memory used for reading and decompressing data. `0` if * there is no limit * @param row_group_indices Input row groups indices - * @param row_mask Boolean column indicating which rows need to be read. All rows read if empty + * @param row_mask Boolean column indicating which rows need to be read * @param mask_data_pages Whether to build and use a data page mask using the row mask * @param column_chunk_data Device spans of column chunk data of payload columns * @param options Parquet reader options @@ -623,7 +623,7 @@ class hybrid_scan_reader { * @brief Materializes a chunk of payload columns and applies the corresponding range of input row * mask to the output table chunk * - * @param row_mask Boolean column indicating which rows need to be read. All rows read if empty + * @param row_mask Boolean column indicating which rows need to be read * * @return Table chunk of materialized filter columns and metadata */ diff --git a/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp index fd3ef6ac5523..b6b9dc1f3842 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp @@ -100,7 +100,7 @@ class hybrid_scan_multifile { * @brief Get all available per-source row group indices from the parquet files * * @param options Parquet reader options - * @return Vector of row group indices, one inner vector per source + * @return Vector of vectors of row group indices, one per source */ [[nodiscard]] std::vector> all_row_groups( parquet_reader_options const& options) const; @@ -108,7 +108,7 @@ class hybrid_scan_multifile { /** * @brief Get the total number of top-level rows in the per-source row groups * - * @param row_group_indices Input per-source row group indices (one inner vector per source) + * @param row_group_indices Span of vectors of input row group indices, one per source * @return Total number of top-level rows across all sources */ [[nodiscard]] size_type total_rows_in_row_groups( @@ -130,9 +130,9 @@ class hybrid_scan_multifile { * Filters the row groups such that only the row groups that start within the byte range are * selected. Note that the last selected row group may end beyond the byte range. * - * @param row_group_indices Input row group indices, one per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param options Parquet reader options - * @return Filtered per-source row group indices (one inner vector per source) + * @return Vector of vectors of filtered row group indices, one per source */ [[nodiscard]] std::vector> filter_row_groups_with_byte_range( cudf::host_span const> row_group_indices, @@ -141,10 +141,10 @@ class hybrid_scan_multifile { /** * @brief Filter the input row groups using column chunk statistics * - * @param row_group_indices Input row group indices, one per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param options Parquet reader options * @param stream CUDA stream used for device memory operations and kernel launches - * @return Filtered row group indices, one per source + * @return Vector of vectors of filtered row group indices, one per source */ [[nodiscard]] std::vector> filter_row_groups_with_stats( cudf::host_span const> row_group_indices, @@ -158,7 +158,7 @@ class hybrid_scan_multifile { * @note Device buffers for bloom filter byte ranges must be allocated using a 32 byte * aligned memory resource * - * @param row_group_indices Input row group indices, one per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param options Parquet reader options * @return Pair of vectors of byte ranges of column chunk with bloom filters and dictionary * pages subject to filter predicate @@ -171,7 +171,7 @@ class hybrid_scan_multifile { * @brief Builds a boolean survival column of size equal to the total number of rows in the row * groups containing all `true` values * - * @param row_group_indices Input per-source row group indices (one inner vector per source) + * @param row_group_indices Span of vectors of input row group indices, one per source * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used to allocate the returned column's device memory * @return An all-true boolean (survival) column spanning all selected rows across all sources @@ -185,7 +185,7 @@ class hybrid_scan_multifile { * @brief Builds a boolean column indicating surviving rows using page-level statistics in the * page index * - * @param row_group_indices Input per-source row group indices (one inner vector per source) + * @param row_group_indices Span of vectors of input row group indices, one per source * @param options Parquet reader options * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used to allocate the returned column's device memory @@ -206,7 +206,7 @@ class hybrid_scan_multifile { * map has one source index per byte range and can be used to regroup byte ranges by datasource * before fetching. * - * @param row_group_indices Input row group indices, one inner vector per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param options Parquet reader options * @return Pair of flattened byte ranges to column chunks of filter columns and their * corresponding source indices @@ -219,7 +219,7 @@ class hybrid_scan_multifile { * @brief Materializes filter columns and updates the input row mask to only the rows that exist * in the output table * - * @param row_group_indices Input row group indices, one inner vector per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param column_chunk_data Flattened device spans of filter column chunk data returned in the * same order as `filter_column_chunks_byte_ranges` * @param[in,out] row_mask Mutable boolean column spanning all selected rows across all sources @@ -247,7 +247,7 @@ class hybrid_scan_multifile { * map has one source index per byte range and can be used to regroup byte ranges by datasource * before fetching. * - * @param row_group_indices Input row group indices, one inner vector per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param options Parquet reader options * @return Pair of flattened byte ranges to column chunks of payload columns and their * corresponding source indices @@ -259,11 +259,11 @@ class hybrid_scan_multifile { /** * @brief Materialize payload columns and applies the row mask to the output table * - * @param row_group_indices Input row group indices, one inner vector per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param column_chunk_data Flattened device spans of payload column chunk data returned in the * same order as `payload_column_chunks_byte_ranges` * @param row_mask Boolean column spanning all selected rows across all sources and indicating - * which rows need to be read. All rows read if empty + * which rows need to be read * @param mask_data_pages Whether to build and use a data page mask using the row mask * @param options Parquet reader options * @param stream CUDA stream used for device memory operations and kernel launches @@ -282,7 +282,7 @@ class hybrid_scan_multifile { /** * @brief Get byte ranges of column chunks of all (or selected) columns * - * @param row_group_indices Input row group indices, one inner vector per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param options Parquet reader options * @return Pair of flattened byte ranges to column chunks of all (or selected) columns and their * corresponding source indices @@ -294,7 +294,7 @@ class hybrid_scan_multifile { /** * @brief Materializes all (or selected) columns and returns the final output table * - * @param row_group_indices Input row group indices, one inner vector per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param column_chunk_data Flattened device spans of column chunk data returned in the same order * as `all_column_chunks_byte_ranges` * @param options Parquet reader options @@ -309,6 +309,119 @@ class hybrid_scan_multifile { rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const; + /** + * @brief Setup chunking information for filter columns and preprocess the input data pages + * + * @param chunk_read_limit Limit on total number of bytes to be returned per table chunk. `0` if + * there is no limit + * @param pass_read_limit Limit on the memory used for reading and decompressing data. `0` if + * there is no limit + * @param row_group_indices Span of vectors of input row group indices, one per source + * @param row_mask Boolean column spanning all selected rows across all sources and indicating + * which rows need to be read + * @param mask_data_pages Whether to build and use a data page mask using the row mask + * @param column_chunk_data Flattened device spans of filter column chunk data returned in the + * same order as `filter_column_chunks_byte_ranges` + * @param options Parquet reader options + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the device memory for the output table chunks + */ + void setup_chunking_for_filter_columns( + std::size_t chunk_read_limit, + std::size_t pass_read_limit, + cudf::host_span const> row_group_indices, + cudf::column_view const& row_mask, + use_data_page_mask mask_data_pages, + cudf::host_span const> column_chunk_data, + parquet_reader_options const& options, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + + /** + * @brief Materializes a chunk of filter columns and updates the corresponding range of input row + * mask to only the rows that exist in the output table + * + * @param[in,out] row_mask Mutable boolean column spanning all selected rows across all sources + * and indicating surviving rows from page pruning. The row mask size must equal the total + * number of rows in the input row groups, and is empty only when there are no such rows + * (yielding an empty output table) + * + * @return Table chunk of materialized filter columns and metadata + */ + [[nodiscard]] table_with_metadata materialize_filter_columns_chunk( + cudf::mutable_column_view& row_mask) const; + + /** + * @brief Setup chunking information for payload columns and preprocess the input data pages + * + * @param chunk_read_limit Limit on total number of bytes to be returned per table chunk. `0` if + * there is no limit + * @param pass_read_limit Limit on the memory used for reading and decompressing data. `0` if + * there is no limit + * @param row_group_indices Span of vectors of input row group indices, one per source + * @param row_mask Boolean column spanning all selected rows across all sources and indicating + * which rows need to be read + * @param mask_data_pages Whether to build and use a data page mask using the row mask + * @param column_chunk_data Flattened device spans of payload column chunk data returned in the + * same order as `payload_column_chunks_byte_ranges` + * @param options Parquet reader options + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the device memory for the output table chunks + */ + void setup_chunking_for_payload_columns( + std::size_t chunk_read_limit, + std::size_t pass_read_limit, + cudf::host_span const> row_group_indices, + cudf::column_view const& row_mask, + use_data_page_mask mask_data_pages, + cudf::host_span const> column_chunk_data, + parquet_reader_options const& options, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + + /** + * @brief Materializes a chunk of payload columns and applies the corresponding range of input row + * mask to the output table chunk + * + * @param row_mask Boolean column spanning all selected rows across all sources and indicating + * which rows need to be read + * + * @return Table chunk of materialized payload columns and metadata + */ + [[nodiscard]] table_with_metadata materialize_payload_columns_chunk( + cudf::column_view const& row_mask) const; + + /** + * @brief Setup chunking information for all (or selected) columns and preprocess the input data + * pages + * + * @param chunk_read_limit Limit on total number of bytes to be returned per table chunk. `0` if + * there is no limit + * @param pass_read_limit Limit on the memory used for reading and decompressing data. `0` if + * there is no limit + * @param row_group_indices Span of vectors of input row group indices, one per source + * @param column_chunk_data Flattened device spans of column chunk data returned in the same order + * as `all_column_chunks_byte_ranges` + * @param options Parquet reader options + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the device memory for the output table chunks + */ + void setup_chunking_for_all_columns( + std::size_t chunk_read_limit, + std::size_t pass_read_limit, + cudf::host_span const> row_group_indices, + cudf::host_span const> column_chunk_data, + parquet_reader_options const& options, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + + /** + * @brief Materializes a chunk of all (or selected) columns and returns the output table chunk + * + * @return Table chunk of materialized all (or selected) columns and metadata + */ + [[nodiscard]] table_with_metadata materialize_all_columns_chunk() const; + /** * @brief Partition row groups into passes such that the amount of GPU memory required to read, * decompress and decode a pass is bounded by the specified limit @@ -320,7 +433,7 @@ class hybrid_scan_multifile { * * @throws std::invalid_argument if no row group indices in the input * - * @param row_group_indices Input row group indices, one per source + * @param row_group_indices Span of vectors of input row group indices, one per source * @param pass_read_limit Memory limit to read and decompress row group data, `0` if there is * no limit (single pass) * @@ -330,6 +443,13 @@ class hybrid_scan_multifile { cudf::host_span const> row_group_indices, std::size_t pass_read_limit) const; + /** + * @brief Check if there is any parquet data left to read for the current chunked setup + * + * @return Boolean indicating if there is any data left to read + */ + [[nodiscard]] bool has_next_table_chunk() const; + private: std::unique_ptr _impl; }; diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 665d431593fa..a1e31c700046 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -500,10 +500,9 @@ table_with_metadata hybrid_scan_reader_impl::materialize_filter_columns( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered"); + CUDF_EXPECTS(std::cmp_equal(row_mask.size(), total_rows_in_row_groups(row_group_indices)), + "Row mask must span across all input row groups"); CUDF_EXPECTS(options.get_filter().has_value(), "Empty input filter expression encountered"); - CUDF_EXPECTS(not row_mask.is_empty(), - "Row mask must be non-empty when materializing filter columns"); prepare_materialization( read_columns_mode::FILTER_COLUMNS, row_group_indices.size(), options, stream, mr); @@ -511,6 +510,14 @@ table_with_metadata hybrid_scan_reader_impl::materialize_filter_columns( // Convert the input expression (must be done after prepare_materialization) _expr_conv = build_converted_expression(options); + // Return early if all rows are pruned + if (are_all_rows_pruned(row_mask, stream)) { + auto const empty_row_groups = + std::vector>(row_group_indices.size(), std::vector{}); + prepare_data(read_mode::READ_ALL, empty_row_groups, {}, {}); + return read_chunk_internal(read_mode::READ_ALL, read_columns_mode::FILTER_COLUMNS, row_mask); + } + auto data_page_mask = thrust::host_vector{}; if (mask_data_pages == use_data_page_mask::YES) { data_page_mask = _extended_metadata->compute_data_page_mask( @@ -531,13 +538,22 @@ table_with_metadata hybrid_scan_reader_impl::materialize_payload_columns( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered"); + CUDF_EXPECTS(std::cmp_equal(row_mask.size(), total_rows_in_row_groups(row_group_indices)), + "Row mask must span across all input row groups"); CUDF_EXPECTS(row_mask.null_count() == 0, "Row mask must not have any nulls when materializing payload column"); prepare_materialization( read_columns_mode::PAYLOAD_COLUMNS, row_group_indices.size(), options, stream, mr); + // Return early if all rows are pruned + if (are_all_rows_pruned(row_mask, stream)) { + auto const empty_row_groups = + std::vector>(row_group_indices.size(), std::vector{}); + prepare_data(read_mode::READ_ALL, empty_row_groups, {}, {}); + return read_chunk_internal(read_mode::READ_ALL, read_columns_mode::PAYLOAD_COLUMNS, row_mask); + } + auto data_page_mask = thrust::host_vector{}; if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) { data_page_mask = _extended_metadata->compute_data_page_mask( @@ -586,10 +602,9 @@ void hybrid_scan_reader_impl::setup_chunking_for_filter_columns( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered"); CUDF_EXPECTS(options.get_filter().has_value(), "Empty input filter expression encountered"); - CUDF_EXPECTS(not row_mask.is_empty(), - "Row mask must be non-empty when setting up chunking for filter columns"); + CUDF_EXPECTS(std::cmp_equal(row_mask.size(), total_rows_in_row_groups(row_group_indices)), + "Row mask must span across all input row groups"); prepare_materialization( read_columns_mode::FILTER_COLUMNS, row_group_indices.size(), options, stream, mr); @@ -600,6 +615,14 @@ void hybrid_scan_reader_impl::setup_chunking_for_filter_columns( // Convert the input expression (must be done after prepare_materialization) _expr_conv = build_converted_expression(options); + // Return early if all rows are pruned + if (are_all_rows_pruned(row_mask, stream)) { + auto const empty_row_groups = + std::vector>(row_group_indices.size(), std::vector{}); + prepare_data(read_mode::CHUNKED_READ, empty_row_groups, {}, {}); + return; + } + auto data_page_mask = thrust::host_vector{}; if (mask_data_pages == use_data_page_mask::YES) { data_page_mask = _extended_metadata->compute_data_page_mask( @@ -639,7 +662,8 @@ void hybrid_scan_reader_impl::setup_chunking_for_payload_columns( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered"); + CUDF_EXPECTS(std::cmp_equal(row_mask.size(), total_rows_in_row_groups(row_group_indices)), + "Row mask must span across all input row groups"); CUDF_EXPECTS(row_mask.null_count() == 0, "Row mask must not have any nulls when materializing payload column"); @@ -649,6 +673,14 @@ void hybrid_scan_reader_impl::setup_chunking_for_payload_columns( _input_pass_read_limit = pass_read_limit; _output_chunk_read_limit = chunk_read_limit; + // Return early if all rows are pruned + if (are_all_rows_pruned(row_mask, stream)) { + auto const empty_row_groups = + std::vector>(row_group_indices.size(), std::vector{}); + prepare_data(read_mode::CHUNKED_READ, empty_row_groups, {}, {}); + return; + } + auto data_page_mask = thrust::host_vector{}; if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) { data_page_mask = _extended_metadata->compute_data_page_mask( @@ -1065,11 +1097,13 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( // Create a table from the output columns. auto read_table = std::make_unique(std::move(out_columns)); - CUDF_EXPECTS(row_mask.is_empty() or row_mask.type().id() == type_id::BOOL8, - "Input row mask must be empty or a boolean column"); + // If the input row mask is empty, all rows are pruned anyway. + if (row_mask.is_empty()) { + _output_chunk_produced = true; + return {std::move(read_table), std::move(out_metadata)}; + } - // If the input row mask is empty, return the table as is. - if (row_mask.is_empty()) { return {std::move(read_table), std::move(out_metadata)}; } + CUDF_EXPECTS(row_mask.type().id() == type_id::BOOL8, "Input row mask must be a boolean column"); // Get the current row mask offset auto const mask_offset = _row_mask_offset; diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp index 0607e684466a..0f1141d84f2e 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -59,47 +59,48 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { parquet_reader_options const& options); /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::parquet_metadatas + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::parquet_metadatas */ [[nodiscard]] std::vector parquet_metadatas() const; /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::page_index_byte_ranges + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::page_index_byte_ranges */ [[nodiscard]] std::vector page_index_byte_ranges() const; /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::setup_page_indexes + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::setup_page_indexes */ void setup_page_indexes( cudf::host_span const> page_index_bytes) const; /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::all_row_groups + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::all_row_groups */ [[nodiscard]] std::vector> all_row_groups( parquet_reader_options const& options) const; /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::total_rows_in_row_groups + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::total_rows_in_row_groups */ [[nodiscard]] std::size_t total_rows_in_row_groups( cudf::host_span const> row_group_indices) const; /** - * @copydoc cudf::io::experimental::hybrid_scan::reset_column_selection + * @copydoc cudf::io::parquet::experimental::hybrid_scan::reset_column_selection */ void reset_column_selection(); /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::filter_row_groups_with_byte_range + * @copydoc + * cudf::io::parquet::experimental::hybrid_scan_multifile::filter_row_groups_with_byte_range */ [[nodiscard]] std::vector> filter_row_groups_with_byte_range( cudf::host_span const> row_group_indices, parquet_reader_options const& options) const; /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::filter_row_groups_with_stats + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::filter_row_groups_with_stats */ [[nodiscard]] std::vector> filter_row_groups_with_stats( cudf::host_span const> row_group_indices, @@ -107,14 +108,14 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::cuda_stream_view stream); /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::secondary_filters_byte_ranges + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::secondary_filters_byte_ranges */ [[nodiscard]] std::pair, std::vector> secondary_filters_byte_ranges(cudf::host_span const> row_group_indices, parquet_reader_options const& options); /** - * @copydoc cudf::io::experimental::hybrid_scan::filter_row_groups_with_dictionary_pages + * @copydoc cudf::io::parquet::experimental::hybrid_scan::filter_row_groups_with_dictionary_pages */ [[nodiscard]] std::vector> filter_row_groups_with_dictionary_pages( cudf::host_span const> dictionary_page_data, @@ -123,7 +124,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::cuda_stream_view stream); /** - * @copydoc cudf::io::experimental::hybrid_scan::filter_row_groups_with_bloom_filters + * @copydoc cudf::io::parquet::experimental::hybrid_scan::filter_row_groups_with_bloom_filters */ [[nodiscard]] std::vector> filter_row_groups_with_bloom_filters( cudf::host_span const> bloom_filter_data, @@ -132,7 +133,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::cuda_stream_view stream); /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::build_all_true_row_mask + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::build_all_true_row_mask */ [[nodiscard]] std::unique_ptr build_all_true_row_mask( cudf::host_span const> row_group_indices, @@ -140,7 +141,8 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::device_async_resource_ref mr); /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::build_row_mask_with_page_index_stats + * @copydoc + * cudf::io::parquet::experimental::hybrid_scan_multifile::build_row_mask_with_page_index_stats */ [[nodiscard]] std::unique_ptr build_row_mask_with_page_index_stats( cudf::host_span const> row_group_indices, @@ -161,7 +163,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { parquet_reader_options const& options); /** - * @copydoc cudf::io::experimental::hybrid_scan::materialize_filter_columns + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::materialize_filter_columns */ [[nodiscard]] table_with_metadata materialize_filter_columns( cudf::host_span const> row_group_indices, @@ -185,7 +187,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { parquet_reader_options const& options); /** - * @copydoc cudf::io::experimental::hybrid_scan::materialize_payload_columns + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::materialize_payload_columns */ [[nodiscard]] table_with_metadata materialize_payload_columns( cudf::host_span const> row_group_indices, @@ -197,14 +199,14 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::device_async_resource_ref mr); /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::all_column_chunks_byte_ranges + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::all_column_chunks_byte_ranges */ [[nodiscard]] std::pair, std::vector> all_column_chunks_byte_ranges(cudf::host_span const> row_group_indices, parquet_reader_options const& options); /** - * @copydoc cudf::io::experimental::hybrid_scan_multifile::materialize_all_columns + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::materialize_all_columns */ [[nodiscard]] table_with_metadata materialize_all_columns( cudf::host_span const> row_group_indices, @@ -214,7 +216,8 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::device_async_resource_ref mr); /** - * @copydoc cudf::io::experimental::hybrid_scan::setup_chunking_for_filter_columns + * @copydoc + * cudf::io::parquet::experimental::hybrid_scan_multifile::setup_chunking_for_filter_columns */ void setup_chunking_for_filter_columns( std::size_t chunk_read_limit, @@ -228,13 +231,15 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::device_async_resource_ref mr); /** - * @copydoc cudf::io::experimental::hybrid_scan::materialize_filter_columns_chunk + * @copydoc + * cudf::io::parquet::experimental::hybrid_scan_multifile::materialize_filter_columns_chunk */ [[nodiscard]] table_with_metadata materialize_filter_columns_chunk( cudf::mutable_column_view& row_mask); /** - * @copydoc cudf::io::experimental::hybrid_scan::setup_chunking_for_payload_columns + * @copydoc + * cudf::io::parquet::experimental::hybrid_scan_multifile::setup_chunking_for_payload_columns */ void setup_chunking_for_payload_columns( std::size_t chunk_read_limit, @@ -248,13 +253,14 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::device_async_resource_ref mr); /** - * @copydoc cudf::io::experimental::hybrid_scan::materialize_payload_columns_chunk + * @copydoc + * cudf::io::parquet::experimental::hybrid_scan_multifile::materialize_payload_columns_chunk */ [[nodiscard]] table_with_metadata materialize_payload_columns_chunk( cudf::column_view const& row_mask); /** - * @copydoc cudf::io::experimental::hybrid_scan::setup_chunking_for_all_columns + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::setup_chunking_for_all_columns */ void setup_chunking_for_all_columns( std::size_t chunk_read_limit, @@ -266,7 +272,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::device_async_resource_ref mr); /** - * @copydoc cudf::io::experimental::hybrid_scan::materialize_all_columns_chunk + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::materialize_all_columns_chunk */ [[nodiscard]] table_with_metadata materialize_all_columns_chunk(); @@ -290,27 +296,10 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { std::size_t pass_read_limit) const; /** - * @copydoc cudf::io::experimental::hybrid_scan::has_next_table_chunk + * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::has_next_table_chunk */ [[nodiscard]] bool has_next_table_chunk(); - /** - * @brief Updates the output row mask such that such that out_row_mask[i + out_row_mask_offset] = - * true if and only if in_row_mask[i] is valid and true - * - * Updates the output row mask to reflect the final valid and surviving rows from the input row - * mask. This is inline with the masking behavior of cudf::detail::apply_boolean_mask - * - * @param in_row_mask Input row mask column - * @param out_row_mask Output row mask column - * @param out_row_mask_offset Offset into the output row mask column - * @param stream CUDA stream - */ - static void update_row_mask(cudf::column_view const& in_row_mask, - cudf::mutable_column_view& out_row_mask, - cudf::size_type out_row_mask_offset, - rmm::cuda_stream_view stream); - private: /** * @brief Enum indicating whether we are reading the filter, payload, or all columns @@ -525,6 +514,33 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { read_columns_mode read_columns_mode, RowMaskView row_mask); + /** + * @brief Check if all rows are pruned (all valid and false) + * + * @param row_mask Input row mask column + * @param stream CUDA stream used for device memory operations and kernel launches + * @return True if all rows are pruned + */ + [[nodiscard]] bool are_all_rows_pruned(cudf::column_view const& row_mask, + rmm::cuda_stream_view stream) const; + + /** + * @brief Updates the output row mask such that out_row_mask[i + out_row_mask_offset] = true if + * and only if in_row_mask[i] is valid and true + * + * Updates the output row mask to reflect the final valid and surviving rows from the input row + * mask. This is inline with the masking behavior of cudf::detail::apply_boolean_mask + * + * @param in_row_mask Input row mask column + * @param out_row_mask Output row mask column + * @param out_row_mask_offset Offset into the output row mask column + * @param stream CUDA stream + */ + void update_row_mask(cudf::column_view const& in_row_mask, + cudf::mutable_column_view& out_row_mask, + cudf::size_type out_row_mask_offset, + rmm::cuda_stream_view stream); + /** * @brief Check if this is the first output chunk * @@ -535,7 +551,6 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { return _file_itm_data._output_chunk_count == 0 and not _output_chunk_produced; } - private: aggregate_reader_metadata* _extended_metadata; std::optional> _filter_columns_names; diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp index 162daacd9d14..7ac700096ad9 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp @@ -171,6 +171,88 @@ table_with_metadata hybrid_scan_multifile::materialize_all_columns( return _impl->materialize_all_columns(row_group_indices, column_chunk_data, options, stream, mr); } +void hybrid_scan_multifile::setup_chunking_for_filter_columns( + std::size_t chunk_read_limit, + std::size_t pass_read_limit, + cudf::host_span const> row_group_indices, + cudf::column_view const& row_mask, + use_data_page_mask mask_data_pages, + cudf::host_span const> column_chunk_data, + parquet_reader_options const& options, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + CUDF_FUNC_RANGE(); + _impl->setup_chunking_for_filter_columns(chunk_read_limit, + pass_read_limit, + row_group_indices, + row_mask, + mask_data_pages, + column_chunk_data, + options, + stream, + mr); +} + +table_with_metadata hybrid_scan_multifile::materialize_filter_columns_chunk( + cudf::mutable_column_view& row_mask) const +{ + CUDF_FUNC_RANGE(); + return _impl->materialize_filter_columns_chunk(row_mask); +} + +void hybrid_scan_multifile::setup_chunking_for_payload_columns( + std::size_t chunk_read_limit, + std::size_t pass_read_limit, + cudf::host_span const> row_group_indices, + cudf::column_view const& row_mask, + use_data_page_mask mask_data_pages, + cudf::host_span const> column_chunk_data, + parquet_reader_options const& options, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + CUDF_FUNC_RANGE(); + _impl->setup_chunking_for_payload_columns(chunk_read_limit, + pass_read_limit, + row_group_indices, + row_mask, + mask_data_pages, + column_chunk_data, + options, + stream, + mr); +} + +table_with_metadata hybrid_scan_multifile::materialize_payload_columns_chunk( + cudf::column_view const& row_mask) const +{ + CUDF_FUNC_RANGE(); + return _impl->materialize_payload_columns_chunk(row_mask); +} + +void hybrid_scan_multifile::setup_chunking_for_all_columns( + std::size_t chunk_read_limit, + std::size_t pass_read_limit, + cudf::host_span const> row_group_indices, + cudf::host_span const> column_chunk_data, + parquet_reader_options const& options, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + CUDF_FUNC_RANGE(); + _impl->setup_chunking_for_all_columns( + chunk_read_limit, pass_read_limit, row_group_indices, column_chunk_data, options, stream, mr); +} + +table_with_metadata hybrid_scan_multifile::materialize_all_columns_chunk() const +{ + CUDF_FUNC_RANGE(); + return _impl->materialize_all_columns_chunk(); +} + +bool hybrid_scan_multifile::has_next_table_chunk() const { return _impl->has_next_table_chunk(); } + std::vector>> hybrid_scan_multifile::construct_row_group_passes( cudf::host_span const> row_group_indices, std::size_t pass_read_limit) const diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_preprocess.cu b/cpp/src/io/parquet/experimental/hybrid_scan_preprocess.cu index 641777b582e3..d86b62e8cbcb 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_preprocess.cu +++ b/cpp/src/io/parquet/experimental/hybrid_scan_preprocess.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,15 +9,18 @@ #include "io/parquet/reader_impl_preprocess_utils.cuh" #include "io/utilities/time_utils.hpp" +#include #include #include #include #include +#include #include #include #include +#include #include #include #include @@ -96,6 +99,13 @@ void hybrid_scan_reader_impl::prepare_row_groups( "For reading larger number of rows, please use chunked_parquet_reader.", std::overflow_error); + // Inclusive scan the number of rows per source + _file_itm_data.exclusive_sum_num_rows_per_source.resize( + _file_itm_data.num_rows_per_source.size()); + std::inclusive_scan(_file_itm_data.num_rows_per_source.cbegin(), + _file_itm_data.num_rows_per_source.cend(), + _file_itm_data.exclusive_sum_num_rows_per_source.begin()); + // check for page indexes _has_page_index = std::all_of(_file_itm_data.row_groups.cbegin(), _file_itm_data.row_groups.cend(), @@ -274,6 +284,52 @@ hybrid_scan_reader_impl::prepare_dictionaries( return {has_compressed_data, std::move(chunks), std::move(pages)}; } +namespace { + +/** + * @brief Computes the updated row mask value such that out_row_mask[i] = true, iff in_row_mask[i] + * is valid and true. This is inline with the masking behavior of cudf::apply_boolean_mask. + */ +struct row_mask_update_fn { + bool is_nullable; + bool const* in_row_mask; + bitmask_type const* in_bitmask; + + __device__ bool operator()(cudf::size_type row_idx) const + { + if (is_nullable and not bit_is_set(in_bitmask, row_idx)) { return false; } + return in_row_mask[row_idx]; + } +}; + +/** + * @brief Checks if a row is pruned (valid and false) + */ +struct is_row_pruned_fn { + bool is_nullable; + bool const* row_mask; + bitmask_type const* bitmask; + __device__ bool operator()(cudf::size_type row_idx) const + { + if (is_nullable and not bit_is_set(bitmask, row_idx)) { return false; } + return not row_mask[row_idx]; + } +}; + +} // namespace + +bool hybrid_scan_reader_impl::are_all_rows_pruned(cudf::column_view const& row_mask, + rmm::cuda_stream_view stream) const +{ + CUDF_EXPECTS(row_mask.type().id() == type_id::BOOL8, + "Input row mask column must be a boolean column"); + return cudf::detail::all_of( + cuda::counting_iterator{0}, + cuda::counting_iterator{row_mask.size()}, + is_row_pruned_fn{row_mask.nullable(), row_mask.begin(), row_mask.null_mask()}, + stream); +} + void hybrid_scan_reader_impl::update_row_mask(cudf::column_view const& in_row_mask, cudf::mutable_column_view& out_row_mask, cudf::size_type out_row_mask_offset, @@ -281,6 +337,7 @@ void hybrid_scan_reader_impl::update_row_mask(cudf::column_view const& in_row_ma { CUDF_FUNC_RANGE(); + // Total number of output row mask rows to be updated from the input auto const total_rows = static_cast(in_row_mask.size()); CUDF_EXPECTS(out_row_mask_offset + total_rows <= out_row_mask.size(), @@ -290,30 +347,23 @@ void hybrid_scan_reader_impl::update_row_mask(cudf::column_view const& in_row_ma CUDF_EXPECTS(in_row_mask.type().id() == type_id::BOOL8, "Input row mask column must be a boolean column"); - // Update output row mask such that out_row_mask[i] = true, iff in_row_mask[i] is valid and true. - // This is inline with the masking behavior of cudf::detail::apply_boolean_mask. - thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - cuda::counting_iterator{total_rows}, - out_row_mask.begin() + out_row_mask_offset, - [is_nullable = in_row_mask.nullable(), - in_row_mask = in_row_mask.begin(), - in_bitmask = in_row_mask.null_mask()] __device__(auto row_idx) { - auto const is_valid = not is_nullable or bit_is_set(in_bitmask, row_idx); - auto const is_true = in_row_mask[row_idx]; - if (is_nullable) { - return is_valid and is_true; - } else { - return is_true; - } - }); + CUDF_CUDA_TRY(cub::DeviceTransform::Transform( + cuda::counting_iterator{0}, + out_row_mask.begin() + out_row_mask_offset, + total_rows, + row_mask_update_fn{in_row_mask.nullable(), in_row_mask.begin(), in_row_mask.null_mask()}, + stream.value())); // Make sure the null mask of the output row mask column is all valid after the update. This is // to correctly assess if a payload column data page can be pruned. An invalid row in the row mask // column means the corresponding data page cannot be pruned. if (out_row_mask.nullable()) { - cudf::set_null_mask(out_row_mask.null_mask(), 0, total_rows, true, stream); - out_row_mask.set_null_count(0); + cudf::set_null_mask(out_row_mask.null_mask(), + out_row_mask_offset, + out_row_mask_offset + total_rows, + true, + stream); + out_row_mask.set_null_count(out_row_mask.null_count(0, out_row_mask.size(), stream)); } } diff --git a/cpp/src/io/parquet/experimental/page_index_filter.cu b/cpp/src/io/parquet/experimental/page_index_filter.cu index 1cf4995bf66a..124b780a86ab 100644 --- a/cpp/src/io/parquet/experimental/page_index_filter.cu +++ b/cpp/src/io/parquet/experimental/page_index_filter.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -1010,7 +1010,7 @@ thrust::host_vector aggregate_reader_metadata::compute_data_page_mask( // Return early if page index is not present if (not has_page_index) { CUDF_LOG_WARN("Encountered missing Parquet page index for one or more output columns"); - return thrust::host_vector{}; + return thrust::host_vector(0, stream); } // Collect column schema indices from the input columns. diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 37bfa47b236b..1dd0d45f6214 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -348,9 +348,10 @@ ConfigureTest( ) ConfigureTest( HYBRID_SCAN_TEST + io/experimental/hybrid_scan_common.cpp io/experimental/hybrid_scan_composer.cpp - io/experimental/hybrid_scan_multifile_composer.cpp io/experimental/hybrid_scan_filters_test.cpp + io/experimental/hybrid_scan_multifile_composer.cpp io/experimental/hybrid_scan_multifile_filters_test.cpp io/experimental/hybrid_scan_multifile_test.cpp io/experimental/hybrid_scan_test.cpp diff --git a/cpp/tests/io/experimental/hybrid_scan_common.cpp b/cpp/tests/io/experimental/hybrid_scan_common.cpp new file mode 100644 index 000000000000..99ee37e6c333 --- /dev/null +++ b/cpp/tests/io/experimental/hybrid_scan_common.cpp @@ -0,0 +1,365 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "hybrid_scan_common.hpp" + +#include "tests/io/parquet_common.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include + +namespace { + +/** + * @brief Creates a strings column with a constant stringified value between 0 and 9999 + * + * @param value String value between 0 and 9999 + * @return Strings column wrapper + */ +cudf::test::strings_column_wrapper constant_strings(cudf::size_type value) +{ + CUDF_EXPECTS(value >= 0 && value <= 9999, "String value must be between 0000 and 9999"); + + auto elements = thrust::make_transform_iterator(cuda::make_constant_iterator(value), + [](auto i) { return std::format("{:04d}", i); }); + return cudf::test::strings_column_wrapper(elements, elements + num_ordered_rows); +} + +/** + * @brief Fail for types other than duration or timestamp + */ +template ())> +cudf::test::fixed_width_column_wrapper descending_low_cardinality() +{ + static_assert( + cudf::is_chrono(), + "Use testdata::descending() to generate descending values for non-temporal types"); +} + +/** + * @brief Creates a duration column wrapper with low cardinality descending values + * + * @tparam T Duration type + * @return Column wrapper + */ +template ())> +cudf::test::fixed_width_column_wrapper descending_low_cardinality() +{ + auto elements = cudf::detail::make_counting_transform_iterator( + 0, [](auto i) { return T((num_ordered_rows - i) / 100); }); + return cudf::test::fixed_width_column_wrapper(elements, elements + num_ordered_rows); +} + +/** + * @brief Creates a timestamp column wrapper with low cardinality descending values + * + * @tparam T Timestamp type + * @return Column wrapper + */ +template ())> +cudf::test::fixed_width_column_wrapper descending_low_cardinality() +{ + auto elements = cudf::detail::make_counting_transform_iterator( + 0, [](auto i) { return T(typename T::duration((num_ordered_rows - i) / 100)); }); + return cudf::test::fixed_width_column_wrapper(elements, elements + num_ordered_rows); +} + +} // namespace + +std::unique_ptr make_list_str_column(std::mt19937& gen, + bool is_str_nullable, + bool is_list_nullable) +{ + auto constexpr num_rows = num_ordered_rows; + auto constexpr string_per_row = 3; + auto constexpr num_string_rows = num_rows * string_per_row; + + std::vector strings{ + "abc", "x", "bananas", "gpu", "minty", "backspace", "", "cayenne", "turbine", "soft"}; + std::uniform_int_distribution uni(0, strings.size() - 1); + auto string_iter = cudf::detail::make_counting_transform_iterator( + 0, [&](cudf::size_type idx) { return strings[uni(gen)]; }); + + std::bernoulli_distribution bn(0.7f); + auto string_valids = cudf::detail::make_counting_transform_iterator( + 0, [&](int index) { return is_str_nullable ? bn(gen) : true; }); + cudf::test::strings_column_wrapper string_col{ + string_iter, string_iter + num_string_rows, string_valids}; + + auto offset_iter = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type idx) { return idx * string_per_row; }); + cudf::test::fixed_width_column_wrapper offsets(offset_iter, + offset_iter + num_rows + 1); + + auto list_valids = + cudf::detail::make_counting_transform_iterator(0, [&](int index) { return index % 100; }); + auto [null_mask, null_count] = [&]() { + if (is_list_nullable) { + return cudf::test::detail::make_null_mask(list_valids, list_valids + num_rows); + } else { + return std::make_pair(rmm::device_buffer{}, 0); + } + }(); + return cudf::make_lists_column( + num_rows, offsets.release(), string_col.release(), null_count, std::move(null_mask)); +} + +multifile_inputs::multifile_inputs(cudf::io::source_info const& source_info) + : datasources{cudf::io::make_datasources(source_info)} +{ + datasource_refs.reserve(datasources.size()); + footer_buffers.reserve(datasources.size()); + footer_byte_spans.reserve(datasources.size()); + + for (auto const& datasource : datasources) { + datasource_refs.emplace_back(*datasource); + footer_buffers.emplace_back(cudf::io::parquet::fetch_footer_to_host(datasource_refs.back())); + footer_byte_spans.emplace_back(*footer_buffers.back()); + } +} + +cudf::io::source_info build_source_info(std::vector> const& file_buffers) +{ + std::vector> spans; + spans.reserve(file_buffers.size()); + for (auto const& buf : file_buffers) { + spans.emplace_back(buf.data(), buf.size()); + } + return cudf::io::source_info(cudf::host_span>{spans}); +} + +void setup_page_indexes(cudf::io::parquet::experimental::hybrid_scan_multifile const& reader, + multifile_inputs const& inputs) +{ + auto const page_index_byte_ranges = reader.page_index_byte_ranges(); + std::vector> page_index_byte_spans; + page_index_byte_spans.reserve(page_index_byte_ranges.size()); + + auto const page_index_buffers = cudf::io::parquet::fetch_page_indexes_to_host( + cudf::host_span const>{inputs.datasource_refs}, + cudf::host_span{page_index_byte_ranges}); + std::transform(page_index_buffers.begin(), + page_index_buffers.end(), + std::back_inserter(page_index_byte_spans), + [](auto const& buffer) { return cudf::host_span{*buffer}; }); + + reader.setup_page_indexes( + cudf::host_span const>{page_index_byte_spans}); +} + +std::vector> group_byte_ranges_by_source( + std::pair, std::vector> const& + byte_ranges_and_source_map, + std::size_t num_sources) +{ + auto const& [byte_ranges, source_map] = byte_ranges_and_source_map; + CUDF_EXPECTS(byte_ranges.size() == source_map.size(), "Invalid source map size"); + + auto byte_ranges_per_source = + std::vector>(num_sources); + std::for_each(byte_ranges.begin(), + byte_ranges.end(), + [&, range_index = std::size_t{0}](auto const& range) mutable { + auto const source_index = source_map[range_index++]; + CUDF_EXPECTS(source_index >= 0 and static_cast(source_index) < + byte_ranges_per_source.size(), + "Invalid byte range source index"); + byte_ranges_per_source[source_index].push_back(range); + }); + return byte_ranges_per_source; +} + +multisource_device_data fetch_multisource_device_data( + multifile_inputs const& inputs, + std::pair, std::vector> const& + byte_ranges_and_source_map, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto const byte_ranges_per_source = + group_byte_ranges_by_source(byte_ranges_and_source_map, inputs.datasources.size()); + auto [buffers, per_source_spans, tasks] = cudf::io::parquet::fetch_byte_ranges_to_device_async( + inputs.datasource_refs, + cudf::host_span const>{byte_ranges_per_source}, + stream, + mr); + tasks.get(); + + auto flat_spans = std::vector>{}; + for (auto const& source_spans : per_source_spans) { + flat_spans.insert(flat_spans.end(), source_spans.begin(), source_spans.end()); + } + + return {std::move(buffers), std::move(per_source_spans), std::move(flat_spans)}; +} + +std::unique_ptr concatenate_tables(std::vector>&& tables, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + if (tables.size() == 1) { return std::move(tables[0]); } + + auto table_views = std::vector{}; + table_views.reserve(tables.size()); + std::transform( + tables.begin(), tables.end(), std::back_inserter(table_views), [](auto const& tbl) { + return tbl->view(); + }); + return cudf::concatenate(table_views, stream, mr); +} + +template +std::pair, std::vector> create_parquet_with_stats( + cudf::size_type str_col_value, + cudf::io::compression_type compression, + rmm::cuda_stream_view stream) +{ + static_assert(NumTableConcats >= 1, "Concatenated table must contain at least one table"); + + auto col0 = testdata::ascending(); + auto col1 = []() { + if constexpr (cudf::is_chrono()) { + return descending_low_cardinality(); + } else { + return testdata::descending(); + } + }(); + + auto col2 = [&]() { + if constexpr (IsConstantStrings) { + return constant_strings(str_col_value); // constant stringified value + } else { + return testdata::ascending(); // ascending strings + } + }(); + + // Output table view + auto output = table_view{{col0, col1, col2}}; + + // Add nullmasks to the columns if specified + std::vector> columns; + if constexpr (IsNullable) { + std::mt19937 gen(0xc0ffee); + std::bernoulli_distribution bn(0.7f); + auto valids = + cudf::detail::make_counting_transform_iterator(0, [&](int index) { return bn(gen); }); + auto const num_rows = static_cast(col0).size(); + + auto const make_null_mask = [stream](auto begin, auto end) { + auto [null_mask, null_count] = cudf::test::detail::make_null_mask_vector(begin, end); + auto d_mask = rmm::device_buffer{ + null_mask.data(), cudf::bitmask_allocation_size_bytes(cudf::distance(begin, end)), stream}; + return std::pair{std::move(d_mask), null_count}; + }; + + columns.emplace_back(col0.release()); + auto [nullmask, nullcount] = make_null_mask(valids, valids + num_rows); + columns.back()->set_null_mask(std::move(nullmask), nullcount); + + columns.emplace_back(col1.release()); + std::tie(nullmask, nullcount) = make_null_mask(valids + num_rows, valids + 2 * num_rows); + columns.back()->set_null_mask(std::move(nullmask), nullcount); + + columns.emplace_back(col2.release()); + std::tie(nullmask, nullcount) = make_null_mask(valids + 2 * num_rows, valids + 3 * num_rows); + columns.back()->set_null_mask(std::move(nullmask), nullcount); + + // Purge non-empty nulls from the strings column only + columns.back() = cudf::purge_nonempty_nulls(columns.back()->view(), stream); + + // Update the output table view with the nullable columns + output = table_view{{columns[0]->view(), columns[1]->view(), columns[2]->view()}}; + } + + auto table = cudf::concatenate(std::vector(NumTableConcats, output), stream); + output = table->view(); + cudf::io::table_input_metadata output_metadata(output); + output_metadata.column_metadata[0].set_name("col0"); + output_metadata.column_metadata[1].set_name("col1"); + output_metadata.column_metadata[2].set_name("col2"); + + std::vector buffer; + cudf::io::parquet_writer_options out_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, output) + .metadata(std::move(output_metadata)) + .row_group_size_rows(page_size_for_ordered_tests) + .max_page_size_rows(page_size_for_ordered_tests / 5) + .compression(compression) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN); + + if constexpr (NumTableConcats > 1) { + out_opts.set_row_group_size_rows(num_ordered_rows); + out_opts.set_max_page_size_rows(page_size_for_ordered_tests); + } + + cudf::io::write_parquet(out_opts, stream); + + return std::pair{std::move(table), std::move(buffer)}; +} + +#define INSTANTIATE_CREATE_PARQUET_WITH_STATS(T, NUM_CONCATS, CONSTANT_STRINGS, NULLABLE) \ + template std::pair, std::vector> \ + create_parquet_with_stats( \ + cudf::size_type, cudf::io::compression_type, rmm::cuda_stream_view) + +#define INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(T) \ + INSTANTIATE_CREATE_PARQUET_WITH_STATS(T, 1, true, false); \ + INSTANTIATE_CREATE_PARQUET_WITH_STATS(T, 1, true, true) + +INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint32_t, 4, true, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(cudf::timestamp_ms, 2, true, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(cudf::duration_ms, 2, true, false); + +INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint8_t, 1, false, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint16_t, 1, false, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint32_t, 1, false, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint64_t, 1, false, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(cudf::duration_ms, 1, false, false); + +INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint8_t, 2, false, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint16_t, 2, false, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint32_t, 2, false, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint64_t, 2, false, false); + +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(int8_t); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(int16_t); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(int32_t); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(int64_t); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(uint8_t); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(uint16_t); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(uint32_t); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(uint64_t); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(float); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(double); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(cudf::timestamp_D); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(cudf::timestamp_ms); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(cudf::timestamp_us); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(cudf::timestamp_ns); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(cudf::duration_ms); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(cudf::duration_us); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(cudf::duration_ns); +INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT(cudf::string_view); + +#undef INSTANTIATE_CREATE_PARQUET_WITH_STATS_DICT +#undef INSTANTIATE_CREATE_PARQUET_WITH_STATS diff --git a/cpp/tests/io/experimental/hybrid_scan_common.hpp b/cpp/tests/io/experimental/hybrid_scan_common.hpp index c31b49a7bd97..583ef520ccd8 100644 --- a/cpp/tests/io/experimental/hybrid_scan_common.hpp +++ b/cpp/tests/io/experimental/hybrid_scan_common.hpp @@ -1,47 +1,32 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include "tests/io/parquet_common.hpp" - -#include - -#include -#include +#include +#include +#include +#include +#include #include #include #include -#include +#include +#include #include +#include -#include - -#include +#include +#include +#include +#include #include -#include #include #include -/** - * @brief Creates a strings column with a constant stringified value between 0 and 9999 - * - * @param value String value between 0 and 9999 - * @return Strings column wrapper - */ - -cudf::test::strings_column_wrapper inline constant_strings(cudf::size_type value) -{ - CUDF_EXPECTS(value >= 0 && value <= 9999, "String value must be between 0000 and 9999"); - - auto elements = thrust::make_transform_iterator(cuda::make_constant_iterator(value), - [](auto i) { return std::format("{:04d}", i); }); - return cudf::test::strings_column_wrapper(elements, elements + num_ordered_rows); -} - /** * @brief Helper to construct a random list column * @@ -51,80 +36,77 @@ cudf::test::strings_column_wrapper inline constant_strings(cudf::size_type value * * @return Unique pointer to the constructed list column */ -inline auto make_list_str_column(std::mt19937& gen, bool is_str_nullable, bool is_list_nullable) -{ - auto constexpr num_rows = num_ordered_rows; - auto constexpr string_per_row = 3; - auto constexpr num_string_rows = num_rows * string_per_row; +[[nodiscard]] std::unique_ptr make_list_str_column(std::mt19937& gen, + bool is_str_nullable, + bool is_list_nullable); - std::vector strings{ - "abc", "x", "bananas", "gpu", "minty", "backspace", "", "cayenne", "turbine", "soft"}; - std::uniform_int_distribution uni(0, strings.size() - 1); - auto string_iter = cudf::detail::make_counting_transform_iterator( - 0, [&](cudf::size_type idx) { return strings[uni(gen)]; }); +/** + * @brief Struct to hold multifile datasources and footer buffers along with their byte spans + */ +struct multifile_inputs { + /** + * @brief Construct datasources, datasource refs, and footer byte spans from source info + */ + explicit multifile_inputs(cudf::io::source_info const& source_info); + + std::vector> datasources; + std::vector> datasource_refs; + std::vector> footer_buffers; + std::vector> footer_byte_spans; +}; - std::bernoulli_distribution bn(0.7f); - auto string_valids = cudf::detail::make_counting_transform_iterator( - 0, [&](int index) { return is_str_nullable ? bn(gen) : true; }); - cudf::test::strings_column_wrapper string_col{ - string_iter, string_iter + num_string_rows, string_valids}; +/** + * @brief Device buffers and spans from multiple input sources + */ +struct multisource_device_data { + std::vector buffers; + std::vector>> per_source_spans; + std::vector> flat_spans; +}; - auto offset_iter = cudf::detail::make_counting_transform_iterator( - 0, [](cudf::size_type idx) { return idx * string_per_row; }); - cudf::test::fixed_width_column_wrapper offsets(offset_iter, - offset_iter + num_rows + 1); +/** + * @brief Construct source info from host buffers + */ +[[nodiscard]] cudf::io::source_info build_source_info( + std::vector> const& file_buffers); - auto list_valids = - cudf::detail::make_counting_transform_iterator(0, [&](int index) { return index % 100; }); - auto [null_mask, null_count] = [&]() { - if (is_list_nullable) { - return cudf::test::detail::make_null_mask(list_valids, list_valids + num_rows); - } else { - return std::make_pair(rmm::device_buffer{}, 0); - } - }(); - return cudf::make_lists_column( - num_rows, offsets.release(), string_col.release(), null_count, std::move(null_mask)); -} +/** + * @brief Fetch and set up page indexes for all sources in a multifile reader + */ +void setup_page_indexes(cudf::io::parquet::experimental::hybrid_scan_multifile const& reader, + multifile_inputs const& inputs); /** - * @brief Fail for types other than duration or timestamp + * @brief Groups a flat byte range list by source using the specified source map */ -template ())> -cudf::test::fixed_width_column_wrapper descending_low_cardinality() -{ - static_assert( - cudf::is_chrono(), - "Use testdata::descending() to generate descending values for non-temporal types"); -} +[[nodiscard]] std::vector> group_byte_ranges_by_source( + std::pair, std::vector> const& + byte_ranges_and_source_map, + std::size_t num_sources); /** - * @brief Creates a duration column wrapper with low cardinality descending values - * - * @tparam T Duration type - * @return Column wrapper + * @brief Fetches byte ranges from multiple sources and returns per-source and flattened spans */ -template ())> -cudf::test::fixed_width_column_wrapper descending_low_cardinality() -{ - auto elements = cudf::detail::make_counting_transform_iterator( - 0, [](auto i) { return T((num_ordered_rows - i) / 100); }); - return cudf::test::fixed_width_column_wrapper(elements, elements + num_ordered_rows); -} +[[nodiscard]] multisource_device_data fetch_multisource_device_data( + multifile_inputs const& inputs, + std::pair, std::vector> const& + byte_ranges_and_source_map, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); /** - * @brief Creates a timestamp column wrapper with low cardinality descending values + * @brief Concatenate a vector of tables and return the resultant table + * + * @param tables Vector of tables to concatenate + * @param stream CUDA stream to use + * @param mr Device memory resource used to allocate the returned table's device memory * - * @tparam T Timestamp type - * @return Column wrapper + * @return Unique pointer to the resultant concatenated table */ -template ())> -cudf::test::fixed_width_column_wrapper descending_low_cardinality() -{ - auto elements = cudf::detail::make_counting_transform_iterator( - 0, [](auto i) { return T(typename T::duration((num_ordered_rows - i) / 100)); }); - return cudf::test::fixed_width_column_wrapper(elements, elements + num_ordered_rows); -} +[[nodiscard]] std::unique_ptr concatenate_tables( + std::vector>&& tables, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); /** * @brief Creates a table and writes it to Parquet host buffer with column level statistics @@ -152,86 +134,7 @@ template -auto create_parquet_with_stats( +[[nodiscard]] std::pair, std::vector> create_parquet_with_stats( cudf::size_type str_col_value = 100, cudf::io::compression_type compression = cudf::io::compression_type::AUTO, - rmm::cuda_stream_view stream = cudf::get_default_stream()) -{ - static_assert(NumTableConcats >= 1, "Concatenated table must contain at least one table"); - - auto col0 = testdata::ascending(); - auto col1 = []() { - if constexpr (cudf::is_chrono()) { - return descending_low_cardinality(); - } else { - return testdata::descending(); - } - }(); - - auto col2 = [&]() { - if constexpr (IsConstantStrings) { - return constant_strings(str_col_value); // constant stringified value - } else { - return testdata::ascending(); // ascending strings - } - }(); - - // Output table view - auto output = table_view{{col0, col1, col2}}; - - // Add nullmasks to the columns if specified - std::vector> columns; - if constexpr (IsNullable) { - std::mt19937 gen(0xc0ffee); - std::bernoulli_distribution bn(0.7f); - auto valids = - cudf::detail::make_counting_transform_iterator(0, [&](int index) { return bn(gen); }); - auto const num_rows = static_cast(col0).size(); - - columns.emplace_back(col0.release()); - auto [nullmask, nullcount] = cudf::test::detail::make_null_mask(valids, valids + num_rows); - columns.back()->set_null_mask(std::move(nullmask), nullcount); - - columns.emplace_back(col1.release()); - std::tie(nullmask, nullcount) = - cudf::test::detail::make_null_mask(valids + num_rows, valids + 2 * num_rows); - columns.back()->set_null_mask(std::move(nullmask), nullcount); - - columns.emplace_back(col2.release()); - std::tie(nullmask, nullcount) = - cudf::test::detail::make_null_mask(valids + 2 * num_rows, valids + 3 * num_rows); - columns.back()->set_null_mask(std::move(nullmask), nullcount); - - // Purge non-empty nulls from the strings column only - cudf::purge_nonempty_nulls(columns.back()->view()); - - // Update the output table view with the nullable columns - output = table_view{{columns[0]->view(), columns[1]->view(), columns[2]->view()}}; - } - - auto table = cudf::concatenate(std::vector(NumTableConcats, output)); - output = table->view(); - cudf::io::table_input_metadata output_metadata(output); - output_metadata.column_metadata[0].set_name("col0"); - output_metadata.column_metadata[1].set_name("col1"); - output_metadata.column_metadata[2].set_name("col2"); - - std::vector buffer; - cudf::io::parquet_writer_options out_opts = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, output) - .metadata(std::move(output_metadata)) - .row_group_size_rows(page_size_for_ordered_tests) - .max_page_size_rows(page_size_for_ordered_tests / 5) - .compression(compression) - .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) - .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN); - - if constexpr (NumTableConcats > 1) { - out_opts.set_row_group_size_rows(num_ordered_rows); - out_opts.set_max_page_size_rows(page_size_for_ordered_tests); - } - - cudf::io::write_parquet(out_opts); - - return std::pair{std::move(table), std::move(buffer)}; -} + rmm::cuda_stream_view stream = cudf::get_default_stream()); diff --git a/cpp/tests/io/experimental/hybrid_scan_composer.cpp b/cpp/tests/io/experimental/hybrid_scan_composer.cpp index 5365988220d9..57a6f6cb7e3f 100644 --- a/cpp/tests/io/experimental/hybrid_scan_composer.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_composer.cpp @@ -1,12 +1,13 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "hybrid_scan_composer.hpp" +#include "hybrid_scan_common.hpp" + #include -#include #include #include #include @@ -125,30 +126,6 @@ auto apply_hybrid_scan_filters(cudf::io::datasource& datasource, current_row_group_indices.end()); } -/* - * @brief Concatenate a vector of tables and return the resultant table - * - * @param tables Vector of tables to concatenate - * @param stream CUDA stream to use - * - * @return Unique pointer to the resultant concatenated table. - */ -std::unique_ptr concatenate_tables(std::vector> tables, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) -{ - if (tables.size() == 1) { return std::move(tables[0]); } - - std::vector table_views; - table_views.reserve(tables.size()); - std::transform( - tables.begin(), tables.end(), std::back_inserter(table_views), [&](auto const& tbl) { - return tbl->view(); - }); - // Construct the final table - return cudf::concatenate(table_views, stream, mr); -} - } // namespace std::tuple, std::unique_ptr> hybrid_scan( diff --git a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp index 10960c185ef9..4a76bd507dd1 100644 --- a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp @@ -1,9 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "hybrid_scan_common.hpp" +#include "tests/io/parquet_common.hpp" #include diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_common.hpp b/cpp/tests/io/experimental/hybrid_scan_multifile_common.hpp deleted file mode 100644 index 2d033539324e..000000000000 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_common.hpp +++ /dev/null @@ -1,79 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include -#include -#include -#include - -#include -#include -#include -#include - -/** - * @brief Struct to hold multifile datasources and footer buffers along with their byte spans - */ -struct multifile_inputs { - /** - * @brief Construct datasources, datasource refs, and footer byte spans from source info - */ - explicit multifile_inputs(cudf::io::source_info const& source_info) - : datasources{cudf::io::make_datasources(source_info)} - { - datasource_refs.reserve(datasources.size()); - footer_buffers.reserve(datasources.size()); - footer_byte_spans.reserve(datasources.size()); - - for (auto const& datasource : datasources) { - datasource_refs.emplace_back(*datasource); - footer_buffers.emplace_back(cudf::io::parquet::fetch_footer_to_host(datasource_refs.back())); - footer_byte_spans.emplace_back(*footer_buffers.back()); - } - } - - std::vector> datasources; - std::vector> datasource_refs; - std::vector> footer_buffers; - std::vector> footer_byte_spans; -}; - -/** - * @brief Construct source info from host buffers - */ -template -cudf::io::source_info build_source_info(Buffers const& file_buffers) -{ - std::vector> spans; - spans.reserve(file_buffers.size()); - for (auto const& buf : file_buffers) { - spans.emplace_back(buf.data(), buf.size()); - } - return cudf::io::source_info(cudf::host_span>{spans}); -} - -/** - * @brief Fetch and set up page indexes for all sources in a multifile reader - */ -inline void setup_page_indexes(cudf::io::parquet::experimental::hybrid_scan_multifile const& reader, - multifile_inputs const& inputs) -{ - auto const page_index_byte_ranges = reader.page_index_byte_ranges(); - std::vector> page_index_byte_spans; - page_index_byte_spans.reserve(page_index_byte_ranges.size()); - - auto const page_index_buffers = cudf::io::parquet::fetch_page_indexes_to_host( - cudf::host_span const>{inputs.datasource_refs}, - cudf::host_span{page_index_byte_ranges}); - std::transform(page_index_buffers.begin(), - page_index_buffers.end(), - std::back_inserter(page_index_byte_spans), - [](auto const& buffer) { return cudf::host_span{*buffer}; }); - - reader.setup_page_indexes( - cudf::host_span const>{page_index_byte_spans}); -} diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_composer.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_composer.cpp index 0c23e35b47c7..d3d52c2c5129 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_composer.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_composer.cpp @@ -5,88 +5,17 @@ #include "hybrid_scan_multifile_composer.hpp" -#include "hybrid_scan_multifile_common.hpp" +#include "hybrid_scan_common.hpp" #include #include -#include -#include -#include -#include - -#include #include #include #include using cudf::io::parquet::experimental::use_data_page_mask; -namespace { - -/** - * @brief Device buffers and spans from multiple input sources - */ -struct multisource_device_data { - std::vector buffers; - std::vector>> per_source_spans; - std::vector> flat_spans; -}; - -/** - * @brief Groups a flat byte range list by source using the specified source map - */ -std::vector> group_byte_ranges_by_source( - std::pair, std::vector> const& - byte_ranges_and_source_map, - std::size_t num_sources) -{ - auto const& [byte_ranges, source_map] = byte_ranges_and_source_map; - CUDF_EXPECTS(byte_ranges.size() == source_map.size(), "Invalid source map size"); - - auto byte_ranges_per_source = - std::vector>(num_sources); - std::for_each(byte_ranges.begin(), - byte_ranges.end(), - [&, range_index = std::size_t{0}](auto const& range) mutable { - auto const source_index = source_map[range_index++]; - CUDF_EXPECTS(source_index >= 0 and static_cast(source_index) < - byte_ranges_per_source.size(), - "Invalid byte range source index"); - byte_ranges_per_source[source_index].push_back(range); - }); - return byte_ranges_per_source; -} - -/** - * @brief Fetches byte ranges from multiple sources and returns per-source and flattened spans - */ -multisource_device_data fetch_multisource_device_data( - multifile_inputs const& inputs, - std::pair, std::vector> const& - byte_ranges_and_source_map, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) -{ - auto const byte_ranges_per_source = - group_byte_ranges_by_source(byte_ranges_and_source_map, inputs.datasources.size()); - auto [buffers, per_source_spans, tasks] = cudf::io::parquet::fetch_byte_ranges_to_device_async( - inputs.datasource_refs, - cudf::host_span const>{byte_ranges_per_source}, - stream, - mr); - tasks.get(); - - auto flat_spans = std::vector>{}; - for (auto const& source_spans : per_source_spans) { - flat_spans.insert(flat_spans.end(), source_spans.begin(), source_spans.end()); - } - - return {std::move(buffers), std::move(per_source_spans), std::move(flat_spans)}; -} - -} // namespace - std::tuple, std::unique_ptr> hybrid_scan_multifile( cudf::io::source_info const& source_info, cudf::ast::operation const& filter_expression, @@ -162,3 +91,112 @@ std::unique_ptr hybrid_scan_multifile_single_step( .materialize_all_columns(row_groups, all_column_chunks.flat_spans, options, stream, mr) .tbl; } + +std::tuple, std::unique_ptr> +chunked_hybrid_scan_multifile(cudf::io::source_info const& source_info, + cudf::ast::operation const& filter_expression, + std::optional> const& payload_column_names, + bool case_sensitive_names, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto options = cudf::io::parquet_reader_options::builder() + .filter(filter_expression) + .case_sensitive_names(case_sensitive_names) + .build(); + if (payload_column_names.has_value()) { options.set_column_names(payload_column_names.value()); } + + auto inputs = multifile_inputs(source_info); + auto reader = + cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; + setup_page_indexes(reader, inputs); + + auto const input_row_groups = reader.all_row_groups(options); + auto const row_groups = reader.filter_row_groups_with_stats(input_row_groups, options, stream); + auto row_mask = reader.build_row_mask_with_page_index_stats(row_groups, options, stream, mr); + + // Non-trivial chunk and pass read limits so chunking and multi-pass reads trigger across the + // single row mask spanning all sources + auto constexpr chunk_read_limit = std::size_t{256 * 1024}; + auto constexpr pass_read_limit = std::size_t{1024 * 1024}; + + auto filter_tables = std::vector>{}; + auto payload_tables = std::vector>{}; + + auto filter_column_chunks = fetch_multisource_device_data( + inputs, reader.filter_column_chunks_byte_ranges(row_groups, options), stream, mr); + auto row_mask_view = row_mask->mutable_view(); + reader.setup_chunking_for_filter_columns(chunk_read_limit, + pass_read_limit, + row_groups, + row_mask_view, + use_data_page_mask::YES, + filter_column_chunks.flat_spans, + options, + stream, + mr); + while (reader.has_next_table_chunk()) { + filter_tables.push_back(reader.materialize_filter_columns_chunk(row_mask_view).tbl); + } + + auto payload_column_chunks = fetch_multisource_device_data( + inputs, reader.payload_column_chunks_byte_ranges(row_groups, options), stream, mr); + reader.setup_chunking_for_payload_columns(chunk_read_limit, + pass_read_limit, + row_groups, + row_mask_view, + use_data_page_mask::YES, + payload_column_chunks.flat_spans, + options, + stream, + mr); + while (reader.has_next_table_chunk()) { + payload_tables.push_back(reader.materialize_payload_columns_chunk(row_mask_view).tbl); + } + + return std::tuple{concatenate_tables(std::move(filter_tables), stream, mr), + concatenate_tables(std::move(payload_tables), stream, mr)}; +} + +std::unique_ptr chunked_hybrid_scan_multifile_single_step( + cudf::io::source_info const& source_info, + cudf::ast::operation const& filter_expression, + std::optional> const& column_names, + bool case_sensitive_names, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto options = cudf::io::parquet_reader_options::builder() + .filter(filter_expression) + .case_sensitive_names(case_sensitive_names) + .build(); + if (column_names.has_value()) { options.set_column_names(column_names.value()); } + + auto inputs = multifile_inputs(source_info); + auto reader = + cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; + setup_page_indexes(reader, inputs); + + auto const input_row_groups = reader.all_row_groups(options); + auto const row_groups = reader.filter_row_groups_with_stats(input_row_groups, options, stream); + + auto constexpr chunk_read_limit = std::size_t{256 * 1024}; + auto constexpr pass_read_limit = std::size_t{1024 * 1024}; + + auto all_column_chunks = fetch_multisource_device_data( + inputs, reader.all_column_chunks_byte_ranges(row_groups, options), stream, mr); + reader.setup_chunking_for_all_columns(chunk_read_limit, + pass_read_limit, + row_groups, + all_column_chunks.flat_spans, + options, + stream, + mr); + + auto tables = std::vector>{}; + while (reader.has_next_table_chunk()) { + tables.push_back(reader.materialize_all_columns_chunk().tbl); + } + + return concatenate_tables(std::move(tables), stream, mr); +} diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_composer.hpp b/cpp/tests/io/experimental/hybrid_scan_multifile_composer.hpp index ccd0716a11bc..50421f6f5ac2 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_composer.hpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_composer.hpp @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -54,3 +55,47 @@ std::unique_ptr hybrid_scan_multifile_single_step( bool case_sensitive_names, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); + +/** + * @brief Read parquet sources with the hybrid scan multifile reader using chunked materialization + * + * Splits the filtered row groups into read passes via `construct_row_group_passes` and chunks the + * filter and payload column materialization within each pass. + * + * @param source_info Input source info containing one or more Parquet sources + * @param filter_expression Filter expression + * @param payload_column_names List of paths of select payload column names, if any + * @param case_sensitive_names Whether column names are case sensitive + * @param stream CUDA stream for hybrid scan reader + * @param mr Device memory resource + * + * @return Tuple of filter and payload tables + */ +std::tuple, std::unique_ptr> +chunked_hybrid_scan_multifile(cudf::io::source_info const& source_info, + cudf::ast::operation const& filter_expression, + std::optional> const& payload_column_names, + bool case_sensitive_names, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/** + * @brief Read parquet sources with the hybrid scan multifile reader in a single step using chunked + * materialization + * + * @param source_info Input source info containing one or more Parquet sources + * @param filter_expression Filter expression + * @param column_names List of column names to read, if any + * @param case_sensitive_names Whether column names are case sensitive + * @param stream CUDA stream + * @param mr Device memory resource + * + * @return Materialized table + */ +std::unique_ptr chunked_hybrid_scan_multifile_single_step( + cudf::io::source_info const& source_info, + cudf::ast::operation const& filter_expression, + std::optional> const& column_names, + bool case_sensitive_names, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp index 37bf9ba4c8fa..6183c375ce7e 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp @@ -1,10 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "hybrid_scan_common.hpp" -#include "hybrid_scan_multifile_common.hpp" +#include "tests/io/parquet_common.hpp" #include @@ -22,7 +22,6 @@ #include #include -#include #include #include #include diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp index 3cb66aec6329..65b04eedfdd5 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp @@ -4,8 +4,8 @@ */ #include "hybrid_scan_common.hpp" -#include "hybrid_scan_multifile_common.hpp" #include "hybrid_scan_multifile_composer.hpp" +#include "tests/io/parquet_common.hpp" #include #include @@ -46,7 +46,8 @@ namespace { */ template void test_hybrid_scan_multifile(std::vector const& columns, - bool case_sensitive_names = true) + bool case_sensitive_names = true, + uint32_t literal_value = 100) { auto const table = cudf::table_view{columns}; cudf::io::table_input_metadata expected_metadata(table); @@ -65,9 +66,9 @@ void test_hybrid_scan_multifile(std::vector const& columns, cudf::io::write_parquet(out_opts); } - auto literal_value = cudf::numeric_scalar(100); - auto literal = cudf::ast::literal(literal_value); - auto col_ref_0 = cudf::ast::column_name_reference(case_sensitive_names ? "col0" : "Col0"); + auto scalar = cudf::numeric_scalar(literal_value); + auto literal = cudf::ast::literal(scalar); + auto col_ref_0 = cudf::ast::column_name_reference(case_sensitive_names ? "col0" : "CoL0"); auto filter_expression = cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref_0, literal); @@ -87,19 +88,43 @@ void test_hybrid_scan_multifile(std::vector const& columns, auto const all_table = hybrid_scan_multifile_single_step( source_info, filter_expression, {}, case_sensitive_names, stream, mr); + auto const [chunked_filter_table, chunked_payload_table] = chunked_hybrid_scan_multifile( + source_info, filter_expression, {}, case_sensitive_names, stream, mr); + + auto const chunked_all_table = chunked_hybrid_scan_multifile_single_step( + source_info, filter_expression, {}, case_sensitive_names, stream, mr); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->select({0}), filter_table->view()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->select({0}), chunked_filter_table->view()); auto payload_column_indices = std::vector(columns.size() - 1); std::iota(payload_column_indices.begin(), payload_column_indices.end(), 1); CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->select(payload_column_indices), payload_table->view()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->select(payload_column_indices), + chunked_payload_table->view()); CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->view(), all_table->view()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->view(), chunked_all_table->view()); } } // namespace struct HybridScanMultifileTest : public cudf::test::BaseFixture {}; +TEST_F(HybridScanMultifileTest, EmptyResult) +{ + std::mt19937 gen(0xc0c0a); + + auto col0 = testdata::ascending(); + auto col1 = make_list_str_column(gen, false, false); + auto col2 = make_list_str_column(gen, false, true); + auto col3 = make_list_str_column(gen, true, false); + auto col4 = make_list_str_column(gen, true, true); + + auto constexpr literal_value = uint32_t(num_ordered_rows); + test_hybrid_scan_multifile({col0, *col1, *col2, *col3, *col4}, false, literal_value); +} + TEST_F(HybridScanMultifileTest, MaterializeLists) { std::mt19937 gen(0xadd); diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index 2ecb1ebfc1f7..f7cd45b926c3 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -1,10 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "hybrid_scan_common.hpp" #include "hybrid_scan_composer.hpp" +#include "tests/io/parquet_common.hpp" #include #include