Improve Parquet reader pass construction

cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp

@@ -793,59 +793,49 @@ hybrid_scan_reader_impl::construct_row_group_passes(
   CUDF_EXPECTS(
     pass_read_limit > 0, "Pass read limit must be greater than 0", std::invalid_argument);
 
-  auto row_groups_info = std::vector<row_group_info>{};
-  row_groups_info.reserve(total_row_groups);
-  size_t start_row = 0;
+  auto row_group_ids   = std::vector<std::pair<size_type, size_type>>{};
+  auto row_group_sizes = std::vector<cudf::io::parquet::detail::row_group_size_info>{};
+  row_group_ids.reserve(total_row_groups);
+  row_group_sizes.reserve(total_row_groups);
+
   std::for_each(cuda::counting_iterator<cudf::size_type>(0),
                 cuda::counting_iterator<cudf::size_type>(row_group_indices.size()),
                 [&](auto const source_index) {
-                  auto const& src_row_groups = row_group_indices[source_index];
-                  std::transform(
-                    src_row_groups.begin(),
-                    src_row_groups.end(),
-                    std::back_inserter(row_groups_info),
-                    [&](auto const rg_index) {
-                      auto const& row_group =
-                        _extended_metadata->get_row_group(rg_index, source_index);
-                      auto const [compressed_size, total_size, num_rows, max_leaf_values] =
-                        _extended_metadata->get_row_group_properties(row_group);
-                      auto rg_info = row_group_info{.index               = rg_index,
-                                                    .start_row           = start_row,
-                                                    .unadjusted_num_rows = num_rows,
-                                                    .source_index        = source_index,
-                                                    .compressed_size     = compressed_size,
-                                                    .max_leaf_values     = max_leaf_values};
-                      start_row += num_rows;
-                      return rg_info;
-                    });
+                  for (auto const rg_index : row_group_indices[source_index]) {
+                    row_group_ids.emplace_back(rg_index, source_index);
+                    // TODO(mh): Compute the row group size information over the selected columns
+                    // instead
+                    row_group_sizes.push_back(_extended_metadata->get_row_group_size_info(
+                      rg_index, source_index, std::nullopt));
+                  }
                 });
 
   auto const comp_read_limit = static_cast<std::size_t>(
     pass_read_limit * cudf::io::parquet::detail::input_limit_compression_reserve);
 
   auto const pass_data =
-    cudf::io::parquet::detail::compute_row_group_passes(row_groups_info, comp_read_limit, 0);
+    cudf::io::parquet::detail::compute_row_group_passes(row_group_sizes, comp_read_limit, 0);
 
   // Convert offset-based pass boundaries back to vectors of row group indices
   auto const& offsets = pass_data.pass_row_group_offsets;
   auto passes         = std::vector<std::vector<cudf::size_type>>{};
   passes.reserve(offsets.size() - 1);
   auto row_group_source_map       = std::vector<cudf::size_type>{};
   auto const has_multiple_sources = row_group_indices.size() > 1;
-  if (has_multiple_sources) { row_group_source_map.reserve(row_groups_info.size()); }
+  if (has_multiple_sources) { row_group_source_map.reserve(row_group_ids.size()); }
   std::transform(offsets.begin(),
                  offsets.end() - 1,
                  offsets.begin() + 1,
                  std::back_inserter(passes),
                  [&](auto const start, auto const end) {
                    auto pass = std::vector<cudf::size_type>{};
                    pass.reserve(end - start);
-                   std::for_each(row_groups_info.begin() + start,
-                                 row_groups_info.begin() + end,
-                                 [&](auto const& rg_info) {
-                                   pass.emplace_back(rg_info.index);
+                   std::for_each(row_group_ids.begin() + start,
+                                 row_group_ids.begin() + end,
+                                 [&](auto const& row_group_id) {
+                                   pass.emplace_back(row_group_id.first);
                                    if (has_multiple_sources) {
-                                     row_group_source_map.emplace_back(rg_info.source_index);
+                                     row_group_source_map.emplace_back(row_group_id.second);
                                    }
                                  });
                    return pass;

cpp/src/io/parquet/reader_impl_chunking.cu

@@ -543,8 +543,19 @@ void reader_impl::compute_input_passes(read_mode mode)
       ? static_cast<size_t>(_input_pass_read_limit * input_limit_compression_reserve)
       : std::numeric_limits<std::size_t>::max();
 
+  // Compute size information for each row group by the columns we are actually going to read.
+  auto row_group_sizes = std::vector<row_group_size_info>{};
+  row_group_sizes.reserve(row_groups_info.size());
+  std::transform(row_groups_info.cbegin(),
+                 row_groups_info.cend(),
+                 std::back_inserter(row_group_sizes),
+                 [&](auto const& row_group) {
+                   return _metadata->get_row_group_size_info(
+                     row_group.index, row_group.source_index, _input_columns);
+                 });
+
   auto pass_data =
-    compute_row_group_passes(row_groups_info, comp_read_limit, _file_itm_data.global_skip_rows);
+    compute_row_group_passes(row_group_sizes, comp_read_limit, _file_itm_data.global_skip_rows);
 
   _file_itm_data.input_pass_row_group_offsets = std::move(pass_data.pass_row_group_offsets);
   _file_itm_data.input_pass_start_row_count   = std::move(pass_data.pass_start_row_counts);

cpp/src/io/parquet/reader_impl_chunking_utils.cu

@@ -953,7 +953,7 @@ rmm::device_uvector<size_t> compute_level_decode_sizes(device_span<ColumnChunkDe
   return level_decode_sizes;
 }
 
-row_group_pass_data compute_row_group_passes(cudf::host_span<row_group_info const> row_groups_info,
+row_group_pass_data compute_row_group_passes(std::span<row_group_size_info const> row_group_sizes,
                                              std::size_t comp_read_limit,
                                              int64_t skip_rows)
 {
@@ -970,14 +970,18 @@ row_group_pass_data compute_row_group_passes(cudf::host_span<row_group_info cons
   std::size_t cur_rg_start             = 0;
   std::size_t cur_row_count            = 0;
 
-  for (std::size_t cur_rg_index = 0; cur_rg_index < row_groups_info.size(); cur_rg_index++) {
-    auto const& rgi = row_groups_info[cur_rg_index];
+  for (std::size_t cur_rg_index = 0; cur_rg_index < row_group_sizes.size(); cur_rg_index++) {
+    auto const& rgi = row_group_sizes[cur_rg_index];
 
     // We must use the effective size of the first row group we are reading to accurately calculate
     // the first non-zero `input_pass_start_row_count` unless we are reading only one row group
-    auto const row_group_rows = (skip_rows and row_groups_info.size() > 1)
-                                  ? (rgi.start_row + rgi.unadjusted_num_rows - skip_rows)
-                                  : rgi.unadjusted_num_rows;
+    auto row_group_rows = rgi.unadjusted_num_rows;
+    if (row_group_sizes.size() > 1) {
+      CUDF_EXPECTS(std::cmp_greater_equal(rgi.unadjusted_num_rows, skip_rows),
+                   "Row groups must contribute non-negative effective rows",
+                   std::invalid_argument);
+      row_group_rows -= skip_rows;
+    }
 
     auto const compressed_rg_size    = rgi.compressed_size;
     auto const row_group_leaf_values = rgi.max_leaf_values;
@@ -1026,8 +1030,8 @@ row_group_pass_data compute_row_group_passes(cudf::host_span<row_group_info cons
   }
 
   // Add the last pass if necessary
-  if (result.pass_row_group_offsets.back() != row_groups_info.size()) {
-    result.pass_row_group_offsets.push_back(row_groups_info.size());
+  if (result.pass_row_group_offsets.back() != row_group_sizes.size()) {
+    result.pass_row_group_offsets.push_back(row_group_sizes.size());
     result.pass_start_row_counts.push_back(cur_row_count);
   }
 

cpp/src/io/parquet/reader_impl_chunking_utils.cuh

@@ -787,13 +787,13 @@ struct row_group_pass_data {
  * @brief Partition row groups into passes based on compressed size, leaf value count,
  * and row count thresholds.
  *
- * @param row_groups_info Span of row group metadata
+ * @param row_group_sizes Span of row group size information
  * @param comp_read_limit Maximum compressed bytes per pass
  * @param skip_rows Number of leading rows to skip (affects the effective size of the first row
  * group)
  * @return A row_group_pass_data containing pass boundary offsets and cumulative row counts
  */
-row_group_pass_data compute_row_group_passes(cudf::host_span<row_group_info const> row_groups_info,
+row_group_pass_data compute_row_group_passes(std::span<row_group_size_info const> row_group_sizes,
                                              std::size_t comp_read_limit,
                                              int64_t skip_rows);
 

cpp/src/io/parquet/reader_impl_helpers.cpp

@@ -1236,6 +1236,49 @@ RowGroup const& aggregate_reader_metadata::get_row_group(size_type row_group_ind
   return per_file_metadata[src_idx].row_groups[row_group_index];
 }
 
+row_group_size_info aggregate_reader_metadata::get_row_group_size_info(
+  size_type row_group_index,
+  size_type src_idx,
+  std::optional<std::span<input_column_info const>> input_columns) const
+{
+  auto const& row_group = get_row_group(row_group_index, src_idx);
+
+  CUDF_EXPECTS(row_group.num_rows >= 0 && std::in_range<size_t>(row_group.num_rows),
+               "Row group has an invalid number of rows",
+               std::invalid_argument);
+  auto size_info =
+    row_group_size_info{.unadjusted_num_rows = static_cast<size_t>(row_group.num_rows)};
+
+  // Helper function to overflow-safeadd compressed sizes
+  auto const add_compressed_size = [&](auto&& current_size, auto&& colchunk_compressed_size) {
+    auto const sum = cuda::add_overflow<size_t>(current_size, colchunk_compressed_size);
+    CUDF_EXPECTS(not sum.overflow,
+                 "Row group compressed size exceeds the supported range",
+                 std::overflow_error);
+    return sum.value;
+  };
+
+  if (input_columns.has_value()) {
+    for (auto const& column : *input_columns) {
+      auto const& column_metadata =
+        get_column_metadata(row_group_index, src_idx, column.schema_idx);
+      size_info.compressed_size =
+        add_compressed_size(size_info.compressed_size, column_metadata.total_compressed_size);
+      size_info.max_leaf_values =
+        std::max<size_t>(size_info.max_leaf_values, column_metadata.num_values);
+    }
+  } else {
+    for (auto const& column_chunk : row_group.columns) {
+      size_info.compressed_size = add_compressed_size(size_info.compressed_size,
+                                                      column_chunk.meta_data.total_compressed_size);
+      size_info.max_leaf_values =
+        std::max<size_t>(size_info.max_leaf_values, column_chunk.meta_data.num_values);
+    }
+  }
+
+  return size_info;
+}
+
 ColumnChunkMetaData const& aggregate_reader_metadata::get_column_metadata(size_type row_group_index,
                                                                           size_type src_idx,
                                                                           int schema_idx) const
@@ -1412,31 +1455,6 @@ std::vector<std::string> aggregate_reader_metadata::get_pandas_index_names() con
   return names;
 }
 
-std::tuple<size_t, size_t, size_t, size_t> aggregate_reader_metadata::get_row_group_properties(
-  RowGroup const& row_group) const
-{
-  auto const compressed_size = std::transform_reduce(
-    row_group.columns.cbegin(),
-    row_group.columns.cend(),
-    size_t{0},
-    std::plus<>(),
-    [](auto const& colchunk) { return colchunk.meta_data.total_compressed_size; });
-
-  auto const total_size = compressed_size + row_group.total_byte_size;
-
-  size_t const max_leaf_values =
-    row_group.columns.empty()
-      ? 0
-      : std::max_element(row_group.columns.cbegin(),
-                         row_group.columns.cend(),
-                         [](auto const& a, auto const& b) {
-                           return a.meta_data.num_values < b.meta_data.num_values;
-                         })
-          ->meta_data.num_values;
-
-  return {compressed_size, total_size, static_cast<size_t>(row_group.num_rows), max_leaf_values};
-}
-
 std::tuple<int64_t,
            std::vector<std::vector<size_type>>,
            std::vector<std::vector<size_t>>,
@@ -1790,10 +1808,6 @@ aggregate_reader_metadata::select_row_groups(
           // Update the number of rows read from this data source
           num_rows_per_source[src_idx] += num_rows_this_row_group;
 
-          // Get row group properties
-          auto const [compressed_size, total_size, num_rows, max_leaf_values] =
-            get_row_group_properties(rg);
-
           // We need the unadjusted start index of this row group to correctly
           // initialize ColumnChunkDesc for this row group in
           // create_global_chunk_info() and calculate the row offset for the first
@@ -1802,10 +1816,8 @@ aggregate_reader_metadata::select_row_groups(
             row_group_info{.index               = rg_idx,
                            .start_row           = row_group_start_row,
                            .source_start_row    = source_row_offsets[rg_idx],
-                           .unadjusted_num_rows = num_rows,
-                           .source_index        = static_cast<cudf::size_type>(src_idx),
-                           .compressed_size     = compressed_size,
-                           .max_leaf_values     = max_leaf_values});
+                           .unadjusted_num_rows = static_cast<size_t>(rg.num_rows),
+                           .source_index        = static_cast<cudf::size_type>(src_idx)});
 
           // If page-level indexes are present, then collect extra chunk and page
           // info. The page indexes rely on absolute row numbers - not adjusted for

cpp/src/io/parquet/reader_impl_helpers.hpp

@@ -16,6 +16,7 @@
 #include <cstddef>
 #include <functional>
 #include <optional>
+#include <span>
 #include <string>
 #include <string_view>
 #include <tuple>
@@ -66,19 +67,26 @@ struct column_chunk_info {
  * @brief The row_group_info class
  */
 struct row_group_info {
-  size_type index;  // row group index within a file. aggregate_reader_metadata::get_row_group() is
-                    // called with index and source_index
-  size_t start_row;
+  size_type index;   // row group index within a file. aggregate_reader_metadata::get_row_group() is
+                     // called with index and source_index
+  size_t start_row;  // global start row of this row group
   size_t source_start_row;     // file-local start row of this row group within its source file
   size_t unadjusted_num_rows;  // number of unadjusted rows in the row group
   size_type source_index;      // file index.
-  size_t compressed_size;      // compressed size of the row group
-  size_t max_leaf_values;      // maximum number of leaf values in the row group
 
   // Optional metadata pulled from the column and offset indexes, if present.
   std::optional<std::vector<column_chunk_info>> column_chunks;
 };
 
+/**
+ * @brief Row group size information for pass partitioning.
+ */
+struct row_group_size_info {
+  size_t unadjusted_num_rows;  // number of unadjusted rows in this row group
+  size_t compressed_size;      // compressed size of the selected columns in this row group
+  size_t max_leaf_values;      // maximum number of leaf values over the selected columns
+};
+
 /**
  * @brief Translates Parquet datatype to cuDF type enum
  */
@@ -434,6 +442,22 @@ class aggregate_reader_metadata {
    */
   [[nodiscard]] RowGroup const& get_row_group(size_type row_group_index, size_type src_idx) const;
 
+  /**
+   * @brief Computes row group size information over selected columns
+   *
+   * When `input_columns` is specified, computes the compressed size and maximum leaf value count
+   * over only those columns. Otherwise, over all columns in the row group.
+   *
+   * @param row_group_index Index of the row group within its source
+   * @param src_idx Index of the input source
+   * @param input_columns Optional selected leaf columns
+   * @return Row group size information
+   */
+  [[nodiscard]] row_group_size_info get_row_group_size_info(
+    size_type row_group_index,
+    size_type src_idx,
+    std::optional<std::span<input_column_info const>> input_columns) const;
+
   /**
    * @brief Check if all row groups have an offset index
    *
@@ -622,18 +646,6 @@ class aggregate_reader_metadata {
    */
   [[nodiscard]] std::vector<std::string> get_pandas_index_names() const;
 
-  /**
-   * @brief Computes the compressed and total size, the number of rows, and the maximum number of
-   * leaf values in the specified row group
-   *
-   * @param row_group The row group
-   *
-   * @return A tuple of row group compressed size, total size, number of rows, and maximum leaf
-   * values
-   */
-  [[nodiscard]] std::tuple<size_t, size_t, size_t, size_t> get_row_group_properties(
-    RowGroup const& rg) const;
-
   /**
    * @brief Filters the row groups using stats and bloom filters based on predicate filter
    *