parquet/arrow/arrow_reader/selection/
boolean.rs

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

parquet/arrow/arrow_reader/selection/
boolean.rs
