Skip to content

Enable chunked row masks in hybrid scan reader #22716

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 7 commits into from
Jun 4, 2026
Merged

Enable chunked row masks in hybrid scan reader #22716

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
f59ebf2
Hybrid scan reader allows chunked row masks
mhaseeb123 May 29, 2026
13bea36
Merge branch 'main' into fea/hybrid-scan-chunked-row-mask
mhaseeb123 Jun 1, 2026
040ed23
Minor improvement
mhaseeb123 Jun 1, 2026
2fe164d
Merge branch 'main' into fea/hybrid-scan-chunked-row-mask
mhaseeb123 Jun 2, 2026
a4c326f
Merge branch 'main' into fea/hybrid-scan-chunked-row-mask
mhaseeb123 Jun 2, 2026
2f75fd8
Add a new test
mhaseeb123 Jun 4, 2026
b71d111
Merge branch 'main' into fea/hybrid-scan-chunked-row-mask
mhaseeb123 Jun 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 24 additions & 16 deletions cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,25 +171,33 @@ void aggregate_reader_metadata::setup_page_index(cudf::host_span<uint8_t const>
file_metadata.setup_page_index(page_index_bytes, min_offset);
}

size_type aggregate_reader_metadata::total_rows_in_row_groups(
std::size_t aggregate_reader_metadata::total_rows_in_row_groups(
cudf::host_span<std::vector<size_type> const> row_group_indices) const
{
std::size_t total_rows = 0;
CUDF_EXPECTS(row_group_indices.size() == per_file_metadata.size(),
"Encountered unexpected number of input row group indices",
std::invalid_argument);

std::for_each(cuda::counting_iterator<std::size_t>{0},
cuda::counting_iterator{row_group_indices.size()},
[&](auto const src_idx) {
auto const& pfm = per_file_metadata[src_idx];
for (auto const row_group_idx : row_group_indices[src_idx]) {
CUDF_EXPECTS(std::cmp_less(row_group_idx, pfm.row_groups.size()),
"Row group index out of bounds");
total_rows += pfm.row_groups[row_group_idx].num_rows;
}
});
CUDF_EXPECTS(std::cmp_less_equal(total_rows, std::numeric_limits<size_type>::max()),

@mhaseeb123 mhaseeb123 May 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed this as we now allow more than 2B rows

"Total number of rows exceeds cudf::size_type's limit");

return static_cast<size_type>(total_rows);
return std::accumulate(

@mhaseeb123 mhaseeb123 May 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No logical differences here except using std::accumulate instead of for_each.

cuda::counting_iterator<std::size_t>{0},
cuda::counting_iterator{row_group_indices.size()},
std::size_t{0},
[&](auto sum, auto const src_idx) {
auto const& file_metadata = per_file_metadata[src_idx];
return std::accumulate(
row_group_indices[src_idx].begin(),
row_group_indices[src_idx].end(),
sum,
[&](auto sum, auto const row_group_idx) {
CUDF_EXPECTS(
std::cmp_greater_equal(row_group_idx, size_type{0}) and
std::cmp_less(row_group_idx, file_metadata.row_groups.size()),
"Encountered out-of-bounds row group index for data source. Row group index: " +
std::to_string(row_group_idx) + ", Source index: " + std::to_string(src_idx) +
", Number of row groups: " + std::to_string(file_metadata.row_groups.size()));
return sum + file_metadata.row_groups[row_group_idx].num_rows;
Comment on lines +239 to +247

@coderabbitai coderabbitai Bot Jun 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Validate row counts before adding them to the unsigned accumulator.

num_rows is metadata-derived and signed. sum + file_metadata.row_groups[row_group_idx].num_rows will wrap on negative metadata and still has no guard for std::size_t overflow, so this can return a bogus total and mis-size later row-mask or buffer work instead of failing fast.

🩹 Suggested guard 
         [&](auto sum, auto const row_group_idx) {
           CUDF_EXPECTS(
             std::cmp_greater_equal(row_group_idx, size_type{0}) and
               std::cmp_less(row_group_idx, file_metadata.row_groups.size()),
             "Encountered out-of-bounds row group index for data source. Row group index: " +
               std::to_string(row_group_idx) + ", Source index: " + std::to_string(src_idx) +
               ", Number of row groups: " + std::to_string(file_metadata.row_groups.size()));
-          return sum + file_metadata.row_groups[row_group_idx].num_rows;
+          auto const num_rows = file_metadata.row_groups[row_group_idx].num_rows;
+          CUDF_EXPECTS(num_rows >= 0,
+                       "Encountered negative row count in row group metadata",
+                       std::invalid_argument);
+          auto const rows = static_cast<std::size_t>(num_rows);
+          CUDF_EXPECTS(sum <= std::numeric_limits<std::size_t>::max() - rows,
+                       "Total number of rows exceeds std::size_t",
+                       std::overflow_error);
+          return sum + rows;
         });
🤖 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_helpers.cpp` around lines 191 -
198, The lambda that sums file_metadata.row_groups[row_group_idx].num_rows into
an unsigned accumulator can accept negative num_rows and overflow; update the
lambda (used with sum and row_group_idx) to first assert that
file_metadata.row_groups[row_group_idx].num_rows is >= 0 (using CUDF_EXPECTS
with context including src_idx and row_group_idx), then convert that value to
the accumulator's unsigned type and check that sum + converted_num_rows does not
exceed std::numeric_limits<size_type>::max() (or otherwise detect overflow) and
fail early (CUDF_EXPECTS/CUDF_FAIL) if it would overflow; only then return the
safe sum.

});
});
}

std::tuple<std::vector<input_column_info>,
Expand Down
2 changes: 1 addition & 1 deletion cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base {
* @param row_group_indices Input row groups indices
* @return Total number of top-level rows in the row groups
*/
[[nodiscard]] size_type 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;

/**
Expand Down
35 changes: 18 additions & 17 deletions cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,8 +185,8 @@ void hybrid_scan_reader_impl::select_columns(read_columns_mode read_columns_mode
_is_all_columns_selected = false;
}

// Reset the rows processed so far
_rows_processed_so_far = 0;
// Reset the materialization step flag
_output_chunk_produced = false;

CUDF_EXPECTS(_input_columns.size() > 0 and _output_buffers.size() > 0, "No columns selected");

Expand All @@ -209,7 +209,7 @@ std::vector<size_type> hybrid_scan_reader_impl::all_row_groups(
return row_groups_indices;
}

size_type hybrid_scan_reader_impl::total_rows_in_row_groups(
std::size_t hybrid_scan_reader_impl::total_rows_in_row_groups(
cudf::host_span<std::vector<size_type> const> row_group_indices) const
{
return _extended_metadata->total_rows_in_row_groups(row_group_indices);
Expand Down Expand Up @@ -513,7 +513,7 @@ table_with_metadata hybrid_scan_reader_impl::materialize_filter_columns(
auto data_page_mask = thrust::host_vector<bool>{};
if (mask_data_pages == use_data_page_mask::YES) {
data_page_mask = _extended_metadata->compute_data_page_mask(
row_mask, row_group_indices, _input_columns, _rows_processed_so_far, stream);
row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
}

prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, data_page_mask);
Expand All @@ -540,7 +540,7 @@ table_with_metadata hybrid_scan_reader_impl::materialize_payload_columns(
auto data_page_mask = thrust::host_vector<bool>{};
if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) {
data_page_mask = _extended_metadata->compute_data_page_mask(
row_mask, row_group_indices, _input_columns, _rows_processed_so_far, stream);
row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
}

prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, data_page_mask);
Expand Down Expand Up @@ -568,8 +568,8 @@ table_with_metadata hybrid_scan_reader_impl::materialize_all_columns(
// Use the main reader's function
auto result = reader_impl::read_chunk_internal(read_mode::READ_ALL);

// base read_chunk_internal() does not update the _rows_processed_so_far
_rows_processed_so_far += result.tbl->num_rows();
// base read_chunk_internal() does not update the output chunk produced flag
_output_chunk_produced = true;

return result;
}
Expand Down Expand Up @@ -602,7 +602,7 @@ void hybrid_scan_reader_impl::setup_chunking_for_filter_columns(
auto data_page_mask = thrust::host_vector<bool>{};
if (mask_data_pages == use_data_page_mask::YES) {
data_page_mask = _extended_metadata->compute_data_page_mask(
row_mask, row_group_indices, _input_columns, _rows_processed_so_far, stream);
row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
}

prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, data_page_mask);
Expand Down Expand Up @@ -651,7 +651,7 @@ void hybrid_scan_reader_impl::setup_chunking_for_payload_columns(
auto data_page_mask = thrust::host_vector<bool>{};
if (not row_mask.is_empty() and mask_data_pages == use_data_page_mask::YES) {
data_page_mask = _extended_metadata->compute_data_page_mask(
row_mask, row_group_indices, _input_columns, _rows_processed_so_far, stream);
row_mask, row_group_indices, _input_columns, _row_mask_offset, stream);
}

prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, data_page_mask);
Expand Down Expand Up @@ -716,8 +716,8 @@ table_with_metadata hybrid_scan_reader_impl::materialize_all_columns_chunk()
// Use the main reader's function for reading all columns
auto result = reader_impl::read_chunk_internal(read_mode::CHUNKED_READ);

// base read_chunk_internal() does not update the _rows_processed_so_far
_rows_processed_so_far += result.tbl->num_rows();
// base read_chunk_internal() does not update the output chunk produced flag
_output_chunk_produced = true;

return result;
}
Expand Down Expand Up @@ -792,6 +792,7 @@ bool hybrid_scan_reader_impl::has_next_table_chunk()

void hybrid_scan_reader_impl::reset_internal_state()
{
_row_mask_offset = 0;
_file_itm_data = file_intermediate_data{};
_file_preprocessed = false;
_has_page_index = false;
Expand Down Expand Up @@ -1042,6 +1043,12 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output(
// 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)}; }

// Get the current row mask offset
auto const mask_offset = _row_mask_offset;
// Update the row mask offset and the output chunk produced flag
_row_mask_offset += read_table->num_rows();
_output_chunk_produced = true;

// For filter columns, apply the filter expression and update the input row mask
if constexpr (std::is_same_v<RowMaskView, cudf::mutable_column_view>) {
CUDF_EXPECTS(read_columns_mode == read_columns_mode::FILTER_COLUMNS, "Invalid read mode");
Expand All @@ -1058,9 +1065,6 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output(
auto output_table = cudf::detail::apply_mask(
read_table->view(), *final_row_mask, cudf::detail::mask_type::RETENTION, _stream, _mr);

auto const mask_offset = _rows_processed_so_far;
_rows_processed_so_far += read_table->num_rows();

// Update the input row mask to reflect the final row mask.
update_row_mask(final_row_mask->view(), row_mask, mask_offset, _stream);

Expand All @@ -1071,9 +1075,6 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output(
else {
CUDF_EXPECTS(read_columns_mode == read_columns_mode::PAYLOAD_COLUMNS, "Invalid read mode");

auto const mask_offset = _rows_processed_so_far;
_rows_processed_so_far += read_table->num_rows();

CUDF_EXPECTS(mask_offset + read_table->num_rows() <= row_mask.size(),
"Encountered invalid sized row mask to apply");
auto effective_row_mask =
Expand Down
7 changes: 4 additions & 3 deletions cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
/**
* @copydoc cudf::io::experimental::hybrid_scan::total_rows_in_row_groups
*/
[[nodiscard]] size_type 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;

/**
Expand Down Expand Up @@ -522,15 +522,16 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl {
*/
[[nodiscard]] bool is_first_output_chunk() const
{
return _file_itm_data._output_chunk_count == 0 and _rows_processed_so_far == 0;
return _file_itm_data._output_chunk_count == 0 and not _output_chunk_produced;
}

private:
aggregate_reader_metadata* _extended_metadata;

std::optional<std::vector<std::string>> _filter_columns_names;

cudf::size_type _rows_processed_so_far{0};
cudf::size_type _row_mask_offset{0};
bool _output_chunk_produced{false};

bool _use_pandas_metadata{false};

Expand Down
16 changes: 9 additions & 7 deletions cpp/src/io/parquet/experimental/page_index_filter.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -995,12 +995,14 @@ thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(

auto const total_rows = total_rows_in_row_groups(row_group_indices);

CUDF_EXPECTS(row_mask_offset + total_rows <= row_mask.size(),
"Mismatch in total rows in input row mask and row groups",
std::invalid_argument);
CUDF_EXPECTS(
std::cmp_less_equal(static_cast<std::size_t>(row_mask_offset) + total_rows, row_mask.size()),
"Encountered a mismatch in number of rows in the row group pass and the row mask size",
std::invalid_argument);

// Return an empty vector if all rows are invalid or all rows are required
if (row_mask.null_count(row_mask_offset, row_mask_offset + total_rows, stream) == total_rows or
if (std::cmp_equal(row_mask.null_count(row_mask_offset, row_mask_offset + total_rows, stream),
total_rows) or
cudf::detail::all_of(row_mask.template begin<bool>() + row_mask_offset,
row_mask.template begin<bool>() + row_mask_offset + total_rows,
cuda::std::identity{},
Expand Down Expand Up @@ -1091,8 +1093,8 @@ thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(
if constexpr (cuda::std::is_same_v<ColumnView, cudf::mutable_column_view>) {
if (row_mask.nullable() and row_mask.null_count() > 0) {
thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
cuda::counting_iterator{row_mask_offset},
cuda::counting_iterator{row_mask_offset + total_rows},
cuda::counting_iterator<cudf::size_type>(row_mask_offset),
cuda::counting_iterator<cudf::size_type>(row_mask_offset + total_rows),
[row_mask = row_mask.template begin<bool>(),
null_mask = row_mask.null_mask()] __device__(auto const row_idx) {
if (not bit_is_set(null_mask, row_idx)) { row_mask[row_idx] = true; }
Expand Down Expand Up @@ -1126,7 +1128,7 @@ thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(
cudf::detail::make_device_uvector_async(host_tree_level_ptrs, stream, mr);

// Build Fenwick tree levels (zeroth level is just the row mask itself)
auto prev_level_size = total_rows;
auto prev_level_size = static_cast<cudf::size_type>(total_rows);
std::for_each(
cuda::counting_iterator<cudf::size_type>{0},
cuda::counting_iterator{num_levels - 1},
Expand Down
Loading
Loading