Add multifile row group pruning with stats and byte ranges

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<std::vector<size_type>> filter_row_groups_with_byte_range(
+    cudf::host_span<std::vector<size_type> 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<std::vector<size_type>> filter_row_groups_with_stats(
+    cudf::host_span<std::vector<size_type> 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<byte_range_info>, std::vector<byte_range_info>>
+  secondary_filters_byte_ranges(cudf::host_span<std::vector<size_type> const> row_group_indices,
+                                parquet_reader_options const& options) const;
+
  private:
   std::unique_ptr<detail::hybrid_scan_reader_impl> _impl;
 };

cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp

@@ -20,6 +20,7 @@
 #include <numeric>
 #include <optional>
 #include <unordered_set>
+#include <utility>
 
 namespace cudf::io::parquet::experimental::detail {
 
@@ -202,6 +203,16 @@ std::vector<std::vector<size_type>> 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;
   }
 

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<std::vector<size_type> const> row_group_indices) const;
@@ -91,22 +91,22 @@ 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<std::vector<cudf::size_type>> filter_row_groups_with_byte_range(
     cudf::host_span<std::vector<size_type> 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<std::vector<size_type>> filter_row_groups_with_stats(
     cudf::host_span<std::vector<size_type> const> row_group_indices,
     parquet_reader_options const& options,
     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<byte_range_info>, std::vector<byte_range_info>>
   secondary_filters_byte_ranges(cudf::host_span<std::vector<size_type> const> row_group_indices,

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<std::vector<size_type>> hybrid_scan_multifile::filter_row_groups_with_byte_range(
+  cudf::host_span<std::vector<size_type> 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<std::vector<size_type>> hybrid_scan_multifile::filter_row_groups_with_stats(
+  cudf::host_span<std::vector<size_type> 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<text::byte_range_info>, std::vector<text::byte_range_info>>
+hybrid_scan_multifile::secondary_filters_byte_ranges(
+  cudf::host_span<std::vector<size_type> 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

cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp

@@ -7,6 +7,7 @@
 
 #include <cudf_test/base_fixture.hpp>
 
+#include <cudf/ast/expressions.hpp>
 #include <cudf/copying.hpp>
 #include <cudf/io/experimental/hybrid_scan_multifile.hpp>
 #include <cudf/io/parquet.hpp>
@@ -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<std::vector<char>> file_buffers;
+  file_buffers.reserve(num_sources);
+  file_buffers.emplace_back(std::get<1>(create_parquet_with_stats<T, 1>()));
+  file_buffers.emplace_back(std::get<1>(create_parquet_with_stats<T, 1>()));
+
+  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<cudf::io::parquet::experimental::hybrid_scan_multifile>(
+      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<cudf::io::parquet::experimental::hybrid_scan_multifile>(
+      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<std::vector<char>> file_buffers;
+  file_buffers.reserve(num_sources);
+  srand(0xc001);
+  file_buffers.emplace_back(std::get<1>(create_parquet_with_stats<T, 1, false>()));
+  srand(0xbeef);
+  file_buffers.emplace_back(std::get<1>(create_parquet_with_stats<T, 1, false>()));
+
+  auto inputs = build_multifile_inputs(file_buffers);
+
+  // Filter - col0 < 50 and col2 > "000010000"
+  auto literal_value0 = cudf::duration_scalar<T>(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<cudf::io::parquet::experimental::hybrid_scan_multifile>(
+    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());
+}