This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new cf744a6c94 perf(parquet): reuse `MaskSelection`'s cached selectors
when converting to selectors (#10443)
cf744a6c94 is described below
commit cf744a6c945b0f27468dc7a2be73e799603bb2ed
Author: Huaijin <[email protected]>
AuthorDate: Tue Aug 4 01:34:26 2026 +0800
perf(parquet): reuse `MaskSelection`'s cached selectors when converting to
selectors (#10443)
# Which issue does this PR close?
- Closes #10422.
# Rationale for this change
Follow up to #10141. `RowSelection::iter()` caches the RLE form of a
mask-backed selection, but the conversions to `RowSelector`s ignored
that cache and re-encoded the bitmap, so a caller that iterates a
selection and then consumes it encodes the whole bitmap twice.
# What changes are included in this PR?
- `MaskSelection::into_selectors` takes the cached `Vec` when it was
populated and otherwise falls back to `mask_to_selectors`. Used by
`RowSelection::into_selectors_vec` (which backs `From<RowSelection>` for
`Vec<RowSelector>` and `VecDeque<RowSelector>`) and by `build_cursor`.
- `MaskSelection::borrowed_selectors` does the same for `intersection`
and `union`, returning `Cow::Borrowed` when the cache is populated and
converting into a temporary when it is not, rather than populating the
cache of a selection the caller still owns.
- `mask_to_selectors` has no callers outside the `selection` module now,
so its re-export is dropped.
# Are these changes tested?
Four new unit tests in `arrow_reader::selection::boolean` cover the
cache being moved out rather than re-encoded (asserted on the `Vec`'s
pointer), the cold-cache fallback, `borrowed_selectors` in both states,
and `intersection`/`union` agreeing either way. `cargo test -p parquet
--all-features` passes.
`row_selector` gains two scenarios over a mask-backed selection of 300k
rows: `mask_iterate_then_consume` iterates and then converts, as
`ParquetAccessPlan::into_overall_row_selection` does, and `mask_consume`
converts with a cold cache.
```
cargo bench -p parquet --bench row_selector -- mask_
```
| benchmark | main | this PR | change |
| --- | --- | --- | --- |
| `mask_iterate_then_consume/run01` | 2.237 ms | 1.202 ms | -46.5% |
| `mask_iterate_then_consume/run04` | 272.5 µs | 132.4 µs | -51.4% |
| `mask_iterate_then_consume/run16` | 73.7 µs | 32.5 µs | -55.9% |
| `mask_iterate_then_consume/run64` | 26.0 µs | 13.6 µs | -47.8% |
| `mask_iterate_then_consume/random` | 784.1 µs | 404.1 µs | -48.7% |
`mask_consume` is unchanged: its medians over three runs per side land
within 1.5% of `main`, inside the run-to-run spread `main` shows on its
own.
# Are there any user-facing changes?
No, the added methods are internal.
---
parquet/benches/row_selector.rs | 68 +++++++++++++-
parquet/src/arrow/arrow_reader/read_plan.rs | 4 +-
.../src/arrow/arrow_reader/selection/boolean.rs | 100 ++++++++++++++++++++-
parquet/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 38fb7122ab..a29a1bf852 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, 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);
diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs
b/parquet/src/arrow/arrow_reader/read_plan.rs
index 04f132e1bc..92ebbe9bae 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 2ac879eba9..bf6c983d1c 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<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