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-11390-e17d04788fcb181bc207d77ef3c9d26640892fba in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit a8ee70c9c4d6f4a7f46884007cdbc22fc8757b77 Author: Huaijin <[email protected]> AuthorDate: Fri Oct 9 14:24:29 2026 +0000 perf: Avoid expanding masks to selectors in mixed `RowSelection` intersection and union (#11390) # Which issue does this PR close? Closes #10423. # Rationale for this change Mixed mask/selector `intersection` and `union` currently expand the bitmap into RLE before merging. An alternating 3M-row bitmap can produce 3M selectors, turning roughly 375 KB into 48 MB of temporary storage. # What changes are included in this PR? Convert the selector operand to a bitmap and reuse the existing bitwise operations. Nonempty mixed inputs return a mask-backed result without materializing the mask's RLE representation. Empty operands preserve the other operand's representation, and unequal lengths retain the existing tail behavior. Remove the unused `borrowed_selectors` helper. Benchmark coverage is submitted separately in #11389 and should land first. # Are these changes tested? All 83 selection tests pass, including mixed-input equivalence, offsets, unequal lengths, empty operands, cached inputs, and result representation. Formatting and Clippy checks pass. Representative local Criterion results on Apple M2 Max, in µs (main → this change), using the equal-length, uncached inputs from #11389: | Case | Intersection | Union | | --- | ---: | ---: | | `alternating_1m` | 5870.78 → 6.63 | 5892.17 → 6.65 | | `alternating_3m` | 19495.49 → 16.83 | 19666.48 → 16.84 | | `compact_1m` | 13.70 → 7.11 | 13.62 → 7.21 | | `fragmented_selectors` | 362.98 → 110.84 | 338.59 → 111.61 | | `chain_mixed_3` | 5756.80 → 18.08 | 5038.74 → 18.11 | Chain timings include all three operations. These measurements exclude input construction and include result allocation and destruction. # Are there any significant user-facing changes? No API signature changes. Mixed nonempty results use bitmap storage even when the input mask already has cached selectors. --------- Co-authored-by: Andrew Lamb <[email protected]> --- .../src/arrow/arrow_reader/selection/algebra.rs | 114 ++++++++++++++++++++- .../src/arrow/arrow_reader/selection/boolean.rs | 55 +++++----- parquet/src/arrow/arrow_reader/selection/mod.rs | 28 ++--- 3 files changed, 153 insertions(+), 44 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/selection/algebra.rs b/parquet/src/arrow/arrow_reader/selection/algebra.rs index 0f9b445fa2..6e16ead75d 100644 --- a/parquet/src/arrow/arrow_reader/selection/algebra.rs +++ b/parquet/src/arrow/arrow_reader/selection/algebra.rs @@ -22,6 +22,7 @@ //! operands: a merge of the [`RowSelector`] runs, and a bitwise variant over //! [`BooleanBuffer`] masks. +use super::boolean::{MaskSelection, boolean_mask_from_selectors}; use super::{MaskRunIter, RowSelection, RowSelectionInner, RowSelector}; use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, MutableBuffer, bit_util}; use std::cmp::Ordering; @@ -276,6 +277,23 @@ pub(super) fn union_row_selections(left: &[RowSelector], right: &[RowSelector]) iter.collect() } +/// Combines a mask and selectors without expanding the mask into runs. +pub(super) fn combine_mixed_selection( + selectors: &[RowSelector], + mask: &MaskSelection, + op: impl FnOnce(&BooleanBuffer, &BooleanBuffer) -> BooleanBuffer, +) -> RowSelection { + // An empty operand leaves the other unchanged. + if selectors.is_empty() { + return RowSelection::from_mask_selection(mask.clone()); + } + if mask.mask().is_empty() { + return RowSelection::from_selectors(selectors.to_vec()); + } + let selectors = boolean_mask_from_selectors(selectors); + RowSelection::from_boolean_buffer(op(mask.mask(), &selectors)) +} + /// Bitwise AND of two mask-backed selections. Longer side's tail passes through. pub(super) fn intersect_masks(l: &BooleanBuffer, r: &BooleanBuffer) -> BooleanBuffer { if l.len() == r.len() { @@ -974,8 +992,10 @@ mod tests { .collect(); let expected_intersection = RowSelection::from_filters(&[BooleanArray::from(intersection_bits)]); - assert_eq!(mask.intersection(&selectors), expected_intersection); - assert_eq!(selectors.intersection(&mask), expected_intersection); + for result in [mask.intersection(&selectors), selectors.intersection(&mask)] { + assert!(result.as_mask().is_some()); + assert_eq!(result, expected_intersection); + } let union_bits: Vec<_> = mask_bits .iter() @@ -983,8 +1003,10 @@ mod tests { .map(|(x, y)| *x || *y) .collect(); let expected_union = RowSelection::from_filters(&[BooleanArray::from(union_bits)]); - assert_eq!(mask.union(&selectors), expected_union); - assert_eq!(selectors.union(&mask), expected_union); + for result in [mask.union(&selectors), selectors.union(&mask)] { + assert!(result.as_mask().is_some()); + assert_eq!(result, expected_union); + } } #[test] @@ -1052,6 +1074,90 @@ mod tests { assert_eq!(actual, expected, "{context}"); } + #[test] + fn test_mixed_mask_selector_algebra_with_offsets_and_uneven_lengths() { + for rows in [1, 7, 8, 63, 64, 65, 137] { + for (mask_len, selector_len) in [(rows, rows), (rows, 137), (137, rows)] { + for offset in [0, 1, 7, 63] { + for mask_bits in [ + vec![false; mask_len], + vec![true; mask_len], + (0..mask_len).map(|i| i < mask_len / 2).collect(), + (0..mask_len).map(|i| i % 2 == 0).collect(), + ] { + let buffer = BooleanBuffer::from_iter( + std::iter::repeat_n(false, offset).chain(mask_bits.iter().copied()), + ) + .slice(offset, mask_len); + let mask = RowSelection::from_boolean_buffer(buffer); + for selector_bits in [ + vec![false; selector_len], + vec![true; selector_len], + (0..selector_len).map(|i| i % 2 == 0).collect(), + (0..selector_len).map(|i| (i / 11) % 2 == 0).collect(), + ] { + let selectors = RowSelection::from_filters(&[BooleanArray::from( + selector_bits.clone(), + )]); + let context = format!( + "mask_len={mask_len} selector_len={selector_len} offset={offset}" + ); + let intersection = + expected_combined(&mask_bits, &selector_bits, |a, b| a && b); + let union = + expected_combined(&mask_bits, &selector_bits, |a, b| a || b); + for (left, right) in [(&mask, &selectors), (&selectors, &mask)] { + assert_mask_eq( + left.intersection(right).as_mask().unwrap(), + &intersection, + &context, + ); + assert_mask_eq( + left.union(right).as_mask().unwrap(), + &union, + &context, + ); + } + } + } + } + } + } + } + + #[test] + fn test_mixed_algebra_empty_operand_preserves_tail() { + let empty = RowSelection::default(); + let buffer = BooleanBuffer::from_iter((0..64).map(|i| i % 3 == 0)); + let mask = RowSelection::from_boolean_buffer(buffer.slice(5, 40)); + let selectors = RowSelection::from(vec![ + RowSelector::skip(3), + RowSelector::select(7), + RowSelector::skip(4), + ]); + let empty_mask = RowSelection::from_boolean_buffer(BooleanBuffer::new_unset(0)); + for (left, right, expected) in [ + (&empty, &mask, &mask), + (&mask, &empty, &mask), + (&empty_mask, &selectors, &selectors), + (&selectors, &empty_mask, &selectors), + (&empty, &empty_mask, &empty_mask), + (&empty_mask, &empty, &empty_mask), + ] { + for result in [left.intersection(right), left.union(right)] { + assert_eq!(&result, expected); + match (result.as_mask(), expected.as_mask()) { + (Some(actual), Some(expected)) => { + assert_eq!(actual.inner().as_ptr(), expected.inner().as_ptr()); + assert_eq!(actual.offset(), expected.offset()); + } + (None, None) => {} + _ => panic!("empty operand changed the other operand's representation"), + } + } + } + } + #[test] #[cfg_attr(miri, ignore)] // Takes too long fn test_mask_algebra_with_offsets() { diff --git a/parquet/src/arrow/arrow_reader/selection/boolean.rs b/parquet/src/arrow/arrow_reader/selection/boolean.rs index 89a454b9c2..e5b698533a 100644 --- a/parquet/src/arrow/arrow_reader/selection/boolean.rs +++ b/parquet/src/arrow/arrow_reader/selection/boolean.rs @@ -26,7 +26,6 @@ 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. @@ -94,14 +93,6 @@ impl MaskSelection { .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() { @@ -450,25 +441,33 @@ mod tests { } #[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, - RowSelectionInner::Selectors(_) => 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, - RowSelectionInner::Selectors(_) => unreachable!(), - }; - match mask.borrowed_selectors() { - Cow::Borrowed(selectors) => assert_eq!(selectors, expected.as_slice()), - Cow::Owned(_) => panic!("expected the cached selectors to be reused"), + fn test_mixed_set_algebra_does_not_materialize_selector_cache() { + let other = RowSelection::from(vec![ + RowSelector::skip(3), + RowSelector::select(129), + RowSelector::skip(5), + ]); + for buffer in [ + BooleanBuffer::from_iter((0..137).map(|i| i < 64)), + BooleanBuffer::from_iter((0..137).map(|i| i % 2 == 0)), + ] { + for cached in [false, true] { + let selection = RowSelection::from_boolean_buffer(buffer.clone()); + if cached { + let _ = selection.iter().count(); + } + let cache_before = cached_selectors_ptr(&selection); + for result in [ + selection.intersection(&other), + other.intersection(&selection), + selection.union(&other), + other.union(&selection), + ] { + assert!(result.as_mask().is_some()); + assert!(cached_selectors_ptr(&result).is_none()); + } + assert_eq!(cached_selectors_ptr(&selection), cache_before); + } } } diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs b/parquet/src/arrow/arrow_reader/selection/mod.rs index 6ce7d71f9f..31c51c3730 100644 --- a/parquet/src/arrow/arrow_reader/selection/mod.rs +++ b/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -44,8 +44,8 @@ mod ranges; mod selector; use algebra::{ - and_then_mask, and_then_row_selections, and_then_selectors_with_mask, intersect_masks, - intersect_row_selections, union_masks, union_row_selections, + and_then_mask, and_then_row_selections, and_then_selectors_with_mask, combine_mixed_selection, + intersect_masks, intersect_row_selections, union_masks, union_row_selections, }; pub use boolean::MaskRunIter; use boolean::{ @@ -547,6 +547,10 @@ impl RowSelection { } /// Compute the intersection of two [`RowSelection`] + /// + /// Nonempty mixed inputs return a mask-backed result. + /// If only one input is empty, the result keeps the other's representation. + /// /// For example: /// self: NNYYYYNNYYNYN /// other: NYNNNNNNY @@ -560,16 +564,18 @@ impl RowSelection { (RowSelectionInner::Selectors(l), RowSelectionInner::Selectors(r)) => { intersect_row_selections(l, r) } - (RowSelectionInner::Selectors(l), RowSelectionInner::Mask(r)) => { - intersect_row_selections(l, &r.borrowed_selectors()) - } - (RowSelectionInner::Mask(l), RowSelectionInner::Selectors(r)) => { - intersect_row_selections(&l.borrowed_selectors(), r) + (RowSelectionInner::Selectors(selectors), RowSelectionInner::Mask(mask)) + | (RowSelectionInner::Mask(mask), RowSelectionInner::Selectors(selectors)) => { + combine_mixed_selection(selectors, mask, intersect_masks) } } } /// Compute the union of two [`RowSelection`] + /// + /// Nonempty mixed inputs return a mask-backed result. + /// If only one input is empty, the result keeps the other's representation. + /// /// For example: /// self: NNYYYYNNYYNYN /// other: NYNNNNNNN @@ -583,11 +589,9 @@ impl RowSelection { (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) + (RowSelectionInner::Selectors(selectors), RowSelectionInner::Mask(mask)) + | (RowSelectionInner::Mask(mask), RowSelectionInner::Selectors(selectors)) => { + combine_mixed_selection(selectors, mask, union_masks) } } }
