From 438e378ab893920cac718438ffbf89838fb610b7 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sun, 26 Jul 2026 23:35:28 +0800 Subject: [PATCH 1/2] perf(parquet): reuse `MaskSelection`'s cached selectors when converting to selectors `RowSelection::iter()` lazily materializes and caches the RLE form of a mask-backed selection, but every path that converted such a selection into `RowSelector`s ignored that cache and re-encoded the bitmap. Add `MaskSelection::into_selectors`, which takes the `OnceLock` value when it was populated and otherwise falls back to `mask_to_selectors`, and use it from `RowSelection::into_selectors_vec` (backing `From` for both `Vec` and `VecDeque`) and from `build_cursor` when lowering a mask-backed selection to the selector cursor. `intersection`/`union` reuse the cache too, via `borrowed_selectors`, which borrows it when populated and otherwise converts into a temporary rather than populating the cache of a selection the caller still owns. `mask_to_selectors` no longer has callers outside the `selection` module, so drop its `pub(crate)` re-export. Benchmarked with the new `mask_iterate_then_consume` /`mask_consume` cases in `row_selector`, which cover a caller that iterates a selection before consuming it, as `ParquetAccessPlan::into_overall_row_selection` does. Over 300k rows the iterate-then-consume conversion drops 46-56% across run lengths, and the cold-cache conversion is unchanged. --- parquet/benches/row_selector.rs | 68 +++++++++++- parquet/src/arrow/arrow_reader/read_plan.rs | 4 +- .../arrow/arrow_reader/selection/boolean.rs | 100 +++++++++++++++++- .../src/arrow/arrow_reader/selection/mod.rs | 39 +++---- 4 files changed, 184 insertions(+), 27 deletions(-) diff --git a/parquet/benches/row_selector.rs b/parquet/benches/row_selector.rs index 38fb7122ab5f..95c30eda6469 100644 --- a/parquet/benches/row_selector.rs +++ b/parquet/benches/row_selector.rs @@ -16,11 +16,16 @@ // under the License. use arrow_array::BooleanArray; +use arrow_buffer::BooleanBuffer; use criterion::*; -use parquet::arrow::arrow_reader::RowSelection; +use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; use rand::Rng; use std::hint; +/// Run lengths for the mask conversion benchmarks. Shorter runs mean more +/// [`RowSelector`]s per row, so the RLE encoding dominates. +const MASK_RUN_LENGTHS: &[usize] = &[1, 4, 16, 64]; + /// Generates a random RowSelection with a specified selection ratio. /// /// # Arguments @@ -39,6 +44,65 @@ fn generate_random_row_selection(total_rows: usize, selection_ratio: f64) -> Boo BooleanArray::from(bools) } +/// Generates a mask alternating between selected and skipped runs of `run_len` rows. +fn generate_run_length_mask(total_rows: usize, run_len: usize) -> BooleanBuffer { + BooleanBuffer::from_iter((0..total_rows).map(|row| (row / run_len) % 2 == 0)) +} + +/// Benchmarks converting a mask-backed [`RowSelection`] into [`RowSelector`]s. +/// +/// `RowSelection::iter` caches the RLE form, so a caller that iterates before +/// consuming should reuse that cache rather than encode the bitmap twice. +/// `mask_consume` is the same conversion with a cold cache. +fn bench_mask_backed_conversion(c: &mut Criterion, total_rows: usize, selection_ratio: f64) { + let mut cases: Vec<(String, BooleanBuffer)> = MASK_RUN_LENGTHS + .iter() + .map(|&run_len| { + ( + format!("run{run_len:02}"), + generate_run_length_mask(total_rows, run_len), + ) + }) + .collect(); + cases.push(( + "random".to_string(), + generate_random_row_selection(total_rows, selection_ratio) + .values() + .clone(), + )); + + for (label, mask) in cases { + let selection = RowSelection::from_boolean_buffer(mask); + + c.bench_with_input( + BenchmarkId::new("mask_iterate_then_consume", &label), + &selection, + |b, selection| { + b.iter(|| { + // `clone` drops the selector cache, so each iteration + // starts from an unconverted selection. + let selection = selection.clone(); + let rows: usize = selection.iter().map(|s| s.row_count).sum(); + hint::black_box(rows); + let selectors: Vec = selection.into(); + hint::black_box(selectors); + }) + }, + ); + + c.bench_with_input( + BenchmarkId::new("mask_consume", &label), + &selection, + |b, selection| { + b.iter(|| { + let selectors: Vec = selection.clone().into(); + hint::black_box(selectors); + }) + }, + ); + } +} + fn criterion_benchmark(c: &mut Criterion) { let total_rows = 300_000; let selection_ratio = 1.0 / 3.0; @@ -82,6 +146,8 @@ fn criterion_benchmark(c: &mut Criterion) { hint::black_box(result); }) }); + + bench_mask_backed_conversion(c, total_rows, selection_ratio); } criterion_group!(benches, criterion_benchmark); diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 04f132e1bc1d..92ebbe9baef8 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -20,7 +20,7 @@ use crate::arrow::array_reader::ArrayReader; use crate::arrow::arrow_reader::selection::{ - LoadedRowRanges, RowSelectionInner, RowSelectionPolicy, RowSelectionStrategy, mask_to_selectors, + LoadedRowRanges, RowSelectionInner, RowSelectionPolicy, RowSelectionStrategy, }; use crate::arrow::arrow_reader::{ ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelectionCursor, RowSelector, @@ -335,7 +335,7 @@ fn build_cursor( RowSelectionCursor::new_selectors(selectors) } (RowSelectionStrategy::Selectors, RowSelectionInner::Mask(mask)) => { - RowSelectionCursor::new_selectors(mask_to_selectors(mask.mask())) + RowSelectionCursor::new_selectors((*mask).into_selectors()) } } } diff --git a/parquet/src/arrow/arrow_reader/selection/boolean.rs b/parquet/src/arrow/arrow_reader/selection/boolean.rs index 2ac879eba9dc..bf6c983d1c4a 100644 --- a/parquet/src/arrow/arrow_reader/selection/boolean.rs +++ b/parquet/src/arrow/arrow_reader/selection/boolean.rs @@ -26,6 +26,7 @@ use super::RowSelector; use arrow_buffer::bit_iterator::BitSliceIterator; use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, Buffer}; +use std::borrow::Cow; use std::sync::OnceLock; /// Mask-backed [`RowSelection`] storage. @@ -92,6 +93,22 @@ impl MaskSelection { .get_or_init(|| mask_to_selectors(&self.mask)) .as_slice() } + + /// Borrows the cached RLE form, converting into a temporary if not cached. + pub(super) fn borrowed_selectors(&self) -> Cow<'_, [RowSelector]> { + match self.selectors.get() { + Some(selectors) => Cow::Borrowed(selectors.as_slice()), + None => Cow::Owned(mask_to_selectors(&self.mask)), + } + } + + /// The RLE form, taking the cache if it was populated. + pub(crate) fn into_selectors(self) -> Vec { + match self.selectors.into_inner() { + Some(selectors) => selectors, + None => mask_to_selectors(&self.mask), + } + } } impl Clone for MaskSelection { @@ -179,7 +196,7 @@ impl Iterator for MaskRunIter<'_> { } /// Materialize a [`BooleanBuffer`] into its RLE form. -pub(crate) fn mask_to_selectors(mask: &BooleanBuffer) -> Vec { +pub(super) fn mask_to_selectors(mask: &BooleanBuffer) -> Vec { let total_rows = mask.len(); if total_rows == 0 { return Vec::new(); @@ -391,6 +408,87 @@ mod tests { ); } + /// Enough runs that the RLE form is a real allocation, so the cache reuse + /// tests can track its pointer across the conversion. + fn interleaved_mask() -> BooleanBuffer { + BooleanBuffer::from((0..256).map(|i| i % 3 == 0).collect::>()) + } + + fn cached_selectors_ptr(selection: &RowSelection) -> Option<*const RowSelector> { + match &selection.inner { + RowSelectionInner::Mask(m) => m.selectors.get().map(|s| s.as_ptr()), + _ => unreachable!(), + } + } + + #[test] + fn test_into_selectors_takes_the_iter_cache() { + let selection = RowSelection::from_boolean_buffer(interleaved_mask()); + let expected: Vec = selection.iter().copied().collect(); + + let cached_ptr = cached_selectors_ptr(&selection).expect("iter populates the cache"); + let selectors: Vec = selection.into(); + + assert_eq!(selectors, expected); + // Moved out of the cache rather than re-encoded from the bitmap. + assert_eq!(selectors.as_ptr(), cached_ptr); + } + + #[test] + fn test_into_selectors_without_cache_still_converts() { + let selection = RowSelection::from_boolean_buffer(interleaved_mask()); + assert!(cached_selectors_ptr(&selection).is_none()); + + let selectors: Vec = selection.into(); + assert_eq!(selectors, mask_to_selectors(&interleaved_mask())); + + // `VecDeque` goes through the same path. + let selection = RowSelection::from_boolean_buffer(interleaved_mask()); + let _ = selection.iter().count(); + let deque: std::collections::VecDeque = selection.into(); + assert_eq!(Vec::from(deque), selectors); + } + + #[test] + fn test_borrowed_selectors_reuses_cache_without_populating_it() { + let selection = RowSelection::from_boolean_buffer(interleaved_mask()); + let mask = match &selection.inner { + RowSelectionInner::Mask(m) => m, + _ => unreachable!(), + }; + + // Uncached: converts into a temporary, leaving the cache empty. + assert!(matches!(mask.borrowed_selectors(), Cow::Owned(_))); + assert!(mask.selectors.get().is_none()); + + let expected: Vec = selection.iter().copied().collect(); + let mask = match &selection.inner { + RowSelectionInner::Mask(m) => m, + _ => unreachable!(), + }; + match mask.borrowed_selectors() { + Cow::Borrowed(selectors) => assert_eq!(selectors, expected.as_slice()), + Cow::Owned(_) => panic!("expected the cached selectors to be reused"), + } + } + + #[test] + fn test_set_algebra_agrees_whether_or_not_the_cache_is_populated() { + let bits: Vec = (0..256).map(|i| i % 3 == 0).collect(); + let other: RowSelection = RowSelection::from_filters(&[BooleanArray::from( + (0..256).map(|i| i % 5 != 0).collect::>(), + )]); + + let cold = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())); + let warm = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits)); + let _ = warm.iter().count(); + + assert_eq!(cold.intersection(&other), warm.intersection(&other)); + assert_eq!(other.intersection(&cold), other.intersection(&warm)); + assert_eq!(cold.union(&other), warm.union(&other)); + assert_eq!(other.union(&cold), other.union(&warm)); + } + #[test] fn test_mask_run_iter_streams_without_cache() { let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs b/parquet/src/arrow/arrow_reader/selection/mod.rs index b8a5bc0f2a74..962f2ed1fe80 100644 --- a/parquet/src/arrow/arrow_reader/selection/mod.rs +++ b/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -48,7 +48,6 @@ use algebra::{ intersect_row_selections, union_masks, union_row_selections, }; pub use boolean::MaskRunIter; -pub(crate) use boolean::mask_to_selectors; use boolean::{ MaskSelection, limit_mask, mask_has_at_least_runs, offset_mask, split_off_mask, trim_mask, }; @@ -300,7 +299,7 @@ impl RowSelection { fn into_selectors_vec(self) -> Vec { match self.inner { RowSelectionInner::Selectors(s) => s, - RowSelectionInner::Mask(m) => mask_to_selectors(m.mask()), + RowSelectionInner::Mask(m) => (*m).into_selectors(), } } @@ -487,12 +486,10 @@ impl RowSelection { intersect_row_selections(l, r) } (RowSelectionInner::Selectors(l), RowSelectionInner::Mask(r)) => { - let r = mask_to_selectors(r.mask()); - intersect_row_selections(l, &r) + intersect_row_selections(l, &r.borrowed_selectors()) } (RowSelectionInner::Mask(l), RowSelectionInner::Selectors(r)) => { - let l = mask_to_selectors(l.mask()); - intersect_row_selections(&l, r) + intersect_row_selections(&l.borrowed_selectors(), r) } } } @@ -504,23 +501,19 @@ impl RowSelection { /// /// returned: NYYYYYNNYYNYN pub fn union(&self, other: &Self) -> Self { - match &self.inner { - RowSelectionInner::Mask(l) => match &other.inner { - RowSelectionInner::Mask(r) => { - Self::from_boolean_buffer(union_masks(l.mask(), r.mask())) - } - RowSelectionInner::Selectors(r) => { - let l = mask_to_selectors(l.mask()); - union_row_selections(&l, r) - } - }, - RowSelectionInner::Selectors(l) => match &other.inner { - RowSelectionInner::Mask(r) => { - let r = mask_to_selectors(r.mask()); - union_row_selections(l, &r) - } - RowSelectionInner::Selectors(r) => union_row_selections(l, r), - }, + match (&self.inner, &other.inner) { + (RowSelectionInner::Mask(l), RowSelectionInner::Mask(r)) => { + Self::from_boolean_buffer(union_masks(l.mask(), r.mask())) + } + (RowSelectionInner::Selectors(l), RowSelectionInner::Selectors(r)) => { + union_row_selections(l, r) + } + (RowSelectionInner::Selectors(l), RowSelectionInner::Mask(r)) => { + union_row_selections(l, &r.borrowed_selectors()) + } + (RowSelectionInner::Mask(l), RowSelectionInner::Selectors(r)) => { + union_row_selections(&l.borrowed_selectors(), r) + } } } From e5f4ec70fefc0db102c7a7ef77eed69d806a300e Mon Sep 17 00:00:00 2001 From: haohuaijin Date: Mon, 27 Jul 2026 16:14:49 +0800 Subject: [PATCH 2/2] update --- parquet/benches/row_selector.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parquet/benches/row_selector.rs b/parquet/benches/row_selector.rs index 95c30eda6469..a29a1bf852eb 100644 --- a/parquet/benches/row_selector.rs +++ b/parquet/benches/row_selector.rs @@ -24,7 +24,7 @@ use std::hint; /// Run lengths for the mask conversion benchmarks. Shorter runs mean more /// [`RowSelector`]s per row, so the RLE encoding dominates. -const MASK_RUN_LENGTHS: &[usize] = &[1, 4, 16, 64]; +const MASK_RUN_LENGTHS: &[usize] = &[1, 4, 16, 32, 48, 64, 96, 128]; /// Generates a random RowSelection with a specified selection ratio. ///