parquet/arrow/arrow_reader/selection/
mod.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//! Logic for selecting which rows to read: [`RowSelection`] and [`RowSelector`]
6//!
7//! This module holds [`RowSelection`] and its public API, which dispatches to
8//! one of the two backings depending on how the selection is stored:
9//!
10//! * `selector`: the run length backing, [`RowSelector`] and its primitives
11//! * `boolean`: the bitmap backing, `MaskSelection` and its primitives
12//!
13//! The remaining modules hold the operations that are common to both:
14//!
15//! * `algebra`: `and_then`, `intersection` and `union`
16//! * `ranges`: mapping a [`RowSelection`] onto page and batch ranges
17//! * `cursor`: iterating a [`RowSelection`] while reading
18
19use crate::file::page_index::offset_index::PageLocation;
20use arrow_array::{Array, BooleanArray};
21use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
22use arrow_select::filter::SlicesIterator;
23use std::cmp::Ordering;
24use std::collections::VecDeque;
25use std::ops::Range;
26
27mod algebra;
28mod boolean;
29mod cursor;
30mod ranges;
31mod selector;
32
33use algebra::{
34    and_then_mask, and_then_row_selections, and_then_selectors_with_mask, combine_mixed_selection,
35    intersect_masks, intersect_row_selections, union_masks, union_row_selections,
36};
37pub use boolean::MaskRunIter;
38use boolean::{
39    MaskSelection, boolean_mask_from_selectors, limit_mask, mask_has_at_least_runs, offset_mask,
40    split_off_mask, trim_mask,
41};
42pub(crate) use cursor::{LoadedRowRanges, MaskCursor, RowSelectionStrategy};
43pub use cursor::{RowSelectionCursor, RowSelectionPolicy};
44use ranges::{
45    expand_to_batch_boundaries_from_selectors, scan_ranges_from_mask, scan_ranges_from_selectors,
46};
47pub use selector::RowSelector;
48use selector::{limit_selectors, offset_selectors, split_off_selectors};
49
50/// [`RowSelection`] represents selecting a subset of rows
51/// when scanning a parquet file.
52///
53/// This is applied prior to reading column data, and can therefore
54/// be used to skip IO to fetch data into memory
55///
56/// A typical use-case would be using the [`PageIndex`] to filter out rows
57/// that don't satisfy a predicate
58///
59/// Depending on the pattern of rows to be selected, [`RowSelection`] has
60/// either a bitmap or an RLE ([`RowSelector`]) based implementation.
61///
62/// # Example
63/// ```
64/// use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
65///
66/// let selectors = vec![
67///     RowSelector::skip(5),
68///     RowSelector::select(5),
69///     RowSelector::select(5),
70///     RowSelector::skip(5),
71/// ];
72///
73/// // Creating a selection will combine adjacent selectors
74/// let selection: RowSelection = selectors.into();
75///
76/// let expected = vec![
77///     RowSelector::skip(5),
78///     RowSelector::select(10),
79///     RowSelector::skip(5),
80/// ];
81///
82/// let actual: Vec<RowSelector> = selection.into();
83/// assert_eq!(actual, expected);
84///
85/// // you can also create a selection from consecutive ranges
86/// let ranges = vec![5..10, 10..15];
87/// let selection =
88///   RowSelection::from_consecutive_ranges(ranges.into_iter(), 20);
89/// let actual: Vec<RowSelector> = selection.into();
90/// assert_eq!(actual, expected);
91///
92/// // or directly from a packed bitmap, when the upstream producer already
93/// // has one. The bitmap is kept as-is rather than run-length-encoded.
94/// use arrow_buffer::BooleanBuffer;
95/// let mask = BooleanBuffer::from(vec![true, false, true, true]);
96/// let selection = RowSelection::from_boolean_buffer(mask);
97/// assert_eq!(selection.row_count(), 3);
98/// ```
99///
100/// An RLE ([`RowSelector`]) backed [`RowSelection`] maintains the following
101/// invariants (they do not apply to the bitmap backed implementation):
102///
103/// * It contains no [`RowSelector`] of 0 rows
104/// * Consecutive [`RowSelector`]s alternate skipping or selecting rows
105///
106/// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData
107#[derive(Default, Clone)]
108pub struct RowSelection {
109    inner: RowSelectionInner,
110}
111
112/// Internal storage for [`RowSelection`].
113#[derive(Debug, Clone)]
114pub(crate) enum RowSelectionInner {
115    Selectors(Vec<RowSelector>),
116    Mask(Box<MaskSelection>),
117}
118
119impl Default for RowSelectionInner {
120    fn default() -> Self {
121        Self::Selectors(Vec::new())
122    }
123}
124
125impl std::fmt::Debug for RowSelection {
126    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
127        match &self.inner {
128            RowSelectionInner::Selectors(s) => f
129                .debug_struct("RowSelection")
130                .field("selectors", s)
131                .finish(),
132            RowSelectionInner::Mask(m) => f
133                .debug_struct("RowSelection")
134                .field("mask_len", &m.mask().len())
135                .finish_non_exhaustive(),
136        }
137    }
138}
139
140impl PartialEq for RowSelection {
141    fn eq(&self, other: &Self) -> bool {
142        match (&self.inner, &other.inner) {
143            (RowSelectionInner::Selectors(a), RowSelectionInner::Selectors(b)) => a == b,
144            (RowSelectionInner::Mask(a), RowSelectionInner::Mask(b)) => a.mask() == b.mask(),
145            (RowSelectionInner::Mask(mask), RowSelectionInner::Selectors(selectors))
146            | (RowSelectionInner::Selectors(selectors), RowSelectionInner::Mask(mask)) => {
147                if selectors
148                    .iter()
149                    .try_fold(0usize, |acc, selector| acc.checked_add(selector.row_count))
150                    != Some(mask.mask().len())
151                {
152                    return false;
153                }
154
155                let mut slices = mask.mask().set_slices().peekable();
156                let mut cursor = 0usize;
157
158                for selector in selectors {
159                    let end = cursor + selector.row_count;
160
161                    if selector.skip {
162                        if slices.peek().is_some_and(|(start, _)| *start < end) {
163                            return false;
164                        }
165                    } else {
166                        match slices.next() {
167                            Some((start, slice_end)) if start == cursor && slice_end == end => {}
168                            _ => return false,
169                        }
170                    }
171
172                    cursor = end;
173                }
174
175                slices.next().is_none()
176            }
177        }
178    }
179}
180
181impl Eq for RowSelection {}
182
183impl RowSelection {
184    /// Not `pub`: unlike `From<Vec<RowSelector>>`, this performs no
185    /// validation/normalization of the selectors (e.g. combining adjacent
186    /// selectors), so callers must uphold the invariants themselves.
187    fn from_selectors(selectors: Vec<RowSelector>) -> Self {
188        Self {
189            inner: RowSelectionInner::Selectors(selectors),
190        }
191    }
192
193    /// Create a [`RowSelection`] from a packed [`BooleanBuffer`].
194    ///
195    /// Each set bit selects a row, each unset bit skips one. Unlike
196    /// [`Self::from_filters`], the bitmap is kept as-is rather than
197    /// eagerly run-length-encoded. [`Self::iter`] materializes and caches the
198    /// RLE form on first use; use [`MaskRunIter`] to stream the RLE form
199    /// directly from the bitmap.
200    pub fn from_boolean_buffer(mask: BooleanBuffer) -> Self {
201        Self {
202            inner: RowSelectionInner::Mask(Box::new(MaskSelection::new(mask))),
203        }
204    }
205
206    fn from_mask_selection(mask: MaskSelection) -> Self {
207        Self {
208            inner: RowSelectionInner::Mask(Box::new(mask)),
209        }
210    }
211
212    /// Returns the underlying mask if this selection is mask-backed.
213    ///
214    /// Public so that engines composing selections (e.g. DataFusion's
215    /// [`ParquetAccessPlan::into_overall_row_selection`]) can concatenate
216    /// mask-backed selections without materialising the RLE form.
217    ///
218    /// [`ParquetAccessPlan::into_overall_row_selection`]: https://docs.rs/datafusion-datasource-parquet/latest/datafusion_datasource_parquet/access_plan/struct.ParquetAccessPlan.html#method.into_overall_row_selection
219    pub fn as_mask(&self) -> Option<&BooleanBuffer> {
220        match &self.inner {
221            RowSelectionInner::Mask(m) => Some(m.mask()),
222            RowSelectionInner::Selectors(_) => None,
223        }
224    }
225
226    /// Consume this selection and return a mask-backed selection.
227    ///
228    /// If this selection is already mask-backed, it is returned unchanged.
229    /// Otherwise, its selectors are converted to a bitmap, preserving all
230    /// selected and skipped rows, including trailing skips.
231    ///
232    /// This can be used to keep the result of [`Self::intersection`] or
233    /// [`Self::union`] mask-backed when the other selection is mask-backed.
234    ///
235    /// # Example
236    /// ```
237    /// use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
238    ///
239    /// let selection = RowSelection::from(vec![
240    ///     RowSelector::skip(2),
241    ///     RowSelector::select(3),
242    ///     RowSelector::skip(1),
243    /// ]).force_mask();
244    ///
245    /// assert!(selection.as_mask().is_some());
246    /// assert_eq!(selection.total_row_count(), 6);
247    /// assert_eq!(selection.row_count(), 3);
248    /// ```
249    pub fn force_mask(self) -> Self {
250        if self.as_mask().is_some() {
251            self
252        } else {
253            Self::from_boolean_buffer(self.into_boolean_buffer())
254        }
255    }
256
257    /// Consume the selection and return its internal storage.
258    pub(crate) fn into_inner(self) -> RowSelectionInner {
259        self.inner
260    }
261
262    /// Consume this selection and return its bitmap representation.
263    pub(crate) fn into_boolean_buffer(self) -> BooleanBuffer {
264        match self.inner {
265            RowSelectionInner::Mask(mask) => mask.into_mask(),
266            RowSelectionInner::Selectors(selectors) => boolean_mask_from_selectors(&selectors),
267        }
268    }
269
270    /// Choose the automatic materialisation strategy without converting between
271    /// selector and mask backing.
272    #[inline]
273    pub(crate) fn auto_selection_strategy(&self, threshold: usize) -> RowSelectionStrategy {
274        match &self.inner {
275            RowSelectionInner::Selectors(selectors) => {
276                let (total_rows, run_count) =
277                    selectors
278                        .iter()
279                        .fold((0usize, 0usize), |(rows, count), selector| {
280                            if selector.row_count > 0 {
281                                (rows + selector.row_count, count + 1)
282                            } else {
283                                (rows, count)
284                            }
285                        });
286
287                if run_count == 0
288                    || auto_min_mask_runs(total_rows, threshold)
289                        .is_some_and(|min_runs| run_count >= min_runs)
290                {
291                    RowSelectionStrategy::Mask
292                } else {
293                    RowSelectionStrategy::Selectors
294                }
295            }
296            RowSelectionInner::Mask(mask) => {
297                let mask = mask.mask();
298                let total_rows = mask.len();
299
300                if total_rows == 0 {
301                    return RowSelectionStrategy::Mask;
302                }
303
304                match auto_min_mask_runs(total_rows, threshold) {
305                    Some(min_runs) if mask_has_at_least_runs(mask, min_runs) => {
306                        RowSelectionStrategy::Mask
307                    }
308                    _ => RowSelectionStrategy::Selectors,
309                }
310            }
311        }
312    }
313
314    #[cfg(test)]
315    fn selectors(&self) -> Vec<RowSelector> {
316        self.iter().copied().collect()
317    }
318
319    fn into_selectors_vec(self) -> Vec<RowSelector> {
320        match self.inner {
321            RowSelectionInner::Selectors(s) => s,
322            RowSelectionInner::Mask(m) => (*m).into_selectors(),
323        }
324    }
325
326    /// Creates a [`RowSelection`] from a slice of [`BooleanArray`]
327    ///
328    /// # Panics
329    ///
330    /// Panics if any of the [`BooleanArray`] contain nulls
331    pub fn from_filters(filters: &[BooleanArray]) -> Self {
332        let mut next_offset = 0;
333        let total_rows = filters.iter().map(|x| x.len()).sum();
334
335        let iter = filters.iter().flat_map(|filter| {
336            let offset = next_offset;
337            next_offset += filter.len();
338            assert_eq!(filter.null_count(), 0);
339            SlicesIterator::new(filter).map(move |(start, end)| start + offset..end + offset)
340        });
341
342        Self::from_consecutive_ranges(iter, total_rows)
343    }
344
345    /// Builds a selection equivalent to [`Self::from_filters`] whose backing
346    /// matches [`RowSelectionPolicy::Auto`]. Selector materialization stops as
347    /// soon as the final mask strategy is known.
348    ///
349    /// # Panics
350    ///
351    /// Panics if any of the [`BooleanArray`] contain nulls.
352    pub(crate) fn from_filters_auto(filters: &[BooleanArray], threshold: usize) -> Self {
353        let total_rows = filters.iter().map(|filter| filter.len()).sum::<usize>();
354
355        // Empty selector-backed selections resolve to Mask under Auto. Preserve
356        // that decision in the backing selected by this constructor.
357        if total_rows == 0 {
358            return Self::from_boolean_buffer(BooleanBuffer::new_unset(0));
359        }
360
361        let Some(min_mask_runs) = auto_min_mask_runs(total_rows, threshold) else {
362            return Self::from_filters(filters);
363        };
364
365        match selectors_below_run_limit(filters, total_rows, min_mask_runs) {
366            Some(selectors) => Self::from_selectors(selectors),
367            None => Self::from_filters_mask(filters),
368        }
369    }
370
371    /// Creates a mask-backed [`RowSelection`] from predicate filters.
372    ///
373    /// # Panics
374    ///
375    /// Panics if any of the [`BooleanArray`] contain nulls.
376    pub(crate) fn from_filters_mask(filters: &[BooleanArray]) -> Self {
377        let mask = match filters {
378            [filter] => {
379                assert_eq!(filter.null_count(), 0);
380                filter.values().clone()
381            }
382            _ => filters_to_boolean_buffer(filters),
383        };
384        Self::from_boolean_buffer(mask)
385    }
386
387    /// Creates a [`RowSelection`] from an iterator of consecutive ranges to keep
388    pub fn from_consecutive_ranges<I: Iterator<Item = Range<usize>>>(
389        ranges: I,
390        total_rows: usize,
391    ) -> Self {
392        let mut selectors: Vec<RowSelector> = Vec::with_capacity(ranges.size_hint().0);
393        let mut last_end = 0;
394        for range in ranges {
395            let len = range.end - range.start;
396            if len == 0 {
397                continue;
398            }
399
400            match range.start.cmp(&last_end) {
401                Ordering::Equal => match selectors.last_mut() {
402                    Some(last) => last.row_count = last.row_count.checked_add(len).unwrap(),
403                    None => selectors.push(RowSelector::select(len)),
404                },
405                Ordering::Greater => {
406                    selectors.push(RowSelector::skip(range.start - last_end));
407                    selectors.push(RowSelector::select(len))
408                }
409                Ordering::Less => panic!("out of order"),
410            }
411            last_end = range.end;
412        }
413
414        if last_end != total_rows {
415            selectors.push(RowSelector::skip(total_rows - last_end))
416        }
417
418        Self::from_selectors(selectors)
419    }
420
421    /// Given an offset index, return the byte ranges for all data pages selected by `self`
422    ///
423    /// This is useful for determining what byte ranges to fetch from underlying storage
424    ///
425    /// Note: this method does not make any effort to combine consecutive ranges, nor coalesce
426    /// ranges that are close together. This is instead delegated to the IO subsystem to optimise,
427    /// e.g. `ObjectStore::get_ranges` in the [`object_store`] crate
428    ///
429    /// [`object_store`]: https://crates.io/crates/object_store
430    pub fn scan_ranges(&self, page_locations: &[PageLocation]) -> Vec<Range<u64>> {
431        match &self.inner {
432            RowSelectionInner::Selectors(selectors) => {
433                scan_ranges_from_selectors(selectors.iter().copied(), page_locations)
434            }
435            RowSelectionInner::Mask(mask) => scan_ranges_from_mask(mask.mask(), page_locations),
436        }
437    }
438
439    /// Returns the complete row ranges of the pages selected by [`Self::scan_ranges`].
440    pub(crate) fn row_ranges_for_selected_pages(
441        &self,
442        page_locations: &[PageLocation],
443        total_rows: usize,
444    ) -> Vec<Range<usize>> {
445        let mut selected_pages = self.scan_ranges(page_locations).into_iter().peekable();
446        let mut row_ranges = Vec::new();
447
448        for (idx, page) in page_locations.iter().enumerate() {
449            let Some(selected_page) = selected_pages.peek() else {
450                break;
451            };
452            if selected_page.start != page.offset as u64 {
453                continue;
454            }
455            selected_pages.next();
456
457            let end = page_locations
458                .get(idx + 1)
459                .map(|next| next.first_row_index as usize)
460                .unwrap_or(total_rows);
461            row_ranges.push(page.first_row_index as usize..end);
462        }
463
464        row_ranges
465    }
466
467    /// Splits off the first `row_count` from this [`RowSelection`]
468    pub fn split_off(&mut self, row_count: usize) -> Self {
469        match std::mem::take(&mut self.inner) {
470            RowSelectionInner::Mask(mask) => {
471                let total = mask.cached_count();
472                let (head, tail) = split_off_mask((*mask).into_mask(), row_count);
473                // Popcount only the head and derive the tail by subtraction, so
474                // repeated splits stay O(bitmap) overall.
475                let (head, tail) = match total {
476                    Some(total) => {
477                        let head_count = if tail.is_empty() {
478                            total
479                        } else {
480                            head.count_set_bits()
481                        };
482                        (
483                            MaskSelection::with_count(head, head_count),
484                            MaskSelection::with_count(tail, total - head_count),
485                        )
486                    }
487                    None => (MaskSelection::new(head), MaskSelection::new(tail)),
488                };
489                self.inner = RowSelectionInner::Mask(Box::new(tail));
490                Self::from_mask_selection(head)
491            }
492            RowSelectionInner::Selectors(selectors) => {
493                let (head, tail) = split_off_selectors(selectors, row_count);
494                self.inner = RowSelectionInner::Selectors(tail);
495                Self::from_selectors(head)
496            }
497        }
498    }
499
500    /// returns a [`RowSelection`] representing rows that are selected in both
501    /// input [`RowSelection`]s.
502    ///
503    /// This is equivalent to the logical `AND` / conjunction of the two
504    /// selections.
505    ///
506    /// # Example
507    /// If `N` means the row is not selected, and `Y` means it is
508    /// selected:
509    ///
510    /// ```text
511    /// self:     NNNNNNNNNNNNYYYYYYYYYYYYYYYYYYYYYYNNNYYYYY
512    /// other:                YYYYYNNNNYYYYYYYYYYYYY   YYNNN
513    ///
514    /// returned: NNNNNNNNNNNNYYYYYNNNNYYYYYYYYYYYYYNNNYYNNN
515    /// ```
516    ///
517    /// # Panics
518    ///
519    /// Panics if `other` does not have a length equal to the number of rows selected
520    /// by this RowSelection
521    ///
522    pub fn and_then(&self, other: &Self) -> Self {
523        match (&self.inner, &other.inner) {
524            (RowSelectionInner::Mask(mask), _) => {
525                Self::from_boolean_buffer(and_then_mask(mask.mask(), other))
526            }
527            (RowSelectionInner::Selectors(first), RowSelectionInner::Selectors(second)) => {
528                and_then_row_selections(first, second)
529            }
530            (RowSelectionInner::Selectors(first), RowSelectionInner::Mask(second)) => {
531                and_then_selectors_with_mask(first, second.mask())
532            }
533        }
534    }
535
536    /// Compute the intersection of two [`RowSelection`]
537    ///
538    /// Nonempty mixed inputs return a mask-backed result.
539    /// If only one input is empty, the result keeps the other's representation.
540    ///
541    /// For example:
542    /// self:      NNYYYYNNYYNYN
543    /// other:     NYNNNNNNY
544    ///
545    /// returned:  NNNNNNNNYYNYN
546    pub fn intersection(&self, other: &Self) -> Self {
547        match (&self.inner, &other.inner) {
548            (RowSelectionInner::Mask(l), RowSelectionInner::Mask(r)) => {
549                Self::from_boolean_buffer(intersect_masks(l.mask(), r.mask()))
550            }
551            (RowSelectionInner::Selectors(l), RowSelectionInner::Selectors(r)) => {
552                intersect_row_selections(l, r)
553            }
554            (RowSelectionInner::Selectors(selectors), RowSelectionInner::Mask(mask))
555            | (RowSelectionInner::Mask(mask), RowSelectionInner::Selectors(selectors)) => {
556                combine_mixed_selection(selectors, mask, intersect_masks)
557            }
558        }
559    }
560
561    /// Compute the union of two [`RowSelection`]
562    ///
563    /// Nonempty mixed inputs return a mask-backed result.
564    /// If only one input is empty, the result keeps the other's representation.
565    ///
566    /// For example:
567    /// self:      NNYYYYNNYYNYN
568    /// other:     NYNNNNNNN
569    ///
570    /// returned:  NYYYYYNNYYNYN
571    pub fn union(&self, other: &Self) -> Self {
572        match (&self.inner, &other.inner) {
573            (RowSelectionInner::Mask(l), RowSelectionInner::Mask(r)) => {
574                Self::from_boolean_buffer(union_masks(l.mask(), r.mask()))
575            }
576            (RowSelectionInner::Selectors(l), RowSelectionInner::Selectors(r)) => {
577                union_row_selections(l, r)
578            }
579            (RowSelectionInner::Selectors(selectors), RowSelectionInner::Mask(mask))
580            | (RowSelectionInner::Mask(mask), RowSelectionInner::Selectors(selectors)) => {
581                combine_mixed_selection(selectors, mask, union_masks)
582            }
583        }
584    }
585
586    /// Returns `true` if this [`RowSelection`] selects any rows
587    pub fn selects_any(&self) -> bool {
588        match &self.inner {
589            RowSelectionInner::Selectors(s) => s.iter().any(|x| !x.skip),
590            RowSelectionInner::Mask(m) => match m.cached_count() {
591                Some(count) => count > 0,
592                None => m.mask().set_indices().next().is_some(),
593            },
594        }
595    }
596
597    /// Trims this [`RowSelection`] removing any trailing skips
598    pub(crate) fn trim(self) -> Self {
599        match self.inner {
600            RowSelectionInner::Mask(m) => {
601                let trimmed = trim_mask(m.mask());
602                let cached_count = m.cached_count();
603                match trimmed {
604                    // Trimming only drops trailing unset bits; the count is unchanged.
605                    Some(mask) => match cached_count {
606                        Some(count) => {
607                            Self::from_mask_selection(MaskSelection::with_count(mask, count))
608                        }
609                        None => Self::from_boolean_buffer(mask),
610                    },
611                    // Nothing to trim, hand the existing box back untouched.
612                    None => Self {
613                        inner: RowSelectionInner::Mask(m),
614                    },
615                }
616            }
617            RowSelectionInner::Selectors(mut selectors) => {
618                while selectors.last().map(|x| x.skip).unwrap_or(false) {
619                    selectors.pop();
620                }
621                Self::from_selectors(selectors)
622            }
623        }
624    }
625
626    /// Applies an offset to this [`RowSelection`], skipping the first `offset` selected rows
627    pub(crate) fn offset(self, offset: usize) -> Self {
628        if offset == 0 {
629            return self;
630        }
631
632        match self.inner {
633            RowSelectionInner::Mask(mask) => {
634                let count = mask.count();
635                let buffer = offset_mask((*mask).into_mask(), offset, count);
636                Self::from_mask_selection(MaskSelection::with_count(
637                    buffer,
638                    count.saturating_sub(offset),
639                ))
640            }
641            RowSelectionInner::Selectors(selectors) => {
642                Self::from_selectors(offset_selectors(selectors, offset))
643            }
644        }
645    }
646
647    /// Limit this [`RowSelection`] to only select `limit` rows
648    pub(crate) fn limit(self, limit: usize) -> Self {
649        match self.inner {
650            RowSelectionInner::Mask(mask) => {
651                let cached = mask.cached_count();
652                let buffer = limit_mask((*mask).into_mask(), limit);
653                match cached {
654                    Some(count) => Self::from_mask_selection(MaskSelection::with_count(
655                        buffer,
656                        count.min(limit),
657                    )),
658                    None => Self::from_boolean_buffer(buffer),
659                }
660            }
661            RowSelectionInner::Selectors(selectors) => {
662                Self::from_selectors(limit_selectors(selectors, limit))
663            }
664        }
665    }
666
667    /// Returns an iterator over the [`RowSelector`]s for this
668    /// [`RowSelection`].
669    ///
670    /// Mask-backed selections materialize a `Vec<RowSelector>` cache on first
671    /// call (one allocation, `O(set_slices)` work) so the iterator can hand out
672    /// `&RowSelector`; the cache is not copied on clone. For single-pass walks
673    /// over mask-backed selections, prefer streaming directly via
674    /// [`Self::as_mask`] + [`MaskRunIter::new`] — that path is allocation-free
675    /// and avoids populating the cache.
676    pub fn iter(&self) -> impl Iterator<Item = &RowSelector> {
677        match &self.inner {
678            RowSelectionInner::Selectors(s) => s.iter(),
679            RowSelectionInner::Mask(m) => m.selectors().iter(),
680        }
681    }
682
683    /// Returns the number of selected rows
684    pub fn row_count(&self) -> usize {
685        match &self.inner {
686            RowSelectionInner::Selectors(s) => {
687                s.iter().filter(|x| !x.skip).map(|x| x.row_count).sum()
688            }
689            RowSelectionInner::Mask(m) => m.count(),
690        }
691    }
692
693    /// Returns the total number of rows spanned by this selection, both
694    /// selected and skipped
695    pub fn total_row_count(&self) -> usize {
696        match &self.inner {
697            RowSelectionInner::Selectors(s) => s.iter().map(|x| x.row_count).sum(),
698            RowSelectionInner::Mask(m) => m.mask().len(),
699        }
700    }
701
702    /// Returns the number of de-selected rows
703    pub fn skipped_row_count(&self) -> usize {
704        match &self.inner {
705            RowSelectionInner::Selectors(s) => {
706                s.iter().filter(|x| x.skip).map(|x| x.row_count).sum()
707            }
708            RowSelectionInner::Mask(m) => m.mask().len() - m.count(),
709        }
710    }
711
712    /// Expands the selection to align with batch boundaries.
713    /// This is needed when using cached array readers to ensure that
714    /// the cached data covers full batches.
715    pub(crate) fn expand_to_batch_boundaries(&self, batch_size: usize, total_rows: usize) -> Self {
716        if batch_size == 0 {
717            return self.clone();
718        }
719
720        match &self.inner {
721            RowSelectionInner::Selectors(selectors) => expand_to_batch_boundaries_from_selectors(
722                selectors.iter().copied(),
723                batch_size,
724                total_rows,
725            ),
726            RowSelectionInner::Mask(mask) => expand_to_batch_boundaries_from_selectors(
727                MaskRunIter::new(mask.mask()),
728                batch_size,
729                total_rows,
730            ),
731        }
732    }
733}
734
735/// Returns the minimum normalized run count for which the Auto selection policy
736/// would prefer a mask to RLE.
737///
738/// Auto prefers masks when the average run length is strictly below `threshold`.
739/// For positive thresholds and totals below `usize::MAX`, the first qualifying
740/// run count is `floor(total_rows / threshold) + 1`.
741///
742/// Returns `None` when selectors are always preferred for a non-empty selection.
743/// Callers handle empty selections separately.
744#[inline]
745fn auto_min_mask_runs(total_rows: usize, threshold: usize) -> Option<usize> {
746    // The strict comparison against a saturated product cannot succeed when
747    // the left-hand side is already usize::MAX.
748    if total_rows == usize::MAX {
749        return None;
750    }
751
752    let min_runs = total_rows.checked_div(threshold)?.checked_add(1)?;
753    (min_runs <= total_rows).then_some(min_runs)
754}
755
756/// Builds normalized selectors while their run count remains below
757/// `min_mask_runs`. Returning `None` drops the partial selector allocation
758/// before the caller constructs a mask.
759fn selectors_below_run_limit(
760    filters: &[BooleanArray],
761    total_rows: usize,
762    min_mask_runs: usize,
763) -> Option<Vec<RowSelector>> {
764    let mut selectors = Vec::new();
765    let mut next_offset = 0usize;
766    let mut last_end = 0usize;
767
768    for filter in filters {
769        assert_eq!(filter.null_count(), 0);
770        let offset = next_offset;
771        next_offset = next_offset.checked_add(filter.len()).unwrap();
772
773        for (start, end) in SlicesIterator::new(filter) {
774            let start = start.checked_add(offset).unwrap();
775            let end = end.checked_add(offset).unwrap();
776
777            if start > last_end {
778                append_normalized_selector(&mut selectors, RowSelector::skip(start - last_end));
779                if selectors.len() >= min_mask_runs {
780                    return None;
781                }
782            }
783
784            append_normalized_selector(&mut selectors, RowSelector::select(end - start));
785            if selectors.len() >= min_mask_runs {
786                return None;
787            }
788            last_end = end;
789        }
790    }
791
792    if last_end != total_rows {
793        append_normalized_selector(&mut selectors, RowSelector::skip(total_rows - last_end));
794        if selectors.len() >= min_mask_runs {
795            return None;
796        }
797    }
798
799    Some(selectors)
800}
801
802/// Appends a selector while maintaining the normalized selector invariants.
803fn append_normalized_selector(selectors: &mut Vec<RowSelector>, selector: RowSelector) {
804    if selector.row_count == 0 {
805        return;
806    }
807
808    match selectors.last_mut() {
809        Some(last) if last.skip == selector.skip => {
810            last.row_count = last.row_count.checked_add(selector.row_count).unwrap()
811        }
812        _ => selectors.push(selector),
813    }
814}
815
816fn filters_to_boolean_buffer(filters: &[BooleanArray]) -> BooleanBuffer {
817    let total_rows = filters.iter().map(|filter| filter.len()).sum();
818    let mut builder = BooleanBufferBuilder::new(total_rows);
819    for filter in filters {
820        assert_eq!(filter.null_count(), 0);
821        builder.append_buffer(filter.values());
822    }
823    builder.finish()
824}
825
826impl From<Vec<RowSelector>> for RowSelection {
827    fn from(selectors: Vec<RowSelector>) -> Self {
828        selectors.into_iter().collect()
829    }
830}
831
832impl From<BooleanBuffer> for RowSelection {
833    fn from(mask: BooleanBuffer) -> Self {
834        Self::from_boolean_buffer(mask)
835    }
836}
837
838impl FromIterator<RowSelector> for RowSelection {
839    fn from_iter<T: IntoIterator<Item = RowSelector>>(iter: T) -> Self {
840        let iter = iter.into_iter();
841
842        let mut selectors = Vec::with_capacity(iter.size_hint().0);
843
844        for selector in iter {
845            append_normalized_selector(&mut selectors, selector);
846        }
847
848        Self::from_selectors(selectors)
849    }
850}
851
852impl From<RowSelection> for Vec<RowSelector> {
853    fn from(r: RowSelection) -> Self {
854        r.into_selectors_vec()
855    }
856}
857
858impl From<RowSelection> for VecDeque<RowSelector> {
859    fn from(r: RowSelection) -> Self {
860        r.into_selectors_vec().into()
861    }
862}
863
864impl FromIterator<RowSelection> for RowSelection {
865    /// Concatenate multiple [`RowSelection`]s in iterator order.
866    ///
867    /// When every input is mask-backed the result stays mask-backed
868    /// (`BooleanBuffer`s are appended); otherwise falls back to flattening
869    /// through the per-`RowSelector` form.
870    fn from_iter<T: IntoIterator<Item = RowSelection>>(iter: T) -> Self {
871        let items: Vec<RowSelection> = iter.into_iter().collect();
872
873        let all_mask = items
874            .iter()
875            .all(|s| matches!(&s.inner, RowSelectionInner::Mask(_)));
876
877        if all_mask {
878            let total_len: usize = items
879                .iter()
880                .map(|s| match &s.inner {
881                    RowSelectionInner::Mask(m) => m.mask().len(),
882                    RowSelectionInner::Selectors(_) => unreachable!(),
883                })
884                .sum();
885            let mut builder = BooleanBufferBuilder::new(total_len);
886            for item in items {
887                match item.into_inner() {
888                    RowSelectionInner::Mask(m) => builder.append_buffer(m.mask()),
889                    RowSelectionInner::Selectors(_) => unreachable!(),
890                }
891            }
892            return Self::from_boolean_buffer(builder.finish());
893        }
894
895        items
896            .into_iter()
897            .flat_map(|s| s.into_selectors_vec())
898            .collect()
899    }
900}
901
902#[cfg(test)]
903mod tests {
904    use super::*;
905    use rand::rngs::StdRng;
906    use rand::{RngExt, SeedableRng};
907
908    const MAX_RANDOM_ROWS: usize = 65_536;
909    const THRESHOLDS: &[usize] = &[0, 1, 8, 16, 31, 32, 33, 64];
910    const SELECTIVITIES: &[usize] = &[0, 1, 5, 15, 50, 90, 99, 100];
911    const FILTER_SHAPES: &[FilterShape] = &[
912        FilterShape::Isolated,
913        FilterShape::Runs,
914        FilterShape::Random,
915        FilterShape::Clustered,
916    ];
917
918    #[derive(Clone, Copy, Debug)]
919    enum FilterShape {
920        Isolated,
921        Runs,
922        Random,
923        Clustered,
924    }
925
926    #[test]
927    fn auto_min_mask_runs_matches_policy_boundaries() {
928        assert_eq!(auto_min_mask_runs(0, 32), None);
929        assert_eq!(auto_min_mask_runs(64, 0), None);
930        assert_eq!(auto_min_mask_runs(64, 1), None);
931        assert_eq!(auto_min_mask_runs(64, 32), Some(3));
932        assert_eq!(auto_min_mask_runs(64, 64), Some(2));
933        assert_eq!(auto_min_mask_runs(64, 65), Some(1));
934        assert_eq!(auto_min_mask_runs(usize::MAX, usize::MAX), None);
935    }
936
937    #[test]
938    fn auto_construction_preserves_global_runs_and_threshold_boundary() {
939        let filters = vec![
940            BooleanArray::from(vec![false, true, true]),
941            BooleanArray::from(Vec::<bool>::new()),
942            BooleanArray::from(vec![true, true, false]),
943            BooleanArray::from(vec![false, false, true]),
944        ];
945
946        for threshold in THRESHOLDS {
947            assert_auto_equivalent(&filters, *threshold, "cross-filter run merge");
948        }
949
950        let run31_filters = split_evenly(&run_mask(65_536, 31, 31), 8_192);
951        let run32_filters = split_evenly(&run_mask(65_536, 32, 32), 8_192);
952
953        assert_eq!(
954            RowSelection::from_filters(&run31_filters).auto_selection_strategy(32),
955            RowSelectionStrategy::Mask
956        );
957        assert_eq!(
958            RowSelection::from_filters(&run32_filters).auto_selection_strategy(32),
959            RowSelectionStrategy::Selectors
960        );
961        assert_auto_equivalent(&run31_filters, 32, "run31 below threshold");
962        assert_auto_equivalent(&run32_filters, 32, "run32 equal to threshold");
963
964        assert_auto_equivalent(&[], 0, "no filters");
965        assert_auto_equivalent(&[BooleanArray::from(Vec::<bool>::new())], 0, "empty filter");
966    }
967
968    #[test]
969    #[cfg_attr(miri, ignore)] // Takes too long
970    fn auto_construction_edge_row_counts() {
971        for rows in [0, 1, 7, 8, 31, 32, 33, MAX_RANDOM_ROWS] {
972            let masks = [
973                ("all skipped", BooleanBuffer::new_unset(rows)),
974                ("all selected", BooleanBuffer::new_set(rows)),
975                (
976                    "alternating",
977                    BooleanBuffer::from_iter((0..rows).map(|row| row % 2 == 0)),
978                ),
979                ("short runs", run_mask(rows, 3, 5)),
980            ];
981
982            for (shape, mask) in masks {
983                let mask = with_bit_offset(mask, 5);
984                let filters = split_evenly(&mask, rows.div_ceil(3).max(1));
985
986                for &threshold in THRESHOLDS {
987                    let context = format!("rows={rows} shape={shape} threshold={threshold}");
988                    assert_auto_equivalent(&filters, threshold, &context);
989                }
990            }
991        }
992    }
993
994    #[test]
995    #[cfg_attr(miri, ignore)] // Takes too long
996    fn auto_construction_randomized_equivalence() {
997        let mut rng = StdRng::seed_from_u64(0x1077_6000_5eed);
998        let mut case_idx = 0usize;
999
1000        for &threshold in THRESHOLDS {
1001            for &selectivity in SELECTIVITIES {
1002                for &shape in FILTER_SHAPES {
1003                    for with_offset in [false, true] {
1004                        let rows = rng.random_range(0..=MAX_RANDOM_ROWS);
1005                        let mask = random_shape(&mut rng, rows, selectivity, shape);
1006                        let bit_offset = if with_offset {
1007                            rng.random_range(1..=63)
1008                        } else {
1009                            0
1010                        };
1011                        let mask = with_bit_offset(mask, bit_offset);
1012                        let filter_count = rng.random_range(1..=32);
1013                        let filters = random_split(&mut rng, &mask, filter_count);
1014                        let context = format!(
1015                            "case={case_idx} rows={rows} selectivity={selectivity} shape={shape:?} \
1016                             filters={filter_count} threshold={threshold} bit_offset={bit_offset}"
1017                        );
1018
1019                        assert_auto_equivalent(&filters, threshold, &context);
1020                        case_idx += 1;
1021                    }
1022                }
1023            }
1024        }
1025    }
1026
1027    fn assert_auto_equivalent(filters: &[BooleanArray], threshold: usize, context: &str) {
1028        let reference = RowSelection::from_filters(filters);
1029        let reference_strategy = reference.auto_selection_strategy(threshold);
1030        let auto_built = RowSelection::from_filters_auto(filters, threshold);
1031        let auto_built_strategy = auto_built.auto_selection_strategy(threshold);
1032        let auto_built_backing = match &auto_built.inner {
1033            RowSelectionInner::Mask(_) => RowSelectionStrategy::Mask,
1034            RowSelectionInner::Selectors(_) => RowSelectionStrategy::Selectors,
1035        };
1036
1037        assert_eq!(reference_strategy, auto_built_backing, "backing: {context}");
1038        assert_eq!(
1039            reference_strategy, auto_built_strategy,
1040            "strategy: {context}"
1041        );
1042        assert_eq!(reference, auto_built, "logical selection: {context}");
1043    }
1044
1045    fn run_mask(rows: usize, selected_run: usize, skipped_run: usize) -> BooleanBuffer {
1046        let period = selected_run + skipped_run;
1047        BooleanBuffer::from_iter((0..rows).map(|row| row % period < selected_run))
1048    }
1049
1050    fn split_evenly(mask: &BooleanBuffer, batch_size: usize) -> Vec<BooleanArray> {
1051        (0..mask.len())
1052            .step_by(batch_size)
1053            .map(|offset| {
1054                let len = batch_size.min(mask.len() - offset);
1055                BooleanArray::new(mask.slice(offset, len), None)
1056            })
1057            .collect()
1058    }
1059
1060    fn random_shape(
1061        rng: &mut StdRng,
1062        rows: usize,
1063        selectivity: usize,
1064        shape: FilterShape,
1065    ) -> BooleanBuffer {
1066        if selectivity == 0 {
1067            return BooleanBuffer::new_unset(rows);
1068        }
1069        if selectivity == 100 {
1070            return BooleanBuffer::new_set(rows);
1071        }
1072
1073        match shape {
1074            FilterShape::Isolated => isolated_mask(rows, selectivity),
1075            FilterShape::Runs => {
1076                let scale = rng.random_range(1..=8);
1077                run_mask(rows, selectivity * scale, (100 - selectivity) * scale)
1078            }
1079            FilterShape::Random => BooleanBuffer::from_iter(
1080                (0..rows).map(|_| rng.random_bool(selectivity as f64 / 100.0)),
1081            ),
1082            FilterShape::Clustered => one_cluster_mask(rng, rows, selectivity),
1083        }
1084    }
1085
1086    fn isolated_mask(rows: usize, selectivity: usize) -> BooleanBuffer {
1087        if selectivity <= 50 {
1088            let period = 100usize.div_ceil(selectivity).max(2);
1089            BooleanBuffer::from_iter((0..rows).map(|row| row % period == 0))
1090        } else {
1091            let period = 100usize.div_ceil(100 - selectivity).max(2);
1092            BooleanBuffer::from_iter((0..rows).map(|row| row % period != 0))
1093        }
1094    }
1095
1096    fn one_cluster_mask(rng: &mut StdRng, rows: usize, selectivity: usize) -> BooleanBuffer {
1097        let selected = rows.saturating_mul(selectivity) / 100;
1098        let start = rng.random_range(0..=rows - selected);
1099        let mut builder = BooleanBufferBuilder::new(rows);
1100        builder.append_n(start, false);
1101        builder.append_n(selected, true);
1102        builder.append_n(rows - start - selected, false);
1103        builder.finish()
1104    }
1105
1106    fn with_bit_offset(mask: BooleanBuffer, offset: usize) -> BooleanBuffer {
1107        if offset == 0 {
1108            return mask;
1109        }
1110
1111        let len = mask.len();
1112        let mut builder = BooleanBufferBuilder::new(offset + len);
1113        builder.append_n(offset, false);
1114        builder.append_buffer(&mask);
1115        builder.finish().slice(offset, len)
1116    }
1117
1118    fn random_split(
1119        rng: &mut StdRng,
1120        mask: &BooleanBuffer,
1121        filter_count: usize,
1122    ) -> Vec<BooleanArray> {
1123        let mut filters = Vec::with_capacity(filter_count);
1124        let mut offset = 0usize;
1125
1126        for index in 0..filter_count {
1127            let remaining = mask.len() - offset;
1128            let len = if index + 1 == filter_count {
1129                remaining
1130            } else {
1131                rng.random_range(0..=remaining)
1132            };
1133            filters.push(BooleanArray::new(mask.slice(offset, len), None));
1134            offset += len;
1135        }
1136
1137        filters
1138    }
1139
1140    #[test]
1141    fn test_force_mask_from_selectors() {
1142        let cases = [
1143            (vec![], vec![]),
1144            (vec![RowSelector::select(9)], vec![true; 9]),
1145            (vec![RowSelector::skip(9)], vec![false; 9]),
1146            (
1147                vec![
1148                    RowSelector::skip(3),
1149                    RowSelector::select(7),
1150                    RowSelector::skip(2),
1151                    RowSelector::select(1),
1152                    RowSelector::skip(4),
1153                ],
1154                vec![
1155                    false, false, false, true, true, true, true, true, true, true, false, false,
1156                    true, false, false, false, false,
1157                ],
1158            ),
1159        ];
1160
1161        for (selectors, bits) in cases {
1162            let original = RowSelection::from(selectors);
1163            let selection = original.clone().force_mask();
1164            assert_eq!(selection.as_mask().unwrap(), &BooleanBuffer::from(bits));
1165            assert_eq!(selection.total_row_count(), original.total_row_count());
1166            assert_eq!(selection.row_count(), original.row_count());
1167            assert_eq!(selection.skipped_row_count(), original.skipped_row_count());
1168            assert_eq!(selection, original);
1169        }
1170    }
1171
1172    #[test]
1173    fn test_force_mask_preserves_existing_mask() {
1174        let mask = BooleanBuffer::from(vec![
1175            false, true, false, true, true, false, true, false, false, true, false,
1176        ])
1177        .slice(3, 7);
1178        let selection = RowSelection::from_boolean_buffer(mask.clone()).force_mask();
1179        assert!(selection.as_mask().unwrap().ptr_eq(&mask));
1180        assert_eq!(selection.row_count(), 4);
1181
1182        // Repeated conversion also preserves the buffer and its bit offset.
1183        let selection = selection.force_mask();
1184        assert!(selection.as_mask().unwrap().ptr_eq(&mask));
1185        assert_eq!(selection.total_row_count(), 7);
1186    }
1187
1188    #[test]
1189    fn test_force_mask_preserves_bitmap_intersection() {
1190        let existing = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
1191            true, false, true, true, false, true, true, false, true,
1192        ]));
1193        let selection = RowSelection::from(vec![
1194            RowSelector::skip(2),
1195            RowSelector::select(5),
1196            RowSelector::skip(2),
1197        ]);
1198        let expected = existing.intersection(&selection);
1199        let actual = existing.intersection(&selection.force_mask());
1200
1201        assert!(actual.as_mask().is_some());
1202        assert_eq!(actual, expected);
1203        assert_eq!(actual.total_row_count(), 9);
1204        assert_eq!(actual.row_count(), 4);
1205    }
1206
1207    #[test]
1208    fn test_total_row_count() {
1209        let selection = RowSelection::from(vec![RowSelector::skip(5), RowSelector::select(3)]);
1210        assert_eq!(selection.total_row_count(), 8);
1211        assert_eq!(selection.row_count(), 3);
1212        assert_eq!(selection.skipped_row_count(), 5);
1213
1214        let selection =
1215            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true]));
1216        assert_eq!(selection.total_row_count(), 3);
1217        assert_eq!(selection.row_count(), 2);
1218        assert_eq!(selection.skipped_row_count(), 1);
1219
1220        let empty = RowSelection::from(vec![]);
1221        assert_eq!(empty.total_row_count(), 0);
1222    }
1223
1224    #[test]
1225    fn test_offset_zero_and_zero_batch_expand_are_identity() {
1226        let selection =
1227            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true]));
1228        assert_eq!(selection.clone().offset(0), selection);
1229        assert_eq!(selection.expand_to_batch_boundaries(0, 3), selection);
1230    }
1231
1232    #[test]
1233    fn test_from_filters() {
1234        let filters = vec![
1235            BooleanArray::from(vec![false, false, false, true, true, true, true]),
1236            BooleanArray::from(vec![true, true, false, false, true, true, true]),
1237            BooleanArray::from(vec![false, false, false, false]),
1238            BooleanArray::from(Vec::<bool>::new()),
1239        ];
1240
1241        let selection = RowSelection::from_filters(&filters[..1]);
1242        assert!(selection.selects_any());
1243        assert_eq!(
1244            selection.selectors(),
1245            vec![RowSelector::skip(3), RowSelector::select(4)]
1246        );
1247
1248        let selection = RowSelection::from_filters(&filters[..2]);
1249        assert!(selection.selects_any());
1250        assert_eq!(
1251            selection.selectors(),
1252            vec![
1253                RowSelector::skip(3),
1254                RowSelector::select(6),
1255                RowSelector::skip(2),
1256                RowSelector::select(3)
1257            ]
1258        );
1259
1260        let selection = RowSelection::from_filters(&filters);
1261        assert!(selection.selects_any());
1262        assert_eq!(
1263            selection.selectors(),
1264            vec![
1265                RowSelector::skip(3),
1266                RowSelector::select(6),
1267                RowSelector::skip(2),
1268                RowSelector::select(3),
1269                RowSelector::skip(4)
1270            ]
1271        );
1272
1273        let selection = RowSelection::from_filters(&filters[2..3]);
1274        assert!(!selection.selects_any());
1275        assert_eq!(selection.selectors(), vec![RowSelector::skip(4)]);
1276    }
1277
1278    #[test]
1279    fn test_iter() {
1280        // use the iter() API to show it does what is expected and
1281        // avoid accidental deletion
1282        let selectors = vec![
1283            RowSelector::select(3),
1284            RowSelector::skip(33),
1285            RowSelector::select(4),
1286        ];
1287
1288        let round_tripped: Vec<RowSelector> = RowSelection::from(selectors.clone())
1289            .iter()
1290            .copied()
1291            .collect();
1292        assert_eq!(selectors, round_tripped);
1293    }
1294
1295    #[test]
1296    fn test_row_count() {
1297        let selection = RowSelection::from(vec![
1298            RowSelector::skip(34),
1299            RowSelector::select(12),
1300            RowSelector::skip(3),
1301            RowSelector::select(35),
1302        ]);
1303
1304        assert_eq!(selection.row_count(), 12 + 35);
1305        assert_eq!(selection.skipped_row_count(), 34 + 3);
1306
1307        let selection = RowSelection::from(vec![RowSelector::select(12), RowSelector::select(35)]);
1308
1309        assert_eq!(selection.row_count(), 12 + 35);
1310        assert_eq!(selection.skipped_row_count(), 0);
1311
1312        let selection = RowSelection::from(vec![RowSelector::skip(34), RowSelector::skip(3)]);
1313
1314        assert_eq!(selection.row_count(), 0);
1315        assert_eq!(selection.skipped_row_count(), 34 + 3);
1316
1317        let selection = RowSelection::from(vec![]);
1318
1319        assert_eq!(selection.row_count(), 0);
1320        assert_eq!(selection.skipped_row_count(), 0);
1321    }
1322
1323    #[test]
1324    fn test_mixed_backing_equality_mismatches() {
1325        let mask =
1326            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true, true]));
1327
1328        // Total row counts differ
1329        let longer = RowSelection::from(vec![
1330            RowSelector::select(1),
1331            RowSelector::skip(1),
1332            RowSelector::select(2),
1333            RowSelector::skip(1),
1334        ]);
1335        assert_ne!(mask, longer);
1336        assert_ne!(longer, mask);
1337
1338        // A selected bit falls inside a skip run
1339        let skip_overlap = RowSelection::from(vec![RowSelector::skip(2), RowSelector::select(2)]);
1340        assert_ne!(mask, skip_overlap);
1341
1342        // Select run boundaries do not line up
1343        let misaligned = RowSelection::from(vec![
1344            RowSelector::select(2),
1345            RowSelector::skip(1),
1346            RowSelector::select(1),
1347        ]);
1348        assert_ne!(mask, misaligned);
1349
1350        let equal = RowSelection::from(vec![
1351            RowSelector::select(1),
1352            RowSelector::skip(1),
1353            RowSelector::select(2),
1354        ]);
1355        assert_eq!(mask, equal);
1356        assert_eq!(equal, mask);
1357    }
1358
1359    #[test]
1360    fn test_from_iter_all_mask_preserves_mask_backing() {
1361        let a_bits = vec![true, false, true, true];
1362        let b_bits = vec![false, true, false];
1363        let c_bits = vec![true, true, false, false, true];
1364
1365        let parts = vec![
1366            RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())),
1367            RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())),
1368            RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())),
1369        ];
1370        let collected: RowSelection = parts.into_iter().collect();
1371
1372        let combined = a_bits
1373            .iter()
1374            .chain(b_bits.iter())
1375            .chain(c_bits.iter())
1376            .copied()
1377            .collect::<Vec<_>>();
1378        let expected = RowSelection::from_filters(&[BooleanArray::from(combined)]);
1379
1380        assert!(collected.as_mask().is_some());
1381        assert_eq!(collected, expected);
1382    }
1383
1384    #[test]
1385    fn test_from_iter_mixed_backing_falls_back_to_selectors() {
1386        let a_bits = vec![true, false, true];
1387        let b_selectors = vec![RowSelector::skip(2), RowSelector::select(3)];
1388        let c_bits = vec![false, true];
1389
1390        let parts = vec![
1391            RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())),
1392            RowSelection::from(b_selectors),
1393            RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())),
1394        ];
1395        let collected: RowSelection = parts.into_iter().collect();
1396
1397        assert!(collected.as_mask().is_none());
1398
1399        let combined_bits = vec![
1400            true, false, true, false, false, true, true, true, false, true,
1401        ];
1402        let expected = RowSelection::from_filters(&[BooleanArray::from(combined_bits)]);
1403        assert_eq!(collected, expected);
1404    }
1405
1406    #[test]
1407    fn test_from_iter_empty_yields_empty_selection() {
1408        let collected: RowSelection = std::iter::empty::<RowSelection>().collect();
1409        assert_eq!(collected, RowSelection::default());
1410        assert!(collected.as_mask().is_some());
1411        assert_eq!(collected.as_mask().unwrap().len(), 0);
1412    }
1413}
parquet/arrow/arrow_reader/selection/mod.rs

parquet/arrow/arrow_reader/selection/
mod.rs
