perf(parquet): reuse MaskSelection's cached selectors when converting to selectors

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, 32, 48, 64, 96, 128];
+
 /// 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<RowSelector> = selection.into();
+                    hint::black_box(selectors);
+                })
+            },
+        );
+
+        c.bench_with_input(
+            BenchmarkId::new("mask_consume", &label),
+            &selection,
+            |b, selection| {
+                b.iter(|| {
+                    let selectors: Vec<RowSelector> = 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);

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())
         }
     }
 }

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<RowSelector> {
+        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<RowSelector> {
+pub(super) fn mask_to_selectors(mask: &BooleanBuffer) -> Vec<RowSelector> {
     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::<Vec<bool>>())
+    }
+
+    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<RowSelector> = selection.iter().copied().collect();
+
+        let cached_ptr = cached_selectors_ptr(&selection).expect("iter populates the cache");
+        let selectors: Vec<RowSelector> = 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<RowSelector> = 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<RowSelector> = 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<RowSelector> = 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<bool> = (0..256).map(|i| i % 3 == 0).collect();
+        let other: RowSelection = RowSelection::from_filters(&[BooleanArray::from(
+            (0..256).map(|i| i % 5 != 0).collect::<Vec<bool>>(),
+        )]);
+
+        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![

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<RowSelector> {
         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)
+            }
         }
     }