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

parquet/arrow/arrow_reader/selection/
mod.rs
