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 all commits
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
37 changes: 21 additions & 16 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 Down Expand Up @@ -381,6 +381,10 @@ std::unique_ptr<cudf::column> hybrid_scan_reader_impl::build_all_true_row_mask(
CUDF_EXPECTS(not row_group_indices.empty(), "Empty input row group indices encountered");

auto const num_rows = total_rows_in_row_groups(row_group_indices);
CUDF_EXPECTS(num_rows < std::numeric_limits<cudf::size_type>::max(),
"Total rows in row groups exceed the cudf's column size limit. Retry with a smaller "
"set of row groups",
std::invalid_argument);
auto true_scalar =
cudf::numeric_scalar<bool>(true, true, stream, cudf::get_current_device_resource_ref());
return cudf::make_column_from_scalar(true_scalar, num_rows, stream, mr);
Comment on lines 383 to 390

@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 | 🟡 Minor | ⚡ Quick win

Allow the max-sized pass through this guard.

Line 387 rejects num_rows == std::numeric_limits<cudf::size_type>::max(), even though that is still a valid column size. That turns the exact boundary case into a false overflow and needlessly breaks a max-sized pass.

Suggested fix 
-  CUDF_EXPECTS(num_rows < std::numeric_limits<cudf::size_type>::max(),
+  CUDF_EXPECTS(num_rows <= static_cast<std::size_t>(std::numeric_limits<cudf::size_type>::max()),
                "Total rows in row groups exceed the cudf's column size limit. Retry with a smaller "
                "set of row groups",
                std::invalid_argument);
@@
-  return cudf::make_column_from_scalar(true_scalar, num_rows, stream, mr);
+  return cudf::make_column_from_scalar(
+    true_scalar, static_cast<cudf::size_type>(num_rows), stream, mr);
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
auto const num_rows = total_rows_in_row_groups(row_group_indices);
CUDF_EXPECTS(num_rows < std::numeric_limits<cudf::size_type>::max(),
"Total rows in row groups exceed the cudf's column size limit. Retry with a smaller "
"set of row groups",
std::invalid_argument);
auto true_scalar =
cudf::numeric_scalar<bool>(true, true, stream, cudf::get_current_device_resource_ref());
return cudf::make_column_from_scalar(true_scalar, num_rows, stream, mr);
auto const num_rows = total_rows_in_row_groups(row_group_indices);
CUDF_EXPECTS(num_rows <= static_cast<std::size_t>(std::numeric_limits<cudf::size_type>::max()),
"Total rows in row groups exceed the cudf's column size limit. Retry with a smaller "
"set of row groups",
std::invalid_argument);
auto true_scalar =
cudf::numeric_scalar<bool>(true, true, stream, cudf::get_current_device_resource_ref());
return cudf::make_column_from_scalar(
true_scalar, static_cast<cudf::size_type>(num_rows), stream, mr);
🤖 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 386 - 393,
The guard in hybrid_scan_impl.cpp incorrectly rejects the boundary case by
checking num_rows < std::numeric_limits<cudf::size_type>::max(); update the
CUDF_EXPECTS so it only fails when num_rows is strictly greater than the max
(i.e., allow equality) — locate the check around the
total_rows_in_row_groups(row_group_indices) computation and change the
comparison to permit num_rows == std::numeric_limits<cudf::size_type>::max(),
leaving the error message and exception type intact; no other logic (true_scalar
creation or make_column_from_scalar call) needs to change.

Expand Down Expand Up @@ -510,7 +514,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 @@ -537,7 +541,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 @@ -565,8 +569,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 @@ -599,7 +603,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 @@ -648,7 +652,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 @@ -713,8 +717,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 @@ -789,6 +793,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 @@ -1039,6 +1044,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 @@ -1055,9 +1066,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 @@ -1068,9 +1076,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
5 changes: 3 additions & 2 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 @@ -523,15 +523,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
27 changes: 11 additions & 16 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 @@ -870,20 +870,11 @@ std::unique_ptr<cudf::column> aggregate_reader_metadata::build_row_mask_with_pag
std::runtime_error);

// Total number of rows
auto const total_rows = std::accumulate(
cuda::counting_iterator<std::size_t>{0},
cuda::counting_iterator{row_group_indices.size()},
std::size_t{0},
[&](auto sum, auto const src_index) {
auto const& rg_indices = row_group_indices[src_index];
return std::accumulate(
rg_indices.begin(), rg_indices.end(), sum, [&](auto subsum, auto const rg_index) {
CUDF_EXPECTS(subsum + per_file_metadata[src_index].row_groups[rg_index].num_rows <=
std::numeric_limits<size_type>::max(),
"Total rows exceed the maximum value");
return subsum + per_file_metadata[src_index].row_groups[rg_index].num_rows;
});
});
auto const total_rows = total_rows_in_row_groups(row_group_indices);
CUDF_EXPECTS(std::cmp_less_equal(total_rows, std::numeric_limits<size_type>::max()),
"Total rows in row groups exceed the cudf's column size limit. Retry with a smaller "
"set of row groups",
std::invalid_argument);

auto const num_columns = output_dtypes.size();

Expand Down Expand Up @@ -994,11 +985,15 @@ thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(
"Input row bitmask should be of type BOOL8");

auto const total_rows = total_rows_in_row_groups(row_group_indices);
CUDF_EXPECTS(std::cmp_less_equal(total_rows, std::numeric_limits<size_type>::max()),
"Total rows in row groups exceed the cudf's column size limit. Retry with a smaller "
"set of row groups",
std::invalid_argument);

CUDF_EXPECTS(
std::cmp_less_equal(static_cast<std::size_t>(row_mask_offset) + total_rows, row_mask.size()),
std::cmp_less_equal(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);
std::overflow_error);

// Return an empty vector if all rows are invalid or all rows are required
if (std::cmp_equal(row_mask.null_count(row_mask_offset, row_mask_offset + total_rows, stream),
Expand Down
Loading
Loading