This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11553-6fbb45f1ac89cc4744fe1eda23ffb07abe401e20 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 42fdb954568438f2455a79fd9dedf08d60027dd2 Author: Huaijin <[email protected]> AuthorDate: Thu Oct 15 10:40:38 2026 +0000 perf(parquet): scan bitmap selections directly for page ranges (#11553) # Which issue does this PR close? <!-- It is not necessary to raise an issue, but it is encouraged if it is a non-trivial enhancement/bug. If this PR relates to an existing issue, but does not fully resolve it, please change the wording below from "Closes #NNN" to "Relates to #NNN" or "Part of #NNN" to avoid false closing an issue. --> - Closes https://github.com/apache/arrow-rs/issues/11515 # Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. Please try to explain the problem you are trying to solve in terms of the user-visible behavior, rather than the implementation. --> Bitmap-backed `RowSelection::scan_ranges` enumerates mask runs to identify selected pages, adding overhead for fragmented selections. # What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> Check each page’s row range directly using `BitChunks`, stopping at the first nonzero 64-bit word. # Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests, if so identify these tests)? If this PR claims a performance improvement, please include evidence such as benchmark results. DO NOT simply paste a list of commands ran to verify the code; this is what our CI is for and it pollutes the PR. ONLY explain if testing was included with this PR, or why it wasn't. --> Added tests for sliced masks, padding, word/page boundaries, and loaded row ranges. Existing `arrow_reader_row_filter` benchmarks show approximately 16–22% lower async reader elapsed time in the four bitmap-backed filter/projection cases. [Benchmark results](https://github.com/apache/arrow-rs/pull/11553#issuecomment-6171902714). # Are there any significant user-facing changes? <!-- Significant user-facing changes include: - Breaking public API change - Significant change to existing functionality (e.g. function returns something different) These need to be explicitly called out to ensure we have proper documentation if required and to ensure we follow proper semantic versioning. DO NOT restate the changes of this PR; if a bug was fixed and already explained above, we don't need this to be restated as a user-facing change. --> --------- Co-authored-by: Ed Seidl <[email protected]> --- parquet/src/arrow/arrow_reader/selection/mod.rs | 8 +- parquet/src/arrow/arrow_reader/selection/ranges.rs | 131 ++++++++++++++++++++- 2 files changed, 132 insertions(+), 7 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs b/parquet/src/arrow/arrow_reader/selection/mod.rs index 31c51c3730..bd1fe2b913 100644 --- a/parquet/src/arrow/arrow_reader/selection/mod.rs +++ b/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -54,7 +54,9 @@ use boolean::{ }; pub(crate) use cursor::{LoadedRowRanges, MaskCursor, RowSelectionStrategy}; pub use cursor::{RowSelectionCursor, RowSelectionPolicy}; -use ranges::{expand_to_batch_boundaries_from_selectors, scan_ranges_from_selectors}; +use ranges::{ + expand_to_batch_boundaries_from_selectors, scan_ranges_from_mask, scan_ranges_from_selectors, +}; pub use selector::RowSelector; use selector::{limit_selectors, offset_selectors, split_off_selectors}; @@ -443,9 +445,7 @@ impl RowSelection { RowSelectionInner::Selectors(selectors) => { scan_ranges_from_selectors(selectors.iter().copied(), page_locations) } - RowSelectionInner::Mask(mask) => { - scan_ranges_from_selectors(MaskRunIter::new(mask.mask()), page_locations) - } + RowSelectionInner::Mask(mask) => scan_ranges_from_mask(mask.mask(), page_locations), } } diff --git a/parquet/src/arrow/arrow_reader/selection/ranges.rs b/parquet/src/arrow/arrow_reader/selection/ranges.rs index ac9ee17d9c..9fa65c3b9d 100644 --- a/parquet/src/arrow/arrow_reader/selection/ranges.rs +++ b/parquet/src/arrow/arrow_reader/selection/ranges.rs @@ -15,19 +15,51 @@ // specific language governing permissions and limitations // under the License. -//! Mapping the [`RowSelector`] runs of a selection onto ranges: the byte ranges +//! Mapping row selections onto ranges: the byte ranges //! of the data pages that must be fetched ([`RowSelection::scan_ranges`]) and //! the expansion of a selection to batch boundaries. //! -//! Both are shared by the selector and mask backings, which stream their runs -//! from a slice and a [`MaskRunIter`] respectively. +//! Page selection checks bitmap-backed selections directly, while batch +//! expansion streams their runs using [`MaskRunIter`]. //! //! [`MaskRunIter`]: crate::arrow::arrow_reader::MaskRunIter use super::{RowSelection, RowSelector}; use crate::file::page_index::offset_index::PageLocation; +use arrow_buffer::BooleanBuffer; +use arrow_buffer::bit_chunk_iterator::BitChunks; use std::ops::Range; +/// Checks each page for a selected bit without expanding the mask into runs. +pub(super) fn scan_ranges_from_mask( + mask: &BooleanBuffer, + page_locations: &[PageLocation], +) -> Vec<Range<u64>> { + let mut ranges = Vec::new(); + let mut row_start = 0; + for (idx, page) in page_locations.iter().enumerate() { + if row_start >= mask.len() { + break; + } + let row_end = page_locations + .get(idx + 1) + .map_or(mask.len(), |next| next.first_row_index as usize) + .min(mask.len()); + + let chunks = BitChunks::new( + mask.values(), + mask.offset() + row_start, + row_end - row_start, + ); + if chunks.iter_padded().any(|word| word != 0) { + let start = page.offset as u64; + ranges.push(start..start + page.compressed_page_size as u64); + } + row_start = row_end; + } + ranges +} + /// Byte ranges of the data pages containing at least one selected row. #[inline] pub(super) fn scan_ranges_from_selectors<I>( @@ -140,6 +172,99 @@ where mod tests { use super::*; + fn page_locations(starts: impl IntoIterator<Item = usize>) -> Vec<PageLocation> { + starts + .into_iter() + .enumerate() + .map(|(idx, row)| PageLocation { + offset: (idx * 23 + 5) as i64, + compressed_page_size: (idx % 13 + 1) as i32, + first_row_index: row as i64, + }) + .collect() + } + + fn assert_mask_scan_ranges( + mask: &BooleanBuffer, + offset: usize, + pages: &[PageLocation], + total_rows: usize, + ) { + // Build the reference from individual bits, independent of word iterators. + let selectors = RowSelection::from_consecutive_ranges( + mask.iter() + .enumerate() + .filter_map(|(row, selected)| selected.then_some(row..row + 1)), + mask.len(), + ); + // Set padding bits to ensure they cannot cause extra pages to load. + let padded = BooleanBuffer::from_iter( + std::iter::repeat_n(true, offset) + .chain(mask.iter()) + .chain(std::iter::repeat_n(true, 64)), + ); + let selection = RowSelection::from_boolean_buffer(padded.slice(offset, mask.len())); + assert_eq!( + selection.scan_ranges(pages), + selectors.scan_ranges(pages), + "mask={mask:?}, offset={offset}, pages={pages:?}" + ); + assert_eq!( + selection.row_ranges_for_selected_pages(pages, total_rows), + selectors.row_ranges_for_selected_pages(pages, total_rows) + ); + assert!(selection.scan_ranges(&[]).is_empty()); + } + + #[test] + fn test_mask_scan_ranges_matches_selectors() { + use rand::{RngExt, SeedableRng, rngs::StdRng}; + + for len in [0, 1, 7, 8, 63, 64, 65, 127, 128, 129, 511, 2049] { + for page_rows in [1, 7, 31, 63, 64, 65, 127, 257] { + // Include pages beyond the selection and a final partial page. + let total_rows = len + 37; + let pages = page_locations((0..total_rows).step_by(page_rows)); + for pattern in 0..7 { + let mut rng = StdRng::seed_from_u64(42); + let mask = BooleanBuffer::from_iter((0..len).map(|row| match pattern { + 0 => false, + 1 => true, + 2 => row % 2 == 0, + 3 => row % page_rows == page_rows - 1, + 4 => rng.random_bool(0.01), + 5 => rng.random_bool(0.5), + _ => rng.random_bool(0.99), + })); + for offset in [0, 1, 3, 7, 31, 63, 64, 65] { + assert_mask_scan_ranges(&mask, offset, &pages, total_rows); + } + } + } + } + } + + #[test] + fn test_mask_scan_ranges_single_bits_and_padding() { + let pages = page_locations([0, 1, 130, 257, 386, 512, 549]); + for len in [130, 257, 511, 512] { + let selected_rows = std::iter::once(None).chain( + [ + 0, 1, 7, 8, 55, 56, 63, 64, 65, 127, 128, 129, 130, 255, 256, 257, 510, 511, + ] + .into_iter() + .filter(|&row| row < len) + .map(Some), + ); + for selected_row in selected_rows { + let mask = BooleanBuffer::from_iter((0..len).map(|row| selected_row == Some(row))); + for offset in [0, 1, 7, 8, 31, 63, 64, 65] { + assert_mask_scan_ranges(&mask, offset, &pages, 600); + } + } + } + } + #[test] fn test_scan_ranges() { let index = vec