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
97impl Clone for MaskSelection {
98    fn clone(&self) -> Self {
99        // Drop the selector cache but keep the cheap count cache.
100        Self {
101            mask: self.mask.clone(),
102            selectors: OnceLock::new(),
103            count: self.count.clone(),
104        }
105    }
106}
107
108/// Streaming RLE view of a [`BooleanBuffer`], yielding owned [`RowSelector`]s
109/// without allocation.
110///
111/// Useful as a zero-cost alternative to [`RowSelection::iter`] for mask-backed
112/// selections, via [`RowSelection::as_mask`]:
113///
114/// ```ignore
115/// if let Some(mask) = selection.as_mask() {
116///     for run in MaskRunIter::new(mask) { ... }
117/// }
118/// ```
119///
120/// [`RowSelection::iter`]: crate::arrow::arrow_reader::RowSelection::iter
121/// [`RowSelection::as_mask`]: crate::arrow::arrow_reader::RowSelection::as_mask
122#[derive(Debug)]
123pub struct MaskRunIter<'a> {
124    slices: BitSliceIterator<'a>,
125    cursor: usize,
126    total: usize,
127    pending: Option<RowSelector>,
128    finished: bool,
129}
130
131impl<'a> MaskRunIter<'a> {
132    /// Create a streaming RLE iterator over a [`BooleanBuffer`].
133    pub fn new(mask: &'a BooleanBuffer) -> Self {
134        Self {
135            slices: mask.set_slices(),
136            cursor: 0,
137            total: mask.len(),
138            pending: None,
139            finished: false,
140        }
141    }
142}
143
144impl Iterator for MaskRunIter<'_> {
145    type Item = RowSelector;
146
147    fn next(&mut self) -> Option<RowSelector> {
148        if let Some(p) = self.pending.take() {
149            return Some(p);
150        }
151        if self.finished {
152            return None;
153        }
154        match self.slices.next() {
155            Some((start, end)) => {
156                let select = RowSelector::select(end - start);
157                if start > self.cursor {
158                    let skip = RowSelector::skip(start - self.cursor);
159                    self.pending = Some(select);
160                    self.cursor = end;
161                    Some(skip)
162                } else {
163                    self.cursor = end;
164                    Some(select)
165                }
166            }
167            None => {
168                self.finished = true;
169                if self.cursor < self.total {
170                    let skip = RowSelector::skip(self.total - self.cursor);
171                    self.cursor = self.total;
172                    Some(skip)
173                } else {
174                    None
175                }
176            }
177        }
178    }
179}
180
181/// Materialize a [`BooleanBuffer`] into its RLE form.
182pub(crate) fn mask_to_selectors(mask: &BooleanBuffer) -> Vec<RowSelector> {
183    let total_rows = mask.len();
184    if total_rows == 0 {
185        return Vec::new();
186    }
187    let mut selectors: Vec<RowSelector> = Vec::new();
188    let mut last_end = 0;
189    for (start, end) in mask.set_slices() {
190        if start > last_end {
191            selectors.push(RowSelector::skip(start - last_end));
192        }
193        selectors.push(RowSelector::select(end - start));
194        last_end = end;
195    }
196    if last_end != total_rows {
197        selectors.push(RowSelector::skip(total_rows - last_end));
198    }
199    selectors
200}
201
202/// Returns whether `mask` contains at least `min_runs` alternating set/unset runs.
203///
204/// Stops as soon as the requested number of runs is found, avoiding a full scan
205/// when callers only need to know whether a boundary has been crossed.
206pub(super) fn mask_has_at_least_runs(mask: &BooleanBuffer, min_runs: usize) -> bool {
207    if min_runs == 0 {
208        return true;
209    }
210
211    let total_rows = mask.len();
212    if total_rows == 0 {
213        return false;
214    }
215
216    let mut run_count = 0;
217    let mut last_end = 0;
218    for (start, end) in mask.set_slices() {
219        run_count += usize::from(start > last_end) + 1;
220        if run_count >= min_runs {
221            return true;
222        }
223        last_end = end;
224    }
225
226    run_count + usize::from(last_end < total_rows) >= min_runs
227}
228
229/// Split a mask into `(head, tail)` at `row_count`, preserving an empty mask tail
230/// when the split point is past the end.
231pub(super) fn split_off_mask(
232    mask: BooleanBuffer,
233    row_count: usize,
234) -> (BooleanBuffer, BooleanBuffer) {
235    let total = mask.len();
236    if row_count >= total {
237        return (mask, BooleanBuffer::new_unset(0));
238    }
239
240    let head = mask.slice(0, row_count);
241    let tail = mask.slice(row_count, total - row_count);
242    (head, tail)
243}
244
245/// Position of the highest set bit in `mask`, scanning bytes from the end.
246fn last_set_bit_position(mask: &BooleanBuffer) -> Option<usize> {
247    let values = mask.values();
248    let offset = mask.offset();
249    let end = offset + mask.len();
250    for byte_idx in (offset / 8..end.div_ceil(8)).rev() {
251        let byte_start = byte_idx * 8;
252        let mut byte = values[byte_idx];
253        if end - byte_start < 8 {
254            byte &= (1u8 << (end - byte_start)) - 1;
255        }
256        if byte_start < offset {
257            byte &= !((1u8 << (offset - byte_start)) - 1);
258        }
259        if byte != 0 {
260            return Some(byte_start + 7 - byte.leading_zeros() as usize - offset);
261        }
262    }
263    None
264}
265
266/// Trims trailing unset bits from a mask-backed selection.
267pub(super) fn trim_mask(mask: &BooleanBuffer) -> Option<BooleanBuffer> {
268    let len = mask.len();
269    // Fast path: final bit set means there is nothing to trim.
270    if len == 0 || mask.value(len - 1) {
271        return None;
272    }
273    let new_len = last_set_bit_position(mask).map_or(0, |pos| pos + 1);
274    Some(mask.slice(0, new_len))
275}
276
277/// Skips the first `offset` selected rows of a mask-backed selection.
278/// `popcount` is the caller's (possibly cached) set-bit count of `mask`.
279pub(super) fn offset_mask(mask: BooleanBuffer, offset: usize, popcount: usize) -> BooleanBuffer {
280    if offset >= popcount {
281        return BooleanBuffer::new_unset(0);
282    }
283    // Position one past the `offset`-th set bit, i.e. the index of the first
284    // selected row to keep.
285    let pos = mask.find_nth_set_bit_position(0, offset);
286    let mut builder = BooleanBufferBuilder::new(mask.len());
287    builder.append_n(pos, false);
288    builder.append_buffer(&mask.slice(pos, mask.len() - pos));
289    builder.finish()
290}
291
292/// Keeps only the first `limit` selected rows of a mask-backed selection.
293pub(super) fn limit_mask(mask: BooleanBuffer, limit: usize) -> BooleanBuffer {
294    // `find_nth_set_bit_position` returns `mask.len()` when there are fewer
295    // than `limit` set bits, so the slice naturally degrades to the original
296    // mask in that case.
297    let cut = mask.find_nth_set_bit_position(0, limit);
298    mask.slice(0, cut)
299}
300
301/// Set bits `[start, start + len)` in a zero-initialized little-endian bitmap.
302fn set_bit_run(buf: &mut [u8], start: usize, len: usize) {
303    if len == 0 {
304        return;
305    }
306    let end = start + len;
307    let first_byte = start / 8;
308    let last_byte = (end - 1) / 8;
309    let start_mask = 0xFFu8 << (start % 8);
310    let end_mask = 0xFFu8 >> (8 - (end - last_byte * 8));
311    if first_byte == last_byte {
312        buf[first_byte] |= start_mask & end_mask;
313    } else {
314        buf[first_byte] |= start_mask;
315        buf[first_byte + 1..last_byte].fill(0xFF);
316        buf[last_byte] |= end_mask;
317    }
318}
319
320/// Build a bitmap from a selector sequence by filling bytes directly.
321///
322/// This sits on the read hot path (`Mask` strategy over a selector-backed
323/// selection) where per-selector `append_n` calls are too slow.
324pub(super) fn boolean_mask_from_selectors(selectors: &[RowSelector]) -> BooleanBuffer {
325    let total_rows: usize = selectors.iter().map(|s| s.row_count).sum();
326    let mut buf = vec![0u8; total_rows.div_ceil(8)];
327    let mut position = 0usize;
328    for selector in selectors {
329        if !selector.skip {
330            set_bit_run(&mut buf, position, selector.row_count);
331        }
332        position += selector.row_count;
333    }
334    BooleanBuffer::new(Buffer::from(buf), 0, total_rows)
335}
336
337#[cfg(test)]
338mod tests {
339    use super::*;
340    use crate::arrow::arrow_reader::selection::{RowSelection, RowSelectionInner};
341    use arrow_array::BooleanArray;
342    use rand::{Rng, rng};
343
344    #[test]
345    fn test_mask_iter_yields_borrowed_selectors() {
346        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
347            false, false, true, true, false, true, false, false,
348        ]));
349
350        let borrowed: Vec<&RowSelector> = selection.iter().collect();
351        assert_eq!(
352            borrowed,
353            vec![
354                &RowSelector::skip(2),
355                &RowSelector::select(2),
356                &RowSelector::skip(1),
357                &RowSelector::select(1),
358                &RowSelector::skip(2),
359            ]
360        );
361    }
362
363    #[test]
364    fn test_mask_iter_clone_drops_cache() {
365        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
366            false, false, true, true, false, true, false, false,
367        ]));
368
369        let _ = selection.iter().count();
370        match &selection.inner {
371            RowSelectionInner::Mask(m) => assert!(m.selectors.get().is_some()),
372            _ => unreachable!(),
373        }
374
375        let cloned = selection.clone();
376        match &cloned.inner {
377            RowSelectionInner::Mask(m) => assert!(m.selectors.get().is_none()),
378            _ => unreachable!(),
379        }
380
381        let round_tripped: Vec<RowSelector> = cloned.iter().copied().collect();
382        assert_eq!(
383            round_tripped,
384            vec![
385                RowSelector::skip(2),
386                RowSelector::select(2),
387                RowSelector::skip(1),
388                RowSelector::select(1),
389                RowSelector::skip(2),
390            ]
391        );
392    }
393
394    #[test]
395    fn test_mask_run_iter_streams_without_cache() {
396        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
397            false, false, true, true, false, true, false, false,
398        ]));
399        let mut iter = MaskRunIter::new(selection.as_mask().unwrap());
400
401        assert_eq!(iter.next(), Some(RowSelector::skip(2)));
402        assert_eq!(iter.next(), Some(RowSelector::select(2)));
403        assert_eq!(iter.next(), Some(RowSelector::skip(1)));
404        assert_eq!(iter.next(), Some(RowSelector::select(1)));
405        assert_eq!(iter.next(), Some(RowSelector::skip(2)));
406        assert_eq!(iter.next(), None);
407        assert_eq!(iter.next(), None);
408
409        let selection =
410            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, true, false]));
411        let mut iter = MaskRunIter::new(selection.as_mask().unwrap());
412        assert_eq!(iter.next(), Some(RowSelector::select(2)));
413        assert_eq!(iter.next(), Some(RowSelector::skip(1)));
414        assert_eq!(iter.next(), None);
415    }
416
417    #[test]
418    fn test_from_boolean_buffer() {
419        let bits = vec![
420            false, false, true, true, false, true, false, false, true, false, false, false, false,
421            false, false, true,
422        ];
423        let buf = BooleanBuffer::from(bits.clone());
424        let selection = RowSelection::from_boolean_buffer(buf.clone());
425
426        assert!(selection.as_mask().is_some());
427        assert_eq!(selection.row_count(), 5);
428        assert_eq!(selection.skipped_row_count(), 11);
429        assert!(selection.selects_any());
430
431        let from_filters = RowSelection::from_filters(&[BooleanArray::from(bits)]);
432        assert_eq!(selection, from_filters);
433
434        let bits_tail = vec![true, false, true, false, false, false];
435        let trimmed = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits_tail)).trim();
436        assert!(trimmed.as_mask().is_some());
437        assert_eq!(trimmed.as_mask().unwrap().len(), 3);
438    }
439
440    #[test]
441    fn test_from_boolean_buffer_empty() {
442        let empty = RowSelection::from_boolean_buffer(BooleanBuffer::from(Vec::<bool>::new()));
443        assert!(empty.as_mask().is_some());
444        assert_eq!(empty.row_count(), 0);
445        assert_eq!(empty.skipped_row_count(), 0);
446        assert!(!empty.selects_any());
447        assert!(empty.selectors().is_empty());
448    }
449
450    #[test]
451    fn test_from_boolean_buffer_all_unset_does_not_select() {
452        let all_zero = RowSelection::from_boolean_buffer(BooleanBuffer::new_unset(1024));
453        assert!(all_zero.as_mask().is_some());
454        assert!(!all_zero.selects_any());
455        assert_eq!(all_zero.row_count(), 0);
456        assert_eq!(all_zero.skipped_row_count(), 1024);
457    }
458
459    #[test]
460    fn test_from_boolean_buffer_via_from_impl() {
461        let buf = BooleanBuffer::from(vec![true, false, true, true]);
462        let a = RowSelection::from(buf.clone());
463        let b = RowSelection::from_boolean_buffer(buf);
464        assert_eq!(a, b);
465        assert!(a.as_mask().is_some());
466    }
467
468    #[test]
469    fn test_mask_backing_clone_preserves_backing() {
470        let buf = BooleanBuffer::from(vec![true, false, true]);
471        let original = RowSelection::from_boolean_buffer(buf);
472        let cloned = original.clone();
473        assert!(cloned.as_mask().is_some());
474        assert_eq!(original, cloned);
475    }
476
477    #[test]
478    fn test_mask_backing_mutation_equivalence() {
479        let bits = vec![true, true, false, false, true, false, true, true];
480
481        let from_mask = {
482            let mut s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
483            let split = s.split_off(3);
484            (split, s)
485        };
486        let from_selectors = {
487            let mut s = RowSelection::from_filters(&[BooleanArray::from(bits.clone())]);
488            let split = s.split_off(3);
489            (split, s)
490        };
491        assert_eq!(from_mask.0, from_selectors.0);
492        assert_eq!(from_mask.1, from_selectors.1);
493        assert!(from_mask.0.as_mask().is_some());
494        assert!(from_mask.1.as_mask().is_some());
495
496        let limited_mask =
497            RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())).limit(3);
498        let limited_sel = RowSelection::from_filters(&[BooleanArray::from(bits.clone())]).limit(3);
499        assert!(limited_mask.as_mask().is_some());
500        assert_eq!(limited_mask, limited_sel);
501
502        let offset_mask =
503            RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())).offset(2);
504        let offset_sel = RowSelection::from_filters(&[BooleanArray::from(bits)]).offset(2);
505        assert!(offset_mask.as_mask().is_some());
506        assert_eq!(offset_mask, offset_sel);
507    }
508
509    #[test]
510    fn test_mask_backing_fuzz_equivalence() {
511        let mut rand = rng();
512        for _ in 0..100 {
513            let len = rand.random_range(0..200);
514            let bits: Vec<_> = (0..len).map(|_| rand.random_bool(0.35)).collect();
515
516            let from_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
517            let from_filters = RowSelection::from_filters(&[BooleanArray::from(bits.clone())]);
518
519            assert_eq!(from_mask, from_filters);
520            assert_eq!(from_mask.row_count(), from_filters.row_count());
521            assert_eq!(
522                from_mask.skipped_row_count(),
523                from_filters.skipped_row_count()
524            );
525            assert_eq!(from_mask.selects_any(), from_filters.selects_any());
526
527            let inner_len: usize = bits.iter().map(|b| *b as usize).sum();
528            let inner_bits: Vec<_> = (0..inner_len).map(|_| rand.random_bool(0.7)).collect();
529            let inner = RowSelection::from_filters(&[BooleanArray::from(inner_bits.clone())]);
530            let inner_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(inner_bits));
531            let and_then_mask = from_mask.and_then(&inner);
532            let and_then_both_masks = from_mask.and_then(&inner_mask);
533            assert!(and_then_mask.as_mask().is_some());
534            assert!(and_then_both_masks.as_mask().is_some());
535            assert_eq!(and_then_mask, from_filters.and_then(&inner));
536            assert_eq!(and_then_both_masks, and_then_mask);
537        }
538    }
539
540    #[test]
541    fn test_mask_offset_past_end_preserves_empty_mask_backing() {
542        let selection =
543            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true]))
544                .offset(2);
545
546        assert!(selection.as_mask().is_some());
547        assert_eq!(selection.as_mask().unwrap().len(), 0);
548        assert_eq!(selection.row_count(), 0);
549        assert_eq!(selection.skipped_row_count(), 0);
550    }
551
552    #[test]
553    fn test_mask_limit_truncates_at_nth_selected_row() {
554        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
555            false, true, false, true, false, true, false,
556        ]))
557        .limit(2);
558
559        let mask = selection.as_mask().unwrap();
560        assert_eq!(mask.len(), 4);
561        let actual_bits: Vec<_> = (0..mask.len()).map(|i| mask.value(i)).collect();
562        assert_eq!(actual_bits, vec![false, true, false, true]);
563    }
564
565    #[test]
566    fn test_mask_split_off_preserves_backing() {
567        let bits: Vec<bool> = (0..40).map(|i| i % 3 == 0).collect();
568        let mut s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
569        let head = s.split_off(15);
570
571        assert!(head.as_mask().is_some());
572        assert!(s.as_mask().is_some());
573
574        let head_sel = RowSelection::from_filters(&[BooleanArray::from(bits[..15].to_vec())]);
575        let tail_sel = RowSelection::from_filters(&[BooleanArray::from(bits[15..].to_vec())]);
576        assert_eq!(head, head_sel);
577        assert_eq!(s, tail_sel);
578    }
579
580    #[test]
581    fn test_mask_split_off_past_end_returns_whole() {
582        let bits = vec![true, false, true];
583        let mut s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
584        let head = s.split_off(100);
585
586        assert!(head.as_mask().is_some());
587        assert_eq!(head.as_mask().unwrap().len(), 3);
588        // `self` keeps its mask backing and is left empty.
589        assert!(s.as_mask().is_some());
590        assert_eq!(s.as_mask().unwrap().len(), 0);
591        assert_eq!(s.row_count(), 0);
592        assert_eq!(s.skipped_row_count(), 0);
593    }
594
595    #[test]
596    fn test_mask_offset_exceeds_selected_returns_empty() {
597        let s =
598            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, true, false, true]));
599        let r = s.offset(10);
600        assert_eq!(r.row_count(), 0);
601        assert_eq!(r.skipped_row_count(), 0);
602
603        let from_selectors =
604            RowSelection::from_filters(&[BooleanArray::from(vec![true, true, false, true])])
605                .offset(10);
606        assert_eq!(r, from_selectors);
607    }
608
609    #[test]
610    fn test_mask_limit_exceeds_selected_returns_all() {
611        let bits = vec![true, true, false, true];
612        let s = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
613        let r = s.limit(10);
614        assert_eq!(r.row_count(), 3);
615
616        let from_selectors = RowSelection::from_filters(&[BooleanArray::from(bits)]).limit(10);
617        assert_eq!(r, from_selectors);
618    }
619
620    #[test]
621    fn test_mask_trim_all_zero_collapses_to_empty() {
622        let s = RowSelection::from_boolean_buffer(BooleanBuffer::new_unset(128));
623        let trimmed = s.trim();
624        assert!(trimmed.as_mask().is_some());
625        assert_eq!(trimmed.as_mask().unwrap().len(), 0);
626    }
627
628    #[test]
629    fn test_boolean_mask_from_selectors_fuzz_equivalence() {
630        let mut rand = rng();
631        for _ in 0..200 {
632            let n_selectors = rand.random_range(0..30);
633            let mut selectors = Vec::with_capacity(n_selectors);
634            for _ in 0..n_selectors {
635                selectors.push(RowSelector {
636                    row_count: rand.random_range(0..40),
637                    skip: rand.random_bool(0.5),
638                });
639            }
640
641            let expected = {
642                let total_rows: usize = selectors.iter().map(|s| s.row_count).sum();
643                let mut builder = BooleanBufferBuilder::new(total_rows);
644                for selector in &selectors {
645                    builder.append_n(selector.row_count, !selector.skip);
646                }
647                builder.finish()
648            };
649
650            assert_eq!(boolean_mask_from_selectors(&selectors), expected);
651        }
652    }
653
654    #[test]
655    fn test_mask_has_at_least_runs() {
656        fn assert_run_count(bits: Vec<bool>, expected_runs: usize) {
657            let mask = BooleanBuffer::from(bits);
658            for min_runs in 0..=expected_runs + 2 {
659                assert_eq!(
660                    mask_has_at_least_runs(&mask, min_runs),
661                    expected_runs >= min_runs,
662                    "expected {expected_runs} runs with boundary {min_runs}"
663                );
664            }
665        }
666
667        assert_run_count(vec![], 0);
668        assert_run_count(vec![false; 8], 1);
669        assert_run_count(vec![true; 8], 1);
670        assert_run_count(vec![false, false, true, true, false], 3);
671        assert_run_count(vec![true, false, true, false, true, false], 6);
672
673        // Exercise the unaligned iterator path as mask-backed selections can be slices.
674        let mask = BooleanBuffer::from(vec![true, false, false, true, true, false, true, true])
675            .slice(1, 6);
676        for min_runs in 0..=6 {
677            assert_eq!(mask_has_at_least_runs(&mask, min_runs), 4 >= min_runs);
678        }
679    }
680
681    #[test]
682    fn test_trim_mask_fuzz_equivalence() {
683        let mut rand = rng();
684        for _ in 0..200 {
685            let len = rand.random_range(0..200);
686            let bits: Vec<bool> = (0..len).map(|_| rand.random_bool(0.3)).collect();
687            let full = BooleanBuffer::from(bits.clone());
688            // Exercise non-zero bit offsets via slicing
689            let start = rand.random_range(0..=len);
690            let slice_len = rand.random_range(0..=(len - start));
691            let mask = full.slice(start, slice_len);
692
693            let expected_len = bits[start..start + slice_len]
694                .iter()
695                .rposition(|&b| b)
696                .map_or(0, |pos| pos + 1);
697
698            match trim_mask(&mask) {
699                Some(trimmed) => {
700                    assert_ne!(expected_len, mask.len());
701                    assert_eq!(trimmed.len(), expected_len);
702                    assert_eq!(trimmed, mask.slice(0, expected_len));
703                }
704                None => assert_eq!(expected_len, mask.len()),
705            }
706        }
707    }
708
709    #[test]
710    fn test_split_off_propagates_cached_count() {
711        let bits = vec![true, false, true, true, false, false, true, false];
712        let mut selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits));
713        // Populate the count cache, then verify both split halves.
714        assert_eq!(selection.row_count(), 4);
715        let head = selection.split_off(3);
716        assert_eq!(head.row_count(), 2);
717        assert_eq!(selection.row_count(), 2);
718        let tail_fresh = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
719            true, false, false, true, false,
720        ]));
721        assert_eq!(selection, tail_fresh);
722
723        // Splitting past the end keeps the whole selection as the head
724        let head = selection.split_off(100);
725        assert_eq!(head.row_count(), 2);
726        assert_eq!(selection.row_count(), 0);
727    }
728
729    #[test]
730    fn test_trim_and_offset_and_limit_preserve_cached_count() {
731        let bits = vec![true, true, false, true, false, false];
732        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
733        assert_eq!(selection.row_count(), 3);
734
735        let trimmed = selection.trim();
736        assert!(trimmed.as_mask().is_some());
737        assert_eq!(trimmed.as_mask().unwrap().len(), 4);
738        assert_eq!(trimmed.row_count(), 3);
739
740        let offset = trimmed.clone().offset(1);
741        assert_eq!(offset.row_count(), 2);
742
743        let limited = trimmed.limit(2);
744        assert_eq!(limited.row_count(), 2);
745    }
746}
parquet/arrow/arrow_reader/selection/boolean.rs

parquet/arrow/arrow_reader/selection/
boolean.rs
