-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Add Hybrid scan page pruning when offset index is absent #23731
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.
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
base: main
Are you sure you want to change the base?
Changes from all commits
e32363f
213f35e
1fd4b66
ccc96b6
4273fa9
2da7634
aa7be76
ff88b49
b267d0b
cc6a199
4f786ef
aac2bbb
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -9,6 +9,7 @@ | |
| #include "hybrid_scan_helpers.hpp" | ||
| #include "io/parquet/reader_impl_chunking_utils.cuh" | ||
| #include "io/parquet/synthetic_column_helpers.hpp" | ||
| #include "page_index_filter_utils.hpp" | ||
|
|
||
| #include <cudf/copying.hpp> | ||
| #include <cudf/detail/stream_compaction.hpp> | ||
|
|
@@ -655,8 +656,8 @@ hybrid_scan_reader_impl::payload_pages_byte_ranges( | |
|
|
||
| // Compute the data page mask | ||
| auto const mask_size = mask_offsets.back(); | ||
| auto data_page_mask = _extended_metadata->compute_data_page_mask( | ||
| row_mask, row_group_indices, _input_columns, 0, stream); | ||
| auto data_page_mask = | ||
| _extended_metadata->compute_data_page_mask(row_mask, row_group_indices, _input_columns, stream); | ||
| CUDF_EXPECTS(data_page_mask.empty() or data_page_mask.size() == mask_size, | ||
| "Computed data page mask does not match offset indexes"); | ||
|
|
||
|
|
@@ -757,8 +758,9 @@ 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) { | ||
| _row_mask = set_nulls_to_true(row_mask, stream); | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Set nulls to true (keep row) for mutable row masks passed to filter column materializers. Payload column materializers only take in non-nullable row masks
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. wouldn't it be faster to not incur this cost of copying and modifying the mask, but instead make the downstream function interpret nulls as true? |
||
| data_page_mask = _extended_metadata->compute_data_page_mask( | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. No need to pass in |
||
| row_mask, row_group_indices, _input_columns, _row_mask_offset, stream); | ||
| _row_mask, row_group_indices, _input_columns, stream); | ||
| } | ||
|
|
||
| prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, data_page_mask); | ||
|
|
@@ -795,8 +797,9 @@ 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) { | ||
| _row_mask = row_mask; | ||
| data_page_mask = _extended_metadata->compute_data_page_mask( | ||
| row_mask, row_group_indices, _input_columns, _row_mask_offset, stream); | ||
| _row_mask, row_group_indices, _input_columns, stream); | ||
| } | ||
|
|
||
| prepare_data(read_mode::READ_ALL, row_group_indices, column_chunk_data, data_page_mask); | ||
|
|
@@ -834,7 +837,7 @@ void hybrid_scan_reader_impl::setup_chunking_for_filter_columns( | |
| std::size_t chunk_read_limit, | ||
| std::size_t pass_read_limit, | ||
| std::span<std::vector<size_type> const> row_group_indices, | ||
| cudf::column_view const& row_mask, | ||
| cudf::mutable_column_view const& row_mask, | ||
| use_data_page_mask mask_data_pages, | ||
| std::span<cudf::device_span<uint8_t const> const> column_chunk_data, | ||
| parquet_reader_options const& options, | ||
|
|
@@ -866,8 +869,9 @@ 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) { | ||
| _row_mask = set_nulls_to_true(row_mask, stream); | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Set nulls to true (keep row) for mutable row masks passed to filter column materializers. Payload column materializers only take in non-nullable row masks |
||
| data_page_mask = _extended_metadata->compute_data_page_mask( | ||
| row_mask, row_group_indices, _input_columns, _row_mask_offset, stream); | ||
| _row_mask, row_group_indices, _input_columns, stream); | ||
| } | ||
|
|
||
| prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, data_page_mask); | ||
|
|
@@ -926,8 +930,9 @@ 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) { | ||
| _row_mask = row_mask; | ||
| data_page_mask = _extended_metadata->compute_data_page_mask( | ||
| row_mask, row_group_indices, _input_columns, _row_mask_offset, stream); | ||
| _row_mask, row_group_indices, _input_columns, stream); | ||
| } | ||
|
|
||
| prepare_data(read_mode::CHUNKED_READ, row_group_indices, column_chunk_data, data_page_mask); | ||
|
|
@@ -1135,7 +1140,6 @@ 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_offset_index = false; | ||
|
|
@@ -1159,6 +1163,10 @@ void hybrid_scan_reader_impl::reset_internal_state() | |
| _output_chunk_read_limit = 0; | ||
| _strings_to_categorical = false; | ||
| _reader_column_schema.reset(); | ||
|
|
||
| _row_mask = column_view{}; | ||
| _row_mask_offset = 0; | ||
|
|
||
| _expr_conv = parquet_filter_normalizer{}; | ||
| _mr = cudf::get_current_device_resource_ref(); | ||
| } | ||
|
|
@@ -1380,9 +1388,9 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( | |
| // Prepend the source and row index columns to filter columns only | ||
| if (read_columns_mode == read_columns_mode::FILTER_COLUMNS) { | ||
| if (_options.prepend_row_index_column) { | ||
| out_columns.emplace( | ||
| out_columns.begin(), | ||
| synthesize_row_index_column(_file_itm_data.row_groups, read_info, _stream, _mr)); | ||
| out_columns.emplace(out_columns.begin(), | ||
| parquet::detail::synthesize_row_index_column( | ||
| _file_itm_data.row_groups, read_info, _stream, _mr)); | ||
| out_metadata.schema_info.emplace(out_metadata.schema_info.begin(), | ||
| column_name_info{.name = "row_index", .is_nullable = false}); | ||
| } | ||
|
|
@@ -1506,6 +1514,73 @@ void hybrid_scan_reader_impl::set_pass_page_mask(std::span<bool const> data_page | |
| mark_buffers_nullable_for_pruned_pages(); | ||
| } | ||
|
|
||
| thrust::host_vector<bool> hybrid_scan_reader_impl::compute_data_page_mask_with_page_headers() | ||
|
vuule marked this conversation as resolved.
|
||
| { | ||
| auto const& pass = *_pass_itm_data; | ||
|
|
||
| // Return an empty vector if all rows are required | ||
| if (are_all_rows_retained(_row_mask, _stream)) { return thrust::host_vector<bool>(0); } | ||
|
|
||
| std::vector<cudf::size_type> page_row_offsets; | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This function is actually simpler than it looks. We are essentially doing the same thing as in Go over all pages and:
Call the |
||
| page_row_offsets.reserve(pass.pages.size() * 2); | ||
|
|
||
| // Maps each data page to its flat-page range; -1 keeps nested pages enabled. | ||
| std::vector<cudf::size_type> row_range_map; | ||
| row_range_map.reserve(pass.pages.size()); | ||
|
|
||
| cudf::size_type previous_chunk_idx = -1; | ||
| auto max_page_size = cudf::size_type{0}; | ||
|
|
||
| for (auto const& page : pass.pages) { | ||
| // Ignore dictionary pages altogether | ||
| if (page.flags & parquet::detail::PAGEINFO_FLAGS_DICTIONARY) { continue; } | ||
|
|
||
| auto const& chunk = pass.chunks[page.chunk_idx]; | ||
|
|
||
| // Don't prune list column pages as rows may span page boundaries when offset index isn't | ||
| // present. | ||
| if (chunk.max_level[parquet::detail::level_type::REPETITION] > 0) { | ||
| row_range_map.push_back(-1); | ||
| continue; | ||
| } | ||
|
|
||
| auto const page_start = chunk.start_row + page.chunk_row; | ||
| auto const page_end = page_start + page.num_rows; | ||
| max_page_size = std::max<cudf::size_type>(max_page_size, page_end - page_start); | ||
|
|
||
| // Starting a new column chunk. Push page start row | ||
| if (page.chunk_idx != previous_chunk_idx) { | ||
| page_row_offsets.push_back(page_start); | ||
| previous_chunk_idx = page.chunk_idx; | ||
| } | ||
|
|
||
| // Push row range index and page end row | ||
| row_range_map.push_back(page_row_offsets.size() - 1); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. could just do page_row_offsets.back() |
||
| page_row_offsets.push_back(page_end); | ||
| } | ||
|
|
||
| auto data_page_mask = thrust::host_vector<bool>{}; | ||
|
|
||
| // Compute the row range mask | ||
| CUDF_EXPECTS(std::cmp_equal(_row_mask.size(), pass.num_rows), | ||
| "Row mask must span across all rows in the pass"); | ||
| auto const row_range_mask = | ||
| compute_row_range_selection_mask(_row_mask, page_row_offsets, max_page_size, _stream); | ||
|
|
||
| if (row_range_mask.empty()) { return data_page_mask; } | ||
|
|
||
| CUDF_EXPECTS(row_range_mask.size() == page_row_offsets.size() - 1, | ||
| "Encountered invalid row range mask size"); | ||
|
|
||
| data_page_mask.reserve(row_range_map.size()); | ||
|
|
||
| // Scatter row range results while retaining list column pages. | ||
| for (auto const range_idx : row_range_map) { | ||
| data_page_mask.push_back(range_idx < 0 ? true : row_range_mask[range_idx]); | ||
| } | ||
| return data_page_mask; | ||
| } | ||
|
|
||
| void hybrid_scan_reader_impl::set_sparse_pass_page_mask( | ||
| std::span<cudf::device_span<uint8_t const> const> page_data) | ||
| { | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -253,7 +253,7 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { | |
| std::size_t chunk_read_limit, | ||
| std::size_t pass_read_limit, | ||
| std::span<std::vector<size_type> const> row_group_indices, | ||
| cudf::column_view const& row_mask, | ||
| cudf::mutable_column_view const& row_mask, | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Take in mutable column view |
||
| use_data_page_mask mask_data_pages, | ||
| std::span<cudf::device_span<uint8_t const> const> column_chunk_data, | ||
| parquet_reader_options const& options, | ||
|
|
@@ -394,6 +394,11 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { | |
| */ | ||
| void set_sparse_pass_page_mask(std::span<cudf::device_span<uint8_t const> const> page_data); | ||
|
|
||
| /** | ||
| * @brief Compute a data page mask from the decoded page headers. | ||
| */ | ||
| [[nodiscard]] thrust::host_vector<bool> compute_data_page_mask_with_page_headers(); | ||
|
|
||
| /** | ||
| * @brief Mark output buffers nullable when page pruning synthesizes null rows | ||
| */ | ||
|
|
@@ -626,6 +631,8 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { | |
|
|
||
| std::optional<std::vector<std::string>> _filter_columns_names; | ||
|
|
||
| cudf::column_view _row_mask{}; | ||
|
|
||
| std::vector<cudf::io::detail::inline_column_buffer> _original_output_buffers_template; | ||
|
|
||
| cudf::size_type _row_mask_offset{0}; | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Breaking change: Make this mutable now