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

parquet/arrow/arrow_reader/selection/
boolean.rs
