Improve Parquet reader pass construction - #23446
Conversation
Pass construction sized each row group by summing `total_compressed_size` over *every* column chunk in the row group, via `aggregate_reader_metadata::get_row_group_properties()`. When only a few columns of a wide schema are read, that over-estimates the pass footprint by roughly the projection ratio, so the reader splits into far more passes than the `pass_read_limit` requires. Each extra pass costs another `read_compressed_data()` round trip, another page-header decode, another dictionary decompression, and smaller decode batches. This was also internally inconsistent: the pass *partition* used the all-columns size while `pass.base_mem_size` - which derives the *subpass* budget - already used the selected columns only. Decouple `compute_row_group_passes()` from `row_group_info` by introducing `row_group_pass_size_info`, and add `make_row_group_pass_size_info()`, which reduces over the `ColumnChunkDesc` array that `create_global_chunk_info()` has already built. Those descriptors exist only for the selected leaf columns, so the resulting per-row-group size is exact rather than estimated, and needs no new metadata plumbing. `max_leaf_values` narrows to the selected columns for the same reason: it guards the `size_type` limit on decoded column length, and unselected columns are never dremel-decoded and never allocate an output buffer. `row_group_info::compressed_size` / `max_leaf_values` had no remaining readers, so drop them along with the now-dead `total_size` element of `get_row_group_properties()`'s return. The hybrid scan reader keeps its existing all-columns estimate here; making its planning API column-aware is a follow-up. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`hybrid_scan_reader::construct_row_group_passes()` estimated each row group's compressed size over every column in the row group, which its own documentation acknowledged as a deliberately conservative approximation. That approximation is worst exactly where hybrid scan is used: a filter pass typically touches one or two columns out of a wide schema, so the planner splits into roughly `total_columns / selected_columns` times more passes than `pass_read_limit` requires. Give the planning API the reader options so it can resolve the actual column selection, and add `aggregate_reader_metadata::get_row_group_pass_size_info()`, which sums `total_compressed_size` and takes the max `num_values` over just the selected leaf column chunks. Schema indices are mapped per source with `map_schema_index()` so mismatched-schema multi-source reads stay correct. Sizing uses `ALL_COLUMNS` - the union of projected payload columns and filter columns - because this call cannot know whether the caller will use the two-stage filter/payload flow or `setup_chunking_for_all_columns()`, and the latter needs the full footprint. `select_columns()` memoizes via the `_is_*_columns_selected` flags and `reset_internal_state()` does not clear them, so planning calls `reset_column_selection()` afterwards. Without it, a later `setup_chunking_for_all_columns()` would short-circuit selection and never rebuild `_output_buffers` / `_output_buffers_template`. `get_row_group_properties()` had no remaining callers and is removed. The JNI passes the wrapper's current options, which already track any filter installed via `setFilter`, so no Java signature change is needed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Size passes from selected columns, simplify row-group sizing, and retain hybrid scan compatibility.
9634e08 to
9515372
Compare
| row_group_ids.reserve(total_row_groups); | ||
| row_group_sizes.reserve(total_row_groups); | ||
|
|
||
| std::for_each(cuda::counting_iterator<cudf::size_type>(0), |
There was a problem hiding this comment.
Updated this loop to collect required info using the new get_row_group_size_info API.
|
|
||
| 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); |
There was a problem hiding this comment.
row_group_sizes are used here
| 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()); } |
There was a problem hiding this comment.
row_group_ids are used here.
| CUDF_EXPECTS(chunks.size() == row_groups_info.size() * num_input_columns, | ||
| "Mismatch between the number of column chunk descriptors and row groups"); | ||
|
|
||
| std::vector<row_group_size_info> row_group_sizes; | ||
| row_group_sizes.reserve(row_groups_info.size()); | ||
|
|
||
| std::transform(cuda::counting_iterator<size_t>(0), | ||
| cuda::counting_iterator<size_t>(row_groups_info.size()), | ||
| std::back_inserter(row_group_sizes), | ||
| [&](auto const rg_idx) { | ||
| auto const& rg_info = row_groups_info[rg_idx]; | ||
| auto const chunks_begin = chunks.begin() + (rg_idx * num_input_columns); | ||
| auto const chunks_end = chunks_begin + num_input_columns; | ||
| size_t compressed_size = 0; | ||
| size_t max_leaf_values = 0; | ||
| std::for_each(chunks_begin, chunks_end, [&](auto const& chunk) { | ||
| compressed_size += chunk.compressed_size; | ||
| max_leaf_values = std::max(max_leaf_values, chunk.num_values); | ||
| }); | ||
| return row_group_size_info{.unadjusted_num_rows = rg_info.unadjusted_num_rows, | ||
| .compressed_size = compressed_size, | ||
| .max_leaf_values = max_leaf_values}; | ||
| }); | ||
|
|
||
| return row_group_sizes; | ||
| } |
There was a problem hiding this comment.
Collect size info from each column chunk (of selected columns) to get the final size of each row group instead of just row_group.size which includes size for all column chunks.
| 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; | ||
| } |
There was a problem hiding this comment.
The rgi.start_row has for some time been always zero so no need to subtract it and made this simpler
| }); | ||
| 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 |
| size_info.compressed_size += column_metadata.total_compressed_size; | ||
| size_info.max_leaf_values = | ||
| std::max(size_info.max_leaf_values, static_cast<size_t>(column_metadata.num_values)); |
There was a problem hiding this comment.
I could have moved these two lines to a lambda to make these loops a bit simpler but the next PR will remove the else path altogether. That is only needed here as column selection not yet passed from the hybrid scan path.
| /** | ||
| * @brief Row group size information for pass partitioning. | ||
| */ | ||
| struct row_group_size_info { |
There was a problem hiding this comment.
Separated size related fields from row_group_info above.
| */ | ||
| [[nodiscard]] std::vector<std::string> get_pandas_index_names() const; | ||
|
|
||
| /** |
There was a problem hiding this comment.
Remove in favor of get_row_group_size_info
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe Parquet reader now computes row-group sizing separately from row-group identity, supports selected-column size aggregation, uses size metadata for pass generation, updates hybrid-scan mapping, and adds projected-column chunking coverage. ChangesParquet pass computation
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cpp/src/io/parquet/reader_impl_chunking_utils.cuh (1)
6-21: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd explicit
#include <span>.
std::span<row_group_size_info const>is used directly in this header'scompute_row_group_passesdeclaration, but<span>isn't included here — it's only available transitively throughreader_impl_chunking.hpp→reader_impl_helpers.hpp. Relying on a transitive include for a directly-used symbol is fragile to future header reshuffling.As per coding guidelines: "include headers directly for every used symbol without unused or incorrectly styled includes."
♻️ Proposed fix
`#include` <rmm/cuda_stream_view.hpp> `#include` <cuda/functional> `#include` <cuda/std/utility> +#include <span> `#include` <thrust/binary_search.h>Also applies to: 796-798
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/io/parquet/reader_impl_chunking_utils.cuh` around lines 6 - 21, Add an explicit standard-library <span> include to the header declaring compute_row_group_passes, since it directly uses std::span<row_group_size_info const>. Keep the existing includes unchanged and ensure the declaration no longer depends on the transitive inclusion through reader_impl_chunking.hpp.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@cpp/src/io/parquet/reader_impl_chunking_utils.cuh`:
- Around line 6-21: Add an explicit standard-library <span> include to the
header declaring compute_row_group_passes, since it directly uses
std::span<row_group_size_info const>. Keep the existing includes unchanged and
ensure the declaration no longer depends on the transitive inclusion through
reader_impl_chunking.hpp.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 46b25313-61b7-4e0e-bb2c-1d7ad3ec4c3b
📒 Files selected for processing (7)
cpp/src/io/parquet/experimental/hybrid_scan_impl.cppcpp/src/io/parquet/reader_impl_chunking.cucpp/src/io/parquet/reader_impl_chunking_utils.cucpp/src/io/parquet/reader_impl_chunking_utils.cuhcpp/src/io/parquet/reader_impl_helpers.cppcpp/src/io/parquet/reader_impl_helpers.hppcpp/tests/io/parquet_chunked_reader_test.cu
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/src/io/parquet/reader_impl_helpers.cpp (1)
1239-1267: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winValidate signed Parquet metadata before converting to
size_t.Line 1247 and Lines 1253-1261 cast/add signed metadata without checking for negative values or overflow. Malformed metadata can wrap to huge or underestimated pass sizes, causing incorrect partitioning or excessive decompression allocations. Reject invalid values and use checked additions before populating
row_group_size_info.As per coding guidelines, validate inputs and prevent invalid memory access.
Suggested validation
- auto size_info = - row_group_size_info{.unadjusted_num_rows = static_cast<size_t>(row_group.num_rows)}; + CUDF_EXPECTS(row_group.num_rows >= 0, "Invalid negative row-group row count"); + auto size_info = + row_group_size_info{.unadjusted_num_rows = static_cast<size_t>(row_group.num_rows)}; + auto add_checked = [](size_t& total, int64_t value) { + CUDF_EXPECTS(value >= 0, "Invalid negative Parquet size"); + auto const converted = static_cast<size_t>(value); + CUDF_EXPECTS(converted <= std::numeric_limits<size_t>::max() - total, + "Parquet size metadata overflow"); + total += converted; + };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/io/parquet/reader_impl_helpers.cpp` around lines 1239 - 1267, Update aggregate_reader_metadata::get_row_group_size_info to validate row_group.num_rows, column metadata num_values, and total_compressed_size are non-negative before conversion or accumulation. Use checked additions for compressed_size and reject values that would overflow the row_group_size_info fields, preserving valid metadata behavior while failing malformed metadata before populating size_info.Source: Coding guidelines
cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp (1)
806-809: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winCompute pass sizes from the selected columns.
_input_columnscontains the active projection, but this call passesstd::nullopt, so projected reads are still sized using unprojected row-group metadata. That can create unnecessarily many input passes and defeats the PR objective of less-conservative selected-column sizing. Pass the selected schema-index span toget_row_group_size_infohere. This is the same unresolved concern previously noted on Line 806.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp` around lines 806 - 809, Update the row-group sizing call in the hybrid scan implementation to pass the selected schema-index span derived from _input_columns instead of std::nullopt. Ensure get_row_group_size_info receives the active projection for each row group while preserving the existing rg_index and source_index arguments.
🧹 Nitpick comments (1)
cpp/src/io/parquet/reader_impl_helpers.cpp (1)
60-66: 📐 Maintainability & Code Quality | 🔵 TrivialObtain the required libcudf C++ approvals before merge.
As per coding guidelines, changes to
cpp/**/*.cpprequire at least two approvals from the cudf-cpp-codeowners before merging.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/io/parquet/reader_impl_helpers.cpp` around lines 60 - 66, Obtain at least two approvals from the cudf-cpp-codeowners before merging this change to find_colchunk_iter_offset in the C++ source.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp`:
- Around line 806-809: Update the row-group sizing call in the hybrid scan
implementation to pass the selected schema-index span derived from
_input_columns instead of std::nullopt. Ensure get_row_group_size_info receives
the active projection for each row group while preserving the existing rg_index
and source_index arguments.
In `@cpp/src/io/parquet/reader_impl_helpers.cpp`:
- Around line 1239-1267: Update
aggregate_reader_metadata::get_row_group_size_info to validate
row_group.num_rows, column metadata num_values, and total_compressed_size are
non-negative before conversion or accumulation. Use checked additions for
compressed_size and reject values that would overflow the row_group_size_info
fields, preserving valid metadata behavior while failing malformed metadata
before populating size_info.
---
Nitpick comments:
In `@cpp/src/io/parquet/reader_impl_helpers.cpp`:
- Around line 60-66: Obtain at least two approvals from the cudf-cpp-codeowners
before merging this change to find_colchunk_iter_offset in the C++ source.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 06180efe-387b-48bf-8e40-87e07e35fc54
📒 Files selected for processing (3)
cpp/src/io/parquet/experimental/hybrid_scan_impl.cppcpp/src/io/parquet/reader_impl_helpers.cppcpp/src/io/parquet/reader_impl_helpers.hpp
Description
Part 1 of 3.
This PR improves Parquet reader's pass construction by making it column chunk aware so that it is more accurate (less conservative) when reading select columns from the file.
Subsequent PRs will propagate these improvements to hybrid scan as well as allow page-level calculations once #23375 merges
Checklist