diff --git a/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp index 70ff792660bf..8da2f535e961 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp @@ -123,6 +123,50 @@ class hybrid_scan_multifile { */ void reset_column_selection() const; + /** + * @brief Filter the row groups using the byte range specified by [`bytes_to_skip`, + * `bytes_to_skip + bytes_to_read`) + * + * 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 options Parquet reader options + * @return Filtered per-source row group indices (one inner vector per source) + */ + [[nodiscard]] std::vector> filter_row_groups_with_byte_range( + cudf::host_span const> row_group_indices, + parquet_reader_options const& options) const; + + /** + * @brief Filter the input row groups using column chunk statistics + * + * @param row_group_indices 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 + */ + [[nodiscard]] std::vector> filter_row_groups_with_stats( + cudf::host_span const> row_group_indices, + parquet_reader_options const& options, + rmm::cuda_stream_view stream) const; + + /** + * @brief Get byte ranges of bloom filters and dictionary pages (secondary filters) for row group + * pruning + * + * @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 options Parquet reader options + * @return Pair of vectors of byte ranges of column chunk with bloom filters and dictionary + * pages subject to filter predicate + */ + [[nodiscard]] std::pair, std::vector> + secondary_filters_byte_ranges(cudf::host_span const> row_group_indices, + parquet_reader_options const& options) const; + private: std::unique_ptr _impl; }; diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp index 626ac249b1bd..4386cd1dea42 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp @@ -20,6 +20,7 @@ #include #include #include +#include namespace cudf::io::parquet::experimental::detail { @@ -202,6 +203,16 @@ std::vector> aggregate_reader_metadata::all_row_groups( if (not opts_row_groups.empty()) { CUDF_EXPECTS(opts_row_groups.size() == per_file_metadata.size(), "Row groups in parquet reader options must specify one vector per data source"); + auto iter = cuda::zip_iterator(opts_row_groups.begin(), per_file_metadata.begin()); + std::for_each(iter, iter + opts_row_groups.size(), [&](auto const& pair) { + auto const& [file_row_groups, file_metadata] = pair; + auto const& row_groups = file_metadata.row_groups; + for (auto const rg_idx : file_row_groups) { + CUDF_EXPECTS(rg_idx >= 0 and std::cmp_less(rg_idx, row_groups.size()), + "Encountered out-of-bounds row group index for data source", + std::invalid_argument); + } + }); return opts_row_groups; } diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp index 64eead4463f7..4dced0a168ad 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp @@ -80,7 +80,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { parquet_reader_options const& options) const; /** - * @copydoc cudf::io::experimental::hybrid_scan::total_rows_in_row_groups + * @copydoc cudf::io::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; @@ -91,14 +91,14 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { void reset_column_selection(); /** - * @copydoc cudf::io::experimental::hybrid_scan::filter_row_groups_with_byte_range + * @copydoc cudf::io::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::filter_row_groups_with_stats + * @copydoc cudf::io::experimental::hybrid_scan_multifile::filter_row_groups_with_stats */ [[nodiscard]] std::vector> filter_row_groups_with_stats( cudf::host_span const> row_group_indices, @@ -106,7 +106,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { rmm::cuda_stream_view stream); /** - * @copydoc cudf::io::experimental::hybrid_scan::secondary_filters_byte_ranges + * @copydoc cudf::io::experimental::hybrid_scan_multifile::secondary_filters_byte_ranges */ [[nodiscard]] std::pair, std::vector> secondary_filters_byte_ranges(cudf::host_span const> row_group_indices, diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp index 31b3cb5a6443..59347367d223 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp @@ -57,4 +57,30 @@ size_type hybrid_scan_multifile::total_rows_in_row_groups( void hybrid_scan_multifile::reset_column_selection() const { _impl->reset_column_selection(); } +std::vector> hybrid_scan_multifile::filter_row_groups_with_byte_range( + cudf::host_span const> row_group_indices, + parquet_reader_options const& options) const +{ + CUDF_FUNC_RANGE(); + return _impl->filter_row_groups_with_byte_range(row_group_indices, options); +} + +std::vector> hybrid_scan_multifile::filter_row_groups_with_stats( + cudf::host_span const> row_group_indices, + parquet_reader_options const& options, + rmm::cuda_stream_view stream) const +{ + CUDF_FUNC_RANGE(); + return _impl->filter_row_groups_with_stats(row_group_indices, options, stream); +} + +std::pair, std::vector> +hybrid_scan_multifile::secondary_filters_byte_ranges( + cudf::host_span const> row_group_indices, + parquet_reader_options const& options) const +{ + CUDF_FUNC_RANGE(); + return _impl->secondary_filters_byte_ranges(row_group_indices, options); +} + } // namespace cudf::io::parquet::experimental 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 1a0c135207f4..29c66093b1cf 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp @@ -7,6 +7,7 @@ #include +#include #include #include #include @@ -218,3 +219,98 @@ TEST_F(HybridScanMultifileFiltersTest, EmptySource) EXPECT_FALSE(page_index_byte_ranges.front().is_empty()); EXPECT_TRUE(page_index_byte_ranges.back().is_empty()); } + +TEST_F(HybridScanMultifileFiltersTest, ErrorFilterRowGroupsWithByteRanges) +{ + using T = uint32_t; + auto constexpr num_sources = 2; + srand(0xb47e); + + std::vector> file_buffers; + file_buffers.reserve(num_sources); + file_buffers.emplace_back(std::get<1>(create_parquet_with_stats())); + file_buffers.emplace_back(std::get<1>(create_parquet_with_stats())); + + auto inputs = build_multifile_inputs(file_buffers); + + // Setting `skip_bytes` or `num_bytes` is ambiguous when reading multiple sources. The reader is + // expected to throw an exception if row groups are filtered using byte range in this case. + { + auto const options = cudf::io::parquet_reader_options::builder().skip_bytes(1000).build(); + auto const reader = std::make_unique( + inputs.footer_byte_spans, options); + auto const row_group_indices = reader->all_row_groups(options); + ASSERT_EQ(row_group_indices.size(), num_sources); + EXPECT_THROW( + std::ignore = reader->filter_row_groups_with_byte_range(row_group_indices, options), + std::invalid_argument); + } + { + auto const options = cudf::io::parquet_reader_options::builder().num_bytes(1000).build(); + auto const reader = std::make_unique( + inputs.footer_byte_spans, options); + auto const row_group_indices = reader->all_row_groups(options); + ASSERT_EQ(row_group_indices.size(), num_sources); + EXPECT_THROW( + std::ignore = reader->filter_row_groups_with_byte_range(row_group_indices, options), + std::invalid_argument); + } +} + +TEST_F(HybridScanMultifileFiltersTest, FilterRowGroupsWithStats) +{ + using T = cudf::duration_ms; + auto constexpr num_sources = 2; + auto constexpr rows_per_row_group = page_size_for_ordered_tests; + + // Two sources, each with 4 row groups and ascending strings in col2 + std::vector> file_buffers; + file_buffers.reserve(num_sources); + srand(0xc001); + file_buffers.emplace_back(std::get<1>(create_parquet_with_stats())); + srand(0xbeef); + file_buffers.emplace_back(std::get<1>(create_parquet_with_stats())); + + auto inputs = build_multifile_inputs(file_buffers); + + // Filter - col0 < 50 and col2 > "000010000" + auto literal_value0 = cudf::duration_scalar(T::rep(50), true, cudf::get_default_stream()); + auto literal0 = cudf::ast::literal(literal_value0); + auto col_ref0 = cudf::ast::column_reference(0); + auto filter1 = cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref0, literal0); + + auto literal_value2 = cudf::string_scalar("000010000", true, cudf::get_default_stream()); + auto literal2 = cudf::ast::literal(literal_value2); + auto col_ref2 = cudf::ast::column_reference(2); + auto filter2 = cudf::ast::operation(cudf::ast::ast_operator::GREATER, literal2, col_ref2); + + auto filter_expression = + cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_AND, filter1, filter2); + + auto options = cudf::io::parquet_reader_options::builder().filter(filter_expression).build(); + auto const reader = std::make_unique( + inputs.footer_byte_spans, options); + + // Each source has 4 row groups (20000 rows / 5000 rows per row group) + auto input_row_group_indices = reader->all_row_groups(options); + ASSERT_EQ(input_row_group_indices.size(), num_sources); + EXPECT_EQ(reader->total_rows_in_row_groups(input_row_group_indices), + num_sources * 4 * rows_per_row_group); + + // Each source prunes down to a single surviving row group + auto stats_filtered = reader->filter_row_groups_with_stats( + input_row_group_indices, options, cudf::get_default_stream()); + ASSERT_EQ(stats_filtered.size(), num_sources); + for (std::size_t i = 0; i < stats_filtered.size(); ++i) { + EXPECT_EQ(stats_filtered[i].size(), 1) << "Source index: " << i; + } + EXPECT_EQ(reader->total_rows_in_row_groups(stats_filtered), num_sources * rows_per_row_group); + + // Custom per-source indices that prune all row groups via stats, including an empty source + input_row_group_indices = {{1, 2}, {}}; + stats_filtered = reader->filter_row_groups_with_stats( + input_row_group_indices, options, cudf::get_default_stream()); + ASSERT_EQ(stats_filtered.size(), num_sources); + EXPECT_TRUE(stats_filtered.front().empty()); + EXPECT_TRUE(stats_filtered.back().empty()); +}