parquet/arrow/arrow_reader/selection/
boolean.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! The bitmap backed representation of a [`RowSelection`] and the primitives
19//! operating on it: conversion to and from the run length ([`RowSelector`])
20//! form, and the transforms backing `split_off`, `trim`, `offset` and `limit`.
21//!
22//! The bitwise set algebra lives in the `algebra` module.
23//!
24//! [`RowSelection`]: crate::arrow::arrow_reader::RowSelection
25
26use super::RowSelector;
27use arrow_buffer::bit_iterator::BitSliceIterator;
28use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, Buffer};
29use std::borrow::Cow;
30use std::sync::OnceLock;
31
32/// Mask-backed [`RowSelection`] storage.
33///
34/// `selectors` is only populated if callers use the borrowed
35/// [`RowSelection::iter`] compatibility API. Internal paths that can stream or
36/// consume the bitmap avoid this cache.
37///
38/// `count` caches the popcount; [`RowSelection::split_off`] propagates it to
39/// both halves so repeated `row_count()` calls do not rescan the bitmap.
40///
41/// [`RowSelection`]: crate::arrow::arrow_reader::RowSelection
42/// [`RowSelection::iter`]: crate::arrow::arrow_reader::RowSelection::iter
43/// [`RowSelection::split_off`]: crate::arrow::arrow_reader::RowSelection::split_off
44#[derive(Debug)]
45pub(crate) struct MaskSelection {
46    mask: BooleanBuffer,
47    selectors: OnceLock<Vec<RowSelector>>,
48    count: OnceLock<usize>,
49}
50
51impl MaskSelection {
52    pub(super) fn new(mask: BooleanBuffer) -> Self {
53        Self {
54            mask,
55            selectors: OnceLock::new(),
56            count: OnceLock::new(),
57        }
58    }
59
60    /// Create a selection whose selected-row count is already known.
61    pub(super) fn with_count(mask: BooleanBuffer, count: usize) -> Self {
62        debug_assert!(count <= mask.len());
63        let cell = OnceLock::new();
64        let _ = cell.set(count);
65        Self {
66            mask,
67            selectors: OnceLock::new(),
68            count: cell,
69        }
70    }
71
72    pub(crate) fn mask(&self) -> &BooleanBuffer {
73        &self.mask
74    }
75
76    pub(crate) fn into_mask(self) -> BooleanBuffer {
77        let Self { mask, .. } = self;
78        mask
79    }
80
81    /// Number of selected rows, computed once and cached.
82    pub(super) fn count(&self) -> usize {
83        *self.count.get_or_init(|| self.mask.count_set_bits())
84    }
85
86    /// The cached selected-row count, if it has been computed.
87    pub(super) fn cached_count(&self) -> Option<usize> {
88        self.count.get().copied()
89    }
90
91    pub(super) fn selectors(&self) -> &[RowSelector] {
92        self.selectors
93            .get_or_init(|| mask_to_selectors(&self.mask))
94            .as_slice()
95    }
96
97    /// Borrows the cached RLE form, converting into a temporary if not cached.
98    pub(super) fn borrowed_selectors(&self) -> Cow<'_, [RowSelector]> {
99        match self.selectors.get() {
100            Some(selectors) => Cow::Borrowed(selectors.as_slice()),
101            None => Cow::Owned(mask_to_selectors(&self.mask)),
102        }
103    }
104
105    /// The RLE form, taking the cache if it was populated.
106    pub(crate) fn into_selectors(self) -> Vec<RowSelector> {
107        match self.selectors.into_inner() {
108            Some(selectors) => selectors,
109            None => mask_to_selectors(&self.mask),
110        }
111    }
112}
113
114impl Clone for MaskSelection {
115    fn clone(&self) -> Self {
116        // Drop the selector cache but keep the cheap count cache.
117        Self {
118            mask: self.mask.clone(),
119            selectors: OnceLock::new(),
120            count: self.count.clone(),
121        }
122    }
123}
124
125/// Streaming RLE view of a [`BooleanBuffer`], yielding owned [`RowSelector`]s
126/// without allocation.
127///
128/// Useful as a zero-cost alternative to [`RowSelection::iter`] for mask-backed
129/// selections, via [`RowSelection::as_mask`]:
130///
131/// ```ignore
132/// if let Some(mask) = selection.as_mask() {
133///     for run in MaskRunIter::new(mask) { ... }
134/// }
135/// ```
136///
137/// [`RowSelection::iter`]: crate::arrow::arrow_reader::RowSelection::iter
138/// [`RowSelection::as_mask`]: crate::arrow::arrow_reader::RowSelection::as_mask
139#[derive(Debug)]
140pub struct MaskRunIter<'a> {
141    slices: BitSliceIterator<'a>,
142    cursor: usize,
143    total: usize,
144    pending: Option<RowSelector>,
145    finished: bool,
146}
147
148impl<'a> MaskRunIter<'a> {
149    /// Create a streaming RLE iterator over a [`BooleanBuffer`].
150    pub fn new(mask: &'a BooleanBuffer) -> Self {
151        Self {
152            slices: mask.set_slices(),
153            cursor: 0,
154            total: mask.len(),
155            pending: None,
156            finished: false,
157        }
158    }
159}
160
161impl Iterator for MaskRunIter<'_> {
162    type Item = RowSelector;
163
164    fn next(&mut self) -> Option<RowSelector> {
165        if let Some(p) = self.pending.take() {
166            return Some(p);
167        }
168        if self.finished {
169            return None;
170        }
171        match self.slices.next() {
172            Some((start, end)) => {
173                let select = RowSelector::select(end - start);
174                if start > self.cursor {
175                    let skip = RowSelector::skip(start - self.cursor);
176                    self.pending = Some(select);
177                    self.cursor = end;
178                    Some(skip)
179                } else {
180                    self.cursor = end;
181                    Some(select)
182                }
183            }
184            None => {
185                self.finished = true;
186                if self.cursor < self.total {
187                    let skip = RowSelector::skip(self.total - self.cursor);
188                    self.cursor = self.total;
189                    Some(skip)
190                } else {
191                    None
192                }
193            }
194        }
195    }
196}
197
198/// Materialize a [`BooleanBuffer`] into its RLE form.
199pub(super) fn mask_to_selectors(mask: &BooleanBuffer) -> Vec<RowSelector> {
200    let total_rows = mask.len();
201    if total_rows == 0 {
202        return Vec::new();
203    }
204    let mut selectors: Vec<RowSelector> = Vec::new();
205    let mut last_end = 0;
206    for (start, end) in mask.set_slices() {
207        if start > last_end {
208            selectors.push(RowSelector::skip(start - last_end));
209        }
210        selectors.push(RowSelector::select(end - start));
211        last_end = end;
212    }
213    if last_end != total_rows {
214        selectors.push(RowSelector::skip(total_rows - last_end));
215    }
216    selectors
217}
218
219/// Returns whether `mask` contains at least `min_runs` alternating set/unset runs.
220///
221/// Stops as soon as the requested number of runs is found, avoiding a full scan
222/// when callers only need to know whether a boundary has been crossed.
223pub(super) fn mask_has_at_least_runs(mask: &BooleanBuffer, min_runs: usize) -> bool {
224    if min_runs == 0 {
225        return true;
226    }
227
228    let total_rows = mask.len();
229    if total_rows == 0 {
230        return false;
231    }
232
233    let mut run_count = 0;
234    let mut last_end = 0;
235    for (start, end) in mask.set_slices() {
236        run_count += usize::from(start > last_end) + 1;
237        if run_count >= min_runs {
238            return true;
239        }
240        last_end = end;
241    }
242
243    run_count + usize::from(last_end < total_rows) >= min_runs
244}
245
246/// Split a mask into `(head, tail)` at `row_count`, preserving an empty mask tail
247/// when the split point is past the end.
248pub(super) fn split_off_mask(
249    mask: BooleanBuffer,
250    row_count: usize,
251) -> (BooleanBuffer, BooleanBuffer) {
252    let total = mask.len();
253    if row_count >= total {
254        return (mask, BooleanBuffer::new_unset(0));
255    }
256
257    let head = mask.slice(0, row_count);
258    let tail = mask.slice(row_count, total - row_count);
259    (head, tail)
260}
261
262/// Position of the highest set bit in `mask`, scanning bytes from the end.
263fn last_set_bit_position(mask: &BooleanBuffer) -> Option<usize> {
264    let values = mask.values();
265    let offset = mask.offset();
266    let end = offset + mask.len();
267    for byte_idx in (offset / 8..end.div_ceil(8)).rev() {
268        let byte_start = byte_idx * 8;
269        let mut byte = values[byte_idx];
270        if end - byte_start < 8 {
271            byte &= (1u8 << (end - byte_start)) - 1;
272        }
273        if byte_start < offset {
274            byte &= !((1u8 << (offset - byte_start)) - 1);
275        }
276        if byte != 0 {
277            return Some(byte_start + 7 - byte.leading_zeros() as usize - offset);
278        }
279    }
280    None
281}
282
283/// Trims trailing unset bits from a mask-backed selection.
284pub(super) fn trim_mask(mask: &BooleanBuffer) -> Option<BooleanBuffer> {
285    let len = mask.len();
286    // Fast path: final bit set means there is nothing to trim.
287    if len == 0 || mask.value(len - 1) {
288        return None;
289    }
290    let new_len = last_set_bit_position(mask).map_or(0, |pos| pos + 1);
291    Some(mask.slice(0, new_len))
292}
293
294/// Skips the first `offset` selected rows of a mask-backed selection.
295/// `popcount` is the caller's (possibly cached) set-bit count of `mask`.
296pub(super) fn offset_mask(mask: BooleanBuffer, offset: usize, popcount: usize) -> BooleanBuffer {
297    if offset >= popcount {
298        return BooleanBuffer::new_unset(0);
299    }
300    // Position one past the `offset`-th set bit, i.e. the index of the first
301    // selected row to keep.
302    let pos = mask.find_nth_set_bit_position(0, offset);
303    let mut builder = BooleanBufferBuilder::new(mask.len());
304    builder.append_n(pos, false);
305    builder.append_buffer(&mask.slice(pos, mask.len() - pos));
306    builder.finish()
307}
308
309/// Keeps only the first `limit` selected rows of a mask-backed selection.
310pub(super) fn limit_mask(mask: BooleanBuffer, limit: usize) -> BooleanBuffer {
311    // `find_nth_set_bit_position` returns `mask.len()` when there are fewer
312    // than `limit` set bits, so the slice naturally degrades to the original
313    // mask in that case.
314    let cut = mask.find_nth_set_bit_position(0, limit);
315    mask.slice(0, cut)
316}
317
318/// Set bits `[start, start + len)` in a zero-initialized little-endian bitmap.
319fn set_bit_run(buf: &mut [u8], start: usize, len: usize) {
320    if len == 0 {
321        return;
322    }
323    let end = start + len;
324    let first_byte = start / 8;
325    let last_byte = (end - 1) / 8;
326    let start_mask = 0xFFu8 << (start % 8);
327    let end_mask = 0xFFu8 >> (8 - (end - last_byte * 8));
328    if first_byte == last_byte {
329        buf[first_byte] |= start_mask & end_mask;
330    } else {
331        buf[first_byte] |= start_mask;
332        buf[first_byte + 1..last_byte].fill(0xFF);
333        buf[last_byte] |= end_mask;
334    }
335}
336
337/// Build a bitmap from a selector sequence by filling bytes directly.
338///
339/// This sits on the read hot path (`Mask` strategy over a selector-backed
340/// selection) where per-selector `append_n` calls are too slow.
341pub(super) fn boolean_mask_from_selectors(selectors: &[RowSelector]) -> BooleanBuffer {
342    let total_rows: usize = selectors.iter().map(|s| s.row_count).sum();
343    let mut buf = vec![0u8; total_rows.div_ceil(8)];
344    let mut position = 0usize;
345    for selector in selectors {
346        if !selector.skip {
347            set_bit_run(&mut buf, position, selector.row_count);
348        }
349        position += selector.row_count;
350    }
351    BooleanBuffer::new(Buffer::from(buf), 0, total_rows)
352}
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357    use crate::arrow::arrow_reader::selection::{RowSelection, RowSelectionInner};
358    use arrow_array::BooleanArray;
359    use rand::{RngExt, rng};
360
361    #[test]
362    fn test_mask_iter_yields_borrowed_selectors() {
363        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
364            false, false, true, true, false, true, false, false,
365        ]));
366
367        let borrowed: Vec<&RowSelector> = selection.iter().collect();
368        assert_eq!(
369            borrowed,
370            vec![
371                &RowSelector::skip(2),
372                &RowSelector::select(2),
373                &RowSelector::skip(1),
374                &RowSelector::select(1),
375                &RowSelector::skip(2),
376            ]
377        );
378    }
379
380    #[test]
381    fn test_mask_iter_clone_drops_cache() {
382        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
383            false, false, true, true, false, true, false, false,
384        ]));
385
386        let _ = selection.iter().count();
387        match &selection.inner {
388            RowSelectionInner::Mask(m) => assert!(m.selectors.get().is_some()),
389            _ => unreachable!(),
390        }
391
392        let cloned = selection.clone();
393        match &cloned.inner {
394            RowSelectionInner::Mask(m) => assert!(m.selectors.get().is_none()),
395            _ => unreachable!(),
396        }
397
398        let round_tripped: Vec<RowSelector> = cloned.iter().copied().collect();
399        assert_eq!(
400            round_tripped,
401            vec![
402                RowSelector::skip(2),
403                RowSelector::select(2),
404                RowSelector::skip(1),
405                RowSelector::select(1),
406                RowSelector::skip(2),
407            ]
408        );
409    }
410
411    /// Enough runs that the RLE form is a real allocation, so the cache reuse
412    /// tests can track its pointer across the conversion.
413    fn interleaved_mask() -> BooleanBuffer {
414        BooleanBuffer::from((0..256).map(|i| i % 3 == 0).collect::<Vec<bool>>())
415    }
416
417    fn cached_selectors_ptr(selection: &RowSelection) -> Option<*const RowSelector> {
418        match &selection.inner {
419            RowSelectionInner::Mask(m) => m.selectors.get().map(|s| s.as_ptr()),
420            _ => unreachable!(),
421        }
422    }
423
424    #[test]
425    fn test_into_selectors_takes_the_iter_cache() {
426        let selection = RowSelection::from_boolean_buffer(interleaved_mask());
427        let expected: Vec<RowSelector> = selection.iter().copied().collect();
428
429        let cached_ptr = cached_selectors_ptr(&selection).expect("iter populates the cache");
430        let selectors: Vec<RowSelector> = selection.into();
431
432        assert_eq!(selectors, expected);
433        // Moved out of the cache rather than re-encoded from the bitmap.
434        assert_eq!(selectors.as_ptr(), cached_ptr);
435    }
436
437    #[test]
438    fn test_into_selectors_without_cache_still_converts() {
439        let selection = RowSelection::from_boolean_buffer(interleaved_mask());
440        assert!(cached_selectors_ptr(&selection).is_none());
441
442        let selectors: Vec<RowSelector> = selection.into();
443        assert_eq!(selectors, mask_to_selectors(&interleaved_mask()));
444
445        // `VecDeque` goes through the same path.
446        let selection = RowSelection::from_boolean_buffer(interleaved_mask());
447        let _ = selection.iter().count();
448        let deque: std::collections::VecDeque<RowSelector> = selection.into();
449        assert_eq!(Vec::from(deque), selectors);
450    }
451
452    #[test]
453    fn test_borrowed_selectors_reuses_cache_without_populating_it() {
454        let selection = RowSelection::from_boolean_buffer(interleaved_mask());
455        let mask = match &selection.inner {
456            RowSelectionInner::Mask(m) => m,
457            _ => unreachable!(),
458        };
459
460        // Uncached: converts into a temporary, leaving the cache empty.
461        assert!(matches!(mask.borrowed_selectors(), Cow::Owned(_)));
462        assert!(mask.selectors.get().is_none());
463
464        let expected: Vec<RowSelector> = selection.iter().copied().collect();
465        let mask = match &selection.inner {
466            RowSelectionInner::Mask(m) => m,
467            _ => unreachable!(),
468        };
469        match mask.borrowed_selectors() {
470            Cow::Borrowed(selectors) => assert_eq!(selectors, expected.as_slice()),
471            Cow::Owned(_) => panic!("expected the cached selectors to be reused"),
472        }
473    }
474
475    #[test]
476    fn test_set_algebra_agrees_whether_or_not_the_cache_is_populated() {
477        let bits: Vec<bool> = (0..256).map(|i| i % 3 == 0).collect();
478        let other: RowSelection = RowSelection::from_filters(&[BooleanArray::from(
479            (0..256).map(|i| i % 5 != 0).collect::<Vec<bool>>(),
480        )]);
481
482        let cold = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
483        let warm = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits));
484        let _ = warm.iter().count();
485
486        assert_eq!(cold.intersection(&other), warm.intersection(&other));
487        assert_eq!(other.intersection(&cold), other.intersection(&warm));
488        assert_eq!(cold.union(&other), warm.union(&other));
489        assert_eq!(other.union(&cold), other.union(&warm));
490    }
491
492    #[test]
493    fn test_mask_run_iter_streams_without_cache() {
494        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
495            false, false, true, true, false, true, false, false,
496        ]));
497        let mut iter = MaskRunIter::new(selection.as_mask().unwrap());
498
499        assert_eq!(iter.next(), Some(RowSelector::skip(2)));
500        assert_eq!(iter.next(), Some(RowSelector::select(2)));
501        assert_eq!(iter.next(), Some(RowSelector::skip(1)));
502        assert_eq!(iter.next(), Some(RowSelector::select(1)));
503        assert_eq!(iter.next(), Some(RowSelector::skip(2)));
504        assert_eq!(iter.next(), None);
505        assert_eq!(iter.next(), None);
506
507        let selection =
508            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, true, false]));
509        let mut iter = MaskRunIter::new(selection.as_mask().unwrap());
510        assert_eq!(iter.next(), Some(RowSelector::select(2)));
511        assert_eq!(iter.next(), Some(RowSelector::skip(1)));
512        assert_eq!(iter.next(), None);
513    }
514
515    #[test]
516    fn test_from_boolean_buffer() {
517        let bits = vec![
518            false, false, true, true, false, true, false, false, true, false, false, false, false,
519            false, false, true,
520        ];
521        let buf = BooleanBuffer::from(bits.clone());
522        let selection = RowSelection::from_boolean_buffer(buf.clone());
523
524        assert!(selection.as_mask().is_some());
525        assert_eq!(selection.row_count(), 5);
526        assert_eq!(selection.skipped_row_count(), 11);
527        assert!(selection.selects_any());
528
529        let from_filters = RowSelection::from_filters(&[BooleanArray::from(bits)]);
530        assert_eq!(selection, from_filters);
531
532        let bits_tail = vec![true, false, true, false, false, false];
533        let trimmed = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits_tail)).trim();
534        assert!(trimmed.as_mask().is_some());
535        assert_eq!(trimmed.as_mask().unwrap().len(), 3);
536    }
537
538    #[test]
539    fn test_from_boolean_buffer_empty() {
540        let empty = RowSelection::from_boolean_buffer(BooleanBuffer::from(Vec::<bool>::new()));
541        assert!(empty.as_mask().is_some());
542        assert_eq!(empty.row_count(), 0);
543        assert_eq!(empty.skipped_row_count(), 0);
544        assert!(!empty.selects_any());
545        assert!(empty.selectors().is_empty());
546    }
547
548    #[test]
549    fn test_from_boolean_buffer_all_unset_does_not_select() {
550        let all_zero = RowSelection::from_boolean_buffer(BooleanBuffer::new_unset(1024));
551        assert!(all_zero.as_mask().is_some());
552        assert!(!all_zero.selects_any());
553        assert_eq!(all_zero.row_count(), 0);
554        assert_eq!(all_zero.skipped_row_count(), 1024);
555    }
556
557    #[test]
558    fn test_from_boolean_buffer_via_from_impl() {
559        let buf = BooleanBuffer::from(vec![true, false, true, true]);
560        let a = RowSelection::from(buf.clone());
561        let b = RowSelection::from_boolean_buffer(buf);
562        assert_eq!(a, b);
563        assert!(a.as_mask().is_some());
564    }
565
566    #[test]
567    fn test_mask_backing_clone_preserves_backing() {
568        let buf = BooleanBuffer::from(vec![true, false, true]);
569        let original = RowSelection::from_boolean_buffer(buf);
570        let cloned = original.clone();
571        assert!(cloned.as_mask().is_some());
572        assert_eq!(original, cloned);
573    }
574
575    #[test]
576    fn test_mask_backing_mutation_equivalence() {
577        let bits = vec![true, true, false, false, true, false, true, true];
578
579        let from_mask = {
580            let mut s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
581            let split = s.split_off(3);
582            (split, s)
583        };
584        let from_selectors = {
585            let mut s = RowSelection::from_filters(&[BooleanArray::from(bits.clone())]);
586            let split = s.split_off(3);
587            (split, s)
588        };
589        assert_eq!(from_mask.0, from_selectors.0);
590        assert_eq!(from_mask.1, from_selectors.1);
591        assert!(from_mask.0.as_mask().is_some());
592        assert!(from_mask.1.as_mask().is_some());
593
594        let limited_mask =
595            RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())).limit(3);
596        let limited_sel = RowSelection::from_filters(&[BooleanArray::from(bits.clone())]).limit(3);
597        assert!(limited_mask.as_mask().is_some());
598        assert_eq!(limited_mask, limited_sel);
599
600        let offset_mask =
601            RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())).offset(2);
602        let offset_sel = RowSelection::from_filters(&[BooleanArray::from(bits)]).offset(2);
603        assert!(offset_mask.as_mask().is_some());
604        assert_eq!(offset_mask, offset_sel);
605    }
606
607    #[test]
608    fn test_mask_backing_fuzz_equivalence() {
609        let mut rand = rng();
610        for _ in 0..100 {
611            let len = rand.random_range(0..200);
612            let bits: Vec<_> = (0..len).map(|_| rand.random_bool(0.35)).collect();
613
614            let from_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
615            let from_filters = RowSelection::from_filters(&[BooleanArray::from(bits.clone())]);
616
617            assert_eq!(from_mask, from_filters);
618            assert_eq!(from_mask.row_count(), from_filters.row_count());
619            assert_eq!(
620                from_mask.skipped_row_count(),
621                from_filters.skipped_row_count()
622            );
623            assert_eq!(from_mask.selects_any(), from_filters.selects_any());
624
625            let inner_len: usize = bits.iter().map(|b| *b as usize).sum();
626            let inner_bits: Vec<_> = (0..inner_len).map(|_| rand.random_bool(0.7)).collect();
627            let inner = RowSelection::from_filters(&[BooleanArray::from(inner_bits.clone())]);
628            let inner_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(inner_bits));
629            let and_then_mask = from_mask.and_then(&inner);
630            let and_then_both_masks = from_mask.and_then(&inner_mask);
631            assert!(and_then_mask.as_mask().is_some());
632            assert!(and_then_both_masks.as_mask().is_some());
633            assert_eq!(and_then_mask, from_filters.and_then(&inner));
634            assert_eq!(and_then_both_masks, and_then_mask);
635        }
636    }
637
638    #[test]
639    fn test_mask_offset_past_end_preserves_empty_mask_backing() {
640        let selection =
641            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true]))
642                .offset(2);
643
644        assert!(selection.as_mask().is_some());
645        assert_eq!(selection.as_mask().unwrap().len(), 0);
646        assert_eq!(selection.row_count(), 0);
647        assert_eq!(selection.skipped_row_count(), 0);
648    }
649
650    #[test]
651    fn test_mask_limit_truncates_at_nth_selected_row() {
652        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
653            false, true, false, true, false, true, false,
654        ]))
655        .limit(2);
656
657        let mask = selection.as_mask().unwrap();
658        assert_eq!(mask.len(), 4);
659        let actual_bits: Vec<_> = (0..mask.len()).map(|i| mask.value(i)).collect();
660        assert_eq!(actual_bits, vec![false, true, false, true]);
661    }
662
663    #[test]
664    fn test_mask_split_off_preserves_backing() {
665        let bits: Vec<bool> = (0..40).map(|i| i % 3 == 0).collect();
666        let mut s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
667        let head = s.split_off(15);
668
669        assert!(head.as_mask().is_some());
670        assert!(s.as_mask().is_some());
671
672        let head_sel = RowSelection::from_filters(&[BooleanArray::from(bits[..15].to_vec())]);
673        let tail_sel = RowSelection::from_filters(&[BooleanArray::from(bits[15..].to_vec())]);
674        assert_eq!(head, head_sel);
675        assert_eq!(s, tail_sel);
676    }
677
678    #[test]
679    fn test_mask_split_off_past_end_returns_whole() {
680        let bits = vec![true, false, true];
681        let mut s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
682        let head = s.split_off(100);
683
684        assert!(head.as_mask().is_some());
685        assert_eq!(head.as_mask().unwrap().len(), 3);
686        // `self` keeps its mask backing and is left empty.
687        assert!(s.as_mask().is_some());
688        assert_eq!(s.as_mask().unwrap().len(), 0);
689        assert_eq!(s.row_count(), 0);
690        assert_eq!(s.skipped_row_count(), 0);
691    }
692
693    #[test]
694    fn test_mask_offset_exceeds_selected_returns_empty() {
695        let s =
696            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, true, false, true]));
697        let r = s.offset(10);
698        assert_eq!(r.row_count(), 0);
699        assert_eq!(r.skipped_row_count(), 0);
700
701        let from_selectors =
702            RowSelection::from_filters(&[BooleanArray::from(vec![true, true, false, true])])
703                .offset(10);
704        assert_eq!(r, from_selectors);
705    }
706
707    #[test]
708    fn test_mask_limit_exceeds_selected_returns_all() {
709        let bits = vec![true, true, false, true];
710        let s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
711        let r = s.limit(10);
712        assert_eq!(r.row_count(), 3);
713
714        let from_selectors = RowSelection::from_filters(&[BooleanArray::from(bits)]).limit(10);
715        assert_eq!(r, from_selectors);
716    }
717
718    #[test]
719    fn test_mask_trim_all_zero_collapses_to_empty() {
720        let s = RowSelection::from_boolean_buffer(BooleanBuffer::new_unset(128));
721        let trimmed = s.trim();
722        assert!(trimmed.as_mask().is_some());
723        assert_eq!(trimmed.as_mask().unwrap().len(), 0);
724    }
725
726    #[test]
727    fn test_boolean_mask_from_selectors_fuzz_equivalence() {
728        let mut rand = rng();
729        for _ in 0..200 {
730            let n_selectors = rand.random_range(0..30);
731            let mut selectors = Vec::with_capacity(n_selectors);
732            for _ in 0..n_selectors {
733                selectors.push(RowSelector {
734                    row_count: rand.random_range(0..40),
735                    skip: rand.random_bool(0.5),
736                });
737            }
738
739            let expected = {
740                let total_rows: usize = selectors.iter().map(|s| s.row_count).sum();
741                let mut builder = BooleanBufferBuilder::new(total_rows);
742                for selector in &selectors {
743                    builder.append_n(selector.row_count, !selector.skip);
744                }
745                builder.finish()
746            };
747
748            assert_eq!(boolean_mask_from_selectors(&selectors), expected);
749        }
750    }
751
752    #[test]
753    fn test_mask_has_at_least_runs() {
754        fn assert_run_count(bits: Vec<bool>, expected_runs: usize) {
755            let mask = BooleanBuffer::from(bits);
756            for min_runs in 0..=expected_runs + 2 {
757                assert_eq!(
758                    mask_has_at_least_runs(&mask, min_runs),
759                    expected_runs >= min_runs,
760                    "expected {expected_runs} runs with boundary {min_runs}"
761                );
762            }
763        }
764
765        assert_run_count(vec![], 0);
766        assert_run_count(vec![false; 8], 1);
767        assert_run_count(vec![true; 8], 1);
768        assert_run_count(vec![false, false, true, true, false], 3);
769        assert_run_count(vec![true, false, true, false, true, false], 6);
770
771        // Exercise the unaligned iterator path as mask-backed selections can be slices.
772        let mask = BooleanBuffer::from(vec![true, false, false, true, true, false, true, true])
773            .slice(1, 6);
774        for min_runs in 0..=6 {
775            assert_eq!(mask_has_at_least_runs(&mask, min_runs), 4 >= min_runs);
776        }
777    }
778
779    #[test]
780    fn test_trim_mask_fuzz_equivalence() {
781        let mut rand = rng();
782        for _ in 0..200 {
783            let len = rand.random_range(0..200);
784            let bits: Vec<bool> = (0..len).map(|_| rand.random_bool(0.3)).collect();
785            let full = BooleanBuffer::from(bits.clone());
786            // Exercise non-zero bit offsets via slicing
787            let start = rand.random_range(0..=len);
788            let slice_len = rand.random_range(0..=(len - start));
789            let mask = full.slice(start, slice_len);
790
791            let expected_len = bits[start..start + slice_len]
792                .iter()
793                .rposition(|&b| b)
794                .map_or(0, |pos| pos + 1);
795
796            match trim_mask(&mask) {
797                Some(trimmed) => {
798                    assert_ne!(expected_len, mask.len());
799                    assert_eq!(trimmed.len(), expected_len);
800                    assert_eq!(trimmed, mask.slice(0, expected_len));
801                }
802                None => assert_eq!(expected_len, mask.len()),
803            }
804        }
805    }
806
807    #[test]
808    fn test_split_off_propagates_cached_count() {
809        let bits = vec![true, false, true, true, false, false, true, false];
810        let mut selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits));
811        // Populate the count cache, then verify both split halves.
812        assert_eq!(selection.row_count(), 4);
813        let head = selection.split_off(3);
814        assert_eq!(head.row_count(), 2);
815        assert_eq!(selection.row_count(), 2);
816        let tail_fresh = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
817            true, false, false, true, false,
818        ]));
819        assert_eq!(selection, tail_fresh);
820
821        // Splitting past the end keeps the whole selection as the head
822        let head = selection.split_off(100);
823        assert_eq!(head.row_count(), 2);
824        assert_eq!(selection.row_count(), 0);
825    }
826
827    #[test]
828    fn test_trim_and_offset_and_limit_preserve_cached_count() {
829        let bits = vec![true, true, false, true, false, false];
830        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
831        assert_eq!(selection.row_count(), 3);
832
833        let trimmed = selection.trim();
834        assert!(trimmed.as_mask().is_some());
835        assert_eq!(trimmed.as_mask().unwrap().len(), 4);
836        assert_eq!(trimmed.row_count(), 3);
837
838        let offset = trimmed.clone().offset(1);
839        assert_eq!(offset.row_count(), 2);
840
841        let limited = trimmed.limit(2);
842        assert_eq!(limited.row_count(), 2);
843    }
844}
parquet/arrow/arrow_reader/selection/boolean.rs

parquet/arrow/arrow_reader/selection/
boolean.rs
