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, intersect_masks,
48    intersect_row_selections, union_masks, union_row_selections,
49};
50pub use boolean::MaskRunIter;
51use boolean::{
52    MaskSelection, limit_mask, mask_has_at_least_runs, offset_mask, split_off_mask, trim_mask,
53};
54pub(crate) use cursor::{LoadedRowRanges, MaskCursor, RowSelectionStrategy};
55pub use cursor::{RowSelectionCursor, RowSelectionPolicy};
56use ranges::{expand_to_batch_boundaries_from_selectors, scan_ranges_from_selectors};
57pub use selector::RowSelector;
58use selector::{limit_selectors, offset_selectors, split_off_selectors};
59
60/// [`RowSelection`] represents selecting a subset of rows
61/// when scanning a parquet file.
62///
63/// This is applied prior to reading column data, and can therefore
64/// be used to skip IO to fetch data into memory
65///
66/// A typical use-case would be using the [`PageIndex`] to filter out rows
67/// that don't satisfy a predicate
68///
69/// Depending on the pattern of rows to be selected, [`RowSelection`] has
70/// either a bitmap or an RLE ([`RowSelector`]) based implementation.
71///
72/// # Example
73/// ```
74/// use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
75///
76/// let selectors = vec![
77///     RowSelector::skip(5),
78///     RowSelector::select(5),
79///     RowSelector::select(5),
80///     RowSelector::skip(5),
81/// ];
82///
83/// // Creating a selection will combine adjacent selectors
84/// let selection: RowSelection = selectors.into();
85///
86/// let expected = vec![
87///     RowSelector::skip(5),
88///     RowSelector::select(10),
89///     RowSelector::skip(5),
90/// ];
91///
92/// let actual: Vec<RowSelector> = selection.into();
93/// assert_eq!(actual, expected);
94///
95/// // you can also create a selection from consecutive ranges
96/// let ranges = vec![5..10, 10..15];
97/// let selection =
98///   RowSelection::from_consecutive_ranges(ranges.into_iter(), 20);
99/// let actual: Vec<RowSelector> = selection.into();
100/// assert_eq!(actual, expected);
101///
102/// // or directly from a packed bitmap, when the upstream producer already
103/// // has one. The bitmap is kept as-is rather than run-length-encoded.
104/// use arrow_buffer::BooleanBuffer;
105/// let mask = BooleanBuffer::from(vec![true, false, true, true]);
106/// let selection = RowSelection::from_boolean_buffer(mask);
107/// assert_eq!(selection.row_count(), 3);
108/// ```
109///
110/// An RLE ([`RowSelector`]) backed [`RowSelection`] maintains the following
111/// invariants (they do not apply to the bitmap backed implementation):
112///
113/// * It contains no [`RowSelector`] of 0 rows
114/// * Consecutive [`RowSelector`]s alternate skipping or selecting rows
115///
116/// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData
117#[derive(Default, Clone)]
118pub struct RowSelection {
119    inner: RowSelectionInner,
120}
121
122/// Internal storage for [`RowSelection`].
123#[derive(Debug, Clone)]
124pub(crate) enum RowSelectionInner {
125    Selectors(Vec<RowSelector>),
126    Mask(Box<MaskSelection>),
127}
128
129impl Default for RowSelectionInner {
130    fn default() -> Self {
131        Self::Selectors(Vec::new())
132    }
133}
134
135impl std::fmt::Debug for RowSelection {
136    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
137        match &self.inner {
138            RowSelectionInner::Selectors(s) => f
139                .debug_struct("RowSelection")
140                .field("selectors", s)
141                .finish(),
142            RowSelectionInner::Mask(m) => f
143                .debug_struct("RowSelection")
144                .field("mask_len", &m.mask().len())
145                .finish_non_exhaustive(),
146        }
147    }
148}
149
150impl PartialEq for RowSelection {
151    fn eq(&self, other: &Self) -> bool {
152        match (&self.inner, &other.inner) {
153            (RowSelectionInner::Selectors(a), RowSelectionInner::Selectors(b)) => a == b,
154            (RowSelectionInner::Mask(a), RowSelectionInner::Mask(b)) => a.mask() == b.mask(),
155            (RowSelectionInner::Mask(mask), RowSelectionInner::Selectors(selectors))
156            | (RowSelectionInner::Selectors(selectors), RowSelectionInner::Mask(mask)) => {
157                if selectors
158                    .iter()
159                    .try_fold(0usize, |acc, selector| acc.checked_add(selector.row_count))
160                    != Some(mask.mask().len())
161                {
162                    return false;
163                }
164
165                let mut slices = mask.mask().set_slices().peekable();
166                let mut cursor = 0usize;
167
168                for selector in selectors {
169                    let end = cursor + selector.row_count;
170
171                    if selector.skip {
172                        if slices.peek().is_some_and(|(start, _)| *start < end) {
173                            return false;
174                        }
175                    } else {
176                        match slices.next() {
177                            Some((start, slice_end)) if start == cursor && slice_end == end => {}
178                            _ => return false,
179                        }
180                    }
181
182                    cursor = end;
183                }
184
185                slices.next().is_none()
186            }
187        }
188    }
189}
190
191impl Eq for RowSelection {}
192
193impl RowSelection {
194    /// Not `pub`: unlike `From<Vec<RowSelector>>`, this performs no
195    /// validation/normalization of the selectors (e.g. combining adjacent
196    /// selectors), so callers must uphold the invariants themselves.
197    fn from_selectors(selectors: Vec<RowSelector>) -> Self {
198        Self {
199            inner: RowSelectionInner::Selectors(selectors),
200        }
201    }
202
203    /// Create a [`RowSelection`] from a packed [`BooleanBuffer`].
204    ///
205    /// Each set bit selects a row, each unset bit skips one. Unlike
206    /// [`Self::from_filters`], the bitmap is kept as-is rather than
207    /// eagerly run-length-encoded. [`Self::iter`] materializes and caches the
208    /// RLE form on first use; use [`MaskRunIter`] to stream the RLE form
209    /// directly from the bitmap.
210    pub fn from_boolean_buffer(mask: BooleanBuffer) -> Self {
211        Self {
212            inner: RowSelectionInner::Mask(Box::new(MaskSelection::new(mask))),
213        }
214    }
215
216    fn from_mask_selection(mask: MaskSelection) -> Self {
217        Self {
218            inner: RowSelectionInner::Mask(Box::new(mask)),
219        }
220    }
221
222    /// Returns the underlying mask if this selection is mask-backed.
223    ///
224    /// Public so that engines composing selections (e.g. DataFusion's
225    /// [`ParquetAccessPlan::into_overall_row_selection`]) can concatenate
226    /// mask-backed selections without materialising the RLE form.
227    ///
228    /// [`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
229    pub fn as_mask(&self) -> Option<&BooleanBuffer> {
230        match &self.inner {
231            RowSelectionInner::Mask(m) => Some(m.mask()),
232            RowSelectionInner::Selectors(_) => None,
233        }
234    }
235
236    /// Consume the selection and return its internal storage.
237    pub(crate) fn into_inner(self) -> RowSelectionInner {
238        self.inner
239    }
240
241    /// Choose the automatic materialisation strategy without converting between
242    /// selector and mask backing.
243    #[inline]
244    pub(crate) fn auto_selection_strategy(&self, threshold: usize) -> RowSelectionStrategy {
245        let (total_rows, effective_count) = match &self.inner {
246            RowSelectionInner::Selectors(selectors) => {
247                selectors.iter().fold((0usize, 0usize), |(rows, count), s| {
248                    if s.row_count > 0 {
249                        (rows + s.row_count, count + 1)
250                    } else {
251                        (rows, count)
252                    }
253                })
254            }
255            RowSelectionInner::Mask(mask) => {
256                let mask = mask.mask();
257                let total_rows = mask.len();
258
259                if total_rows == 0 {
260                    return RowSelectionStrategy::Mask;
261                }
262
263                // A mask is preferred when:
264                //
265                // total_rows < run_count * threshold
266                //
267                // Therefore only scan until the first run count that can make
268                // the inequality true. Fragmented masks normally reach this
269                // boundary near the start instead of enumerating every run.
270                let min_mask_runs = total_rows
271                    .checked_div(threshold)
272                    .and_then(|max_selector_runs| max_selector_runs.checked_add(1));
273
274                return match min_mask_runs {
275                    Some(min_runs) if mask_has_at_least_runs(mask, min_runs) => {
276                        RowSelectionStrategy::Mask
277                    }
278                    _ => RowSelectionStrategy::Selectors,
279                };
280            }
281        };
282
283        if effective_count == 0 {
284            return RowSelectionStrategy::Mask;
285        }
286
287        if total_rows < effective_count.saturating_mul(threshold) {
288            RowSelectionStrategy::Mask
289        } else {
290            RowSelectionStrategy::Selectors
291        }
292    }
293
294    #[cfg(test)]
295    fn selectors(&self) -> Vec<RowSelector> {
296        self.iter().copied().collect()
297    }
298
299    fn into_selectors_vec(self) -> Vec<RowSelector> {
300        match self.inner {
301            RowSelectionInner::Selectors(s) => s,
302            RowSelectionInner::Mask(m) => (*m).into_selectors(),
303        }
304    }
305
306    /// Creates a [`RowSelection`] from a slice of [`BooleanArray`]
307    ///
308    /// # Panics
309    ///
310    /// Panics if any of the [`BooleanArray`] contain nulls
311    pub fn from_filters(filters: &[BooleanArray]) -> Self {
312        let mut next_offset = 0;
313        let total_rows = filters.iter().map(|x| x.len()).sum();
314
315        let iter = filters.iter().flat_map(|filter| {
316            let offset = next_offset;
317            next_offset += filter.len();
318            assert_eq!(filter.null_count(), 0);
319            SlicesIterator::new(filter).map(move |(start, end)| start + offset..end + offset)
320        });
321
322        Self::from_consecutive_ranges(iter, total_rows)
323    }
324
325    /// Creates a [`RowSelection`] from an iterator of consecutive ranges to keep
326    pub fn from_consecutive_ranges<I: Iterator<Item = Range<usize>>>(
327        ranges: I,
328        total_rows: usize,
329    ) -> Self {
330        let mut selectors: Vec<RowSelector> = Vec::with_capacity(ranges.size_hint().0);
331        let mut last_end = 0;
332        for range in ranges {
333            let len = range.end - range.start;
334            if len == 0 {
335                continue;
336            }
337
338            match range.start.cmp(&last_end) {
339                Ordering::Equal => match selectors.last_mut() {
340                    Some(last) => last.row_count = last.row_count.checked_add(len).unwrap(),
341                    None => selectors.push(RowSelector::select(len)),
342                },
343                Ordering::Greater => {
344                    selectors.push(RowSelector::skip(range.start - last_end));
345                    selectors.push(RowSelector::select(len))
346                }
347                Ordering::Less => panic!("out of order"),
348            }
349            last_end = range.end;
350        }
351
352        if last_end != total_rows {
353            selectors.push(RowSelector::skip(total_rows - last_end))
354        }
355
356        Self::from_selectors(selectors)
357    }
358
359    /// Given an offset index, return the byte ranges for all data pages selected by `self`
360    ///
361    /// This is useful for determining what byte ranges to fetch from underlying storage
362    ///
363    /// Note: this method does not make any effort to combine consecutive ranges, nor coalesce
364    /// ranges that are close together. This is instead delegated to the IO subsystem to optimise,
365    /// e.g. `ObjectStore::get_ranges` in the [`object_store`] crate
366    ///
367    /// [`object_store`]: https://crates.io/crates/object_store
368    pub fn scan_ranges(&self, page_locations: &[PageLocation]) -> Vec<Range<u64>> {
369        match &self.inner {
370            RowSelectionInner::Selectors(selectors) => {
371                scan_ranges_from_selectors(selectors.iter().copied(), page_locations)
372            }
373            RowSelectionInner::Mask(mask) => {
374                scan_ranges_from_selectors(MaskRunIter::new(mask.mask()), page_locations)
375            }
376        }
377    }
378
379    /// Returns the complete row ranges of the pages selected by [`Self::scan_ranges`].
380    pub(crate) fn row_ranges_for_selected_pages(
381        &self,
382        page_locations: &[PageLocation],
383        total_rows: usize,
384    ) -> Vec<Range<usize>> {
385        let mut selected_pages = self.scan_ranges(page_locations).into_iter().peekable();
386        let mut row_ranges = Vec::new();
387
388        for (idx, page) in page_locations.iter().enumerate() {
389            let Some(selected_page) = selected_pages.peek() else {
390                break;
391            };
392            if selected_page.start != page.offset as u64 {
393                continue;
394            }
395            selected_pages.next();
396
397            let end = page_locations
398                .get(idx + 1)
399                .map(|next| next.first_row_index as usize)
400                .unwrap_or(total_rows);
401            row_ranges.push(page.first_row_index as usize..end);
402        }
403
404        row_ranges
405    }
406
407    /// Splits off the first `row_count` from this [`RowSelection`]
408    pub fn split_off(&mut self, row_count: usize) -> Self {
409        match std::mem::take(&mut self.inner) {
410            RowSelectionInner::Mask(mask) => {
411                let total = mask.cached_count();
412                let (head, tail) = split_off_mask((*mask).into_mask(), row_count);
413                // Popcount only the head and derive the tail by subtraction, so
414                // repeated splits stay O(bitmap) overall.
415                let (head, tail) = match total {
416                    Some(total) => {
417                        let head_count = if tail.is_empty() {
418                            total
419                        } else {
420                            head.count_set_bits()
421                        };
422                        (
423                            MaskSelection::with_count(head, head_count),
424                            MaskSelection::with_count(tail, total - head_count),
425                        )
426                    }
427                    None => (MaskSelection::new(head), MaskSelection::new(tail)),
428                };
429                self.inner = RowSelectionInner::Mask(Box::new(tail));
430                Self::from_mask_selection(head)
431            }
432            RowSelectionInner::Selectors(selectors) => {
433                let (head, tail) = split_off_selectors(selectors, row_count);
434                self.inner = RowSelectionInner::Selectors(tail);
435                Self::from_selectors(head)
436            }
437        }
438    }
439
440    /// returns a [`RowSelection`] representing rows that are selected in both
441    /// input [`RowSelection`]s.
442    ///
443    /// This is equivalent to the logical `AND` / conjunction of the two
444    /// selections.
445    ///
446    /// # Example
447    /// If `N` means the row is not selected, and `Y` means it is
448    /// selected:
449    ///
450    /// ```text
451    /// self:     NNNNNNNNNNNNYYYYYYYYYYYYYYYYYYYYYYNNNYYYYY
452    /// other:                YYYYYNNNNYYYYYYYYYYYYY   YYNNN
453    ///
454    /// returned: NNNNNNNNNNNNYYYYYNNNNYYYYYYYYYYYYYNNNYYNNN
455    /// ```
456    ///
457    /// # Panics
458    ///
459    /// Panics if `other` does not have a length equal to the number of rows selected
460    /// by this RowSelection
461    ///
462    pub fn and_then(&self, other: &Self) -> Self {
463        match (&self.inner, &other.inner) {
464            (RowSelectionInner::Mask(mask), _) => {
465                Self::from_boolean_buffer(and_then_mask(mask.mask(), other))
466            }
467            (RowSelectionInner::Selectors(first), RowSelectionInner::Selectors(second)) => {
468                and_then_row_selections(first, second)
469            }
470            (RowSelectionInner::Selectors(first), RowSelectionInner::Mask(second)) => {
471                and_then_selectors_with_mask(first, second.mask())
472            }
473        }
474    }
475
476    /// Compute the intersection of two [`RowSelection`]
477    /// For example:
478    /// self:      NNYYYYNNYYNYN
479    /// other:     NYNNNNNNY
480    ///
481    /// returned:  NNNNNNNNYYNYN
482    pub fn intersection(&self, other: &Self) -> Self {
483        match (&self.inner, &other.inner) {
484            (RowSelectionInner::Mask(l), RowSelectionInner::Mask(r)) => {
485                Self::from_boolean_buffer(intersect_masks(l.mask(), r.mask()))
486            }
487            (RowSelectionInner::Selectors(l), RowSelectionInner::Selectors(r)) => {
488                intersect_row_selections(l, r)
489            }
490            (RowSelectionInner::Selectors(l), RowSelectionInner::Mask(r)) => {
491                intersect_row_selections(l, &r.borrowed_selectors())
492            }
493            (RowSelectionInner::Mask(l), RowSelectionInner::Selectors(r)) => {
494                intersect_row_selections(&l.borrowed_selectors(), r)
495            }
496        }
497    }
498
499    /// Compute the union of two [`RowSelection`]
500    /// For example:
501    /// self:      NNYYYYNNYYNYN
502    /// other:     NYNNNNNNN
503    ///
504    /// returned:  NYYYYYNNYYNYN
505    pub fn union(&self, other: &Self) -> Self {
506        match (&self.inner, &other.inner) {
507            (RowSelectionInner::Mask(l), RowSelectionInner::Mask(r)) => {
508                Self::from_boolean_buffer(union_masks(l.mask(), r.mask()))
509            }
510            (RowSelectionInner::Selectors(l), RowSelectionInner::Selectors(r)) => {
511                union_row_selections(l, r)
512            }
513            (RowSelectionInner::Selectors(l), RowSelectionInner::Mask(r)) => {
514                union_row_selections(l, &r.borrowed_selectors())
515            }
516            (RowSelectionInner::Mask(l), RowSelectionInner::Selectors(r)) => {
517                union_row_selections(&l.borrowed_selectors(), r)
518            }
519        }
520    }
521
522    /// Returns `true` if this [`RowSelection`] selects any rows
523    pub fn selects_any(&self) -> bool {
524        match &self.inner {
525            RowSelectionInner::Selectors(s) => s.iter().any(|x| !x.skip),
526            RowSelectionInner::Mask(m) => match m.cached_count() {
527                Some(count) => count > 0,
528                None => m.mask().set_indices().next().is_some(),
529            },
530        }
531    }
532
533    /// Trims this [`RowSelection`] removing any trailing skips
534    pub(crate) fn trim(self) -> Self {
535        match self.inner {
536            RowSelectionInner::Mask(m) => {
537                let trimmed = trim_mask(m.mask());
538                let cached_count = m.cached_count();
539                match trimmed {
540                    // Trimming only drops trailing unset bits; the count is unchanged.
541                    Some(mask) => match cached_count {
542                        Some(count) => {
543                            Self::from_mask_selection(MaskSelection::with_count(mask, count))
544                        }
545                        None => Self::from_boolean_buffer(mask),
546                    },
547                    // Nothing to trim, hand the existing box back untouched.
548                    None => Self {
549                        inner: RowSelectionInner::Mask(m),
550                    },
551                }
552            }
553            RowSelectionInner::Selectors(mut selectors) => {
554                while selectors.last().map(|x| x.skip).unwrap_or(false) {
555                    selectors.pop();
556                }
557                Self::from_selectors(selectors)
558            }
559        }
560    }
561
562    /// Applies an offset to this [`RowSelection`], skipping the first `offset` selected rows
563    pub(crate) fn offset(self, offset: usize) -> Self {
564        if offset == 0 {
565            return self;
566        }
567
568        match self.inner {
569            RowSelectionInner::Mask(mask) => {
570                let count = mask.count();
571                let buffer = offset_mask((*mask).into_mask(), offset, count);
572                Self::from_mask_selection(MaskSelection::with_count(
573                    buffer,
574                    count.saturating_sub(offset),
575                ))
576            }
577            RowSelectionInner::Selectors(selectors) => {
578                Self::from_selectors(offset_selectors(selectors, offset))
579            }
580        }
581    }
582
583    /// Limit this [`RowSelection`] to only select `limit` rows
584    pub(crate) fn limit(self, limit: usize) -> Self {
585        match self.inner {
586            RowSelectionInner::Mask(mask) => {
587                let cached = mask.cached_count();
588                let buffer = limit_mask((*mask).into_mask(), limit);
589                match cached {
590                    Some(count) => Self::from_mask_selection(MaskSelection::with_count(
591                        buffer,
592                        count.min(limit),
593                    )),
594                    None => Self::from_boolean_buffer(buffer),
595                }
596            }
597            RowSelectionInner::Selectors(selectors) => {
598                Self::from_selectors(limit_selectors(selectors, limit))
599            }
600        }
601    }
602
603    /// Returns an iterator over the [`RowSelector`]s for this
604    /// [`RowSelection`].
605    ///
606    /// Mask-backed selections materialize a `Vec<RowSelector>` cache on first
607    /// call (one allocation, `O(set_slices)` work) so the iterator can hand out
608    /// `&RowSelector`; the cache is not copied on clone. For single-pass walks
609    /// over mask-backed selections, prefer streaming directly via
610    /// [`Self::as_mask`] + [`MaskRunIter::new`] — that path is allocation-free
611    /// and avoids populating the cache.
612    pub fn iter(&self) -> impl Iterator<Item = &RowSelector> {
613        match &self.inner {
614            RowSelectionInner::Selectors(s) => s.iter(),
615            RowSelectionInner::Mask(m) => m.selectors().iter(),
616        }
617    }
618
619    /// Returns the number of selected rows
620    pub fn row_count(&self) -> usize {
621        match &self.inner {
622            RowSelectionInner::Selectors(s) => {
623                s.iter().filter(|x| !x.skip).map(|x| x.row_count).sum()
624            }
625            RowSelectionInner::Mask(m) => m.count(),
626        }
627    }
628
629    /// Returns the total number of rows spanned by this selection, both
630    /// selected and skipped
631    pub fn total_row_count(&self) -> usize {
632        match &self.inner {
633            RowSelectionInner::Selectors(s) => s.iter().map(|x| x.row_count).sum(),
634            RowSelectionInner::Mask(m) => m.mask().len(),
635        }
636    }
637
638    /// Returns the number of de-selected rows
639    pub fn skipped_row_count(&self) -> usize {
640        match &self.inner {
641            RowSelectionInner::Selectors(s) => {
642                s.iter().filter(|x| x.skip).map(|x| x.row_count).sum()
643            }
644            RowSelectionInner::Mask(m) => m.mask().len() - m.count(),
645        }
646    }
647
648    /// Expands the selection to align with batch boundaries.
649    /// This is needed when using cached array readers to ensure that
650    /// the cached data covers full batches.
651    pub(crate) fn expand_to_batch_boundaries(&self, batch_size: usize, total_rows: usize) -> Self {
652        if batch_size == 0 {
653            return self.clone();
654        }
655
656        match &self.inner {
657            RowSelectionInner::Selectors(selectors) => expand_to_batch_boundaries_from_selectors(
658                selectors.iter().copied(),
659                batch_size,
660                total_rows,
661            ),
662            RowSelectionInner::Mask(mask) => expand_to_batch_boundaries_from_selectors(
663                MaskRunIter::new(mask.mask()),
664                batch_size,
665                total_rows,
666            ),
667        }
668    }
669}
670
671impl From<Vec<RowSelector>> for RowSelection {
672    fn from(selectors: Vec<RowSelector>) -> Self {
673        selectors.into_iter().collect()
674    }
675}
676
677impl From<BooleanBuffer> for RowSelection {
678    fn from(mask: BooleanBuffer) -> Self {
679        Self::from_boolean_buffer(mask)
680    }
681}
682
683impl FromIterator<RowSelector> for RowSelection {
684    fn from_iter<T: IntoIterator<Item = RowSelector>>(iter: T) -> Self {
685        let iter = iter.into_iter();
686
687        // Capacity before filter
688        let mut selectors = Vec::with_capacity(iter.size_hint().0);
689
690        let mut filtered = iter.filter(|x| x.row_count != 0);
691        if let Some(x) = filtered.next() {
692            selectors.push(x);
693        }
694
695        for s in filtered {
696            if s.row_count == 0 {
697                continue;
698            }
699
700            // Combine consecutive selectors
701            let last = selectors.last_mut().unwrap();
702            if last.skip == s.skip {
703                last.row_count = last.row_count.checked_add(s.row_count).unwrap();
704            } else {
705                selectors.push(s)
706            }
707        }
708
709        Self::from_selectors(selectors)
710    }
711}
712
713impl From<RowSelection> for Vec<RowSelector> {
714    fn from(r: RowSelection) -> Self {
715        r.into_selectors_vec()
716    }
717}
718
719impl From<RowSelection> for VecDeque<RowSelector> {
720    fn from(r: RowSelection) -> Self {
721        r.into_selectors_vec().into()
722    }
723}
724
725impl FromIterator<RowSelection> for RowSelection {
726    /// Concatenate multiple [`RowSelection`]s in iterator order.
727    ///
728    /// When every input is mask-backed the result stays mask-backed
729    /// (`BooleanBuffer`s are appended); otherwise falls back to flattening
730    /// through the per-`RowSelector` form.
731    fn from_iter<T: IntoIterator<Item = RowSelection>>(iter: T) -> Self {
732        let items: Vec<RowSelection> = iter.into_iter().collect();
733
734        let all_mask = items
735            .iter()
736            .all(|s| matches!(&s.inner, RowSelectionInner::Mask(_)));
737
738        if all_mask {
739            let total_len: usize = items
740                .iter()
741                .map(|s| match &s.inner {
742                    RowSelectionInner::Mask(m) => m.mask().len(),
743                    RowSelectionInner::Selectors(_) => unreachable!(),
744                })
745                .sum();
746            let mut builder = BooleanBufferBuilder::new(total_len);
747            for item in items {
748                match item.into_inner() {
749                    RowSelectionInner::Mask(m) => builder.append_buffer(m.mask()),
750                    RowSelectionInner::Selectors(_) => unreachable!(),
751                }
752            }
753            return Self::from_boolean_buffer(builder.finish());
754        }
755
756        items
757            .into_iter()
758            .flat_map(|s| s.into_selectors_vec())
759            .collect()
760    }
761}
762
763#[cfg(test)]
764mod tests {
765    use super::*;
766
767    #[test]
768    fn test_total_row_count() {
769        let selection = RowSelection::from(vec![RowSelector::skip(5), RowSelector::select(3)]);
770        assert_eq!(selection.total_row_count(), 8);
771        assert_eq!(selection.row_count(), 3);
772        assert_eq!(selection.skipped_row_count(), 5);
773
774        let selection =
775            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true]));
776        assert_eq!(selection.total_row_count(), 3);
777        assert_eq!(selection.row_count(), 2);
778        assert_eq!(selection.skipped_row_count(), 1);
779
780        let empty = RowSelection::from(vec![]);
781        assert_eq!(empty.total_row_count(), 0);
782    }
783
784    #[test]
785    fn test_offset_zero_and_zero_batch_expand_are_identity() {
786        let selection =
787            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true]));
788        assert_eq!(selection.clone().offset(0), selection);
789        assert_eq!(selection.expand_to_batch_boundaries(0, 3), selection);
790    }
791
792    #[test]
793    fn test_from_filters() {
794        let filters = vec![
795            BooleanArray::from(vec![false, false, false, true, true, true, true]),
796            BooleanArray::from(vec![true, true, false, false, true, true, true]),
797            BooleanArray::from(vec![false, false, false, false]),
798            BooleanArray::from(Vec::<bool>::new()),
799        ];
800
801        let selection = RowSelection::from_filters(&filters[..1]);
802        assert!(selection.selects_any());
803        assert_eq!(
804            selection.selectors(),
805            vec![RowSelector::skip(3), RowSelector::select(4)]
806        );
807
808        let selection = RowSelection::from_filters(&filters[..2]);
809        assert!(selection.selects_any());
810        assert_eq!(
811            selection.selectors(),
812            vec![
813                RowSelector::skip(3),
814                RowSelector::select(6),
815                RowSelector::skip(2),
816                RowSelector::select(3)
817            ]
818        );
819
820        let selection = RowSelection::from_filters(&filters);
821        assert!(selection.selects_any());
822        assert_eq!(
823            selection.selectors(),
824            vec![
825                RowSelector::skip(3),
826                RowSelector::select(6),
827                RowSelector::skip(2),
828                RowSelector::select(3),
829                RowSelector::skip(4)
830            ]
831        );
832
833        let selection = RowSelection::from_filters(&filters[2..3]);
834        assert!(!selection.selects_any());
835        assert_eq!(selection.selectors(), vec![RowSelector::skip(4)]);
836    }
837
838    #[test]
839    fn test_iter() {
840        // use the iter() API to show it does what is expected and
841        // avoid accidental deletion
842        let selectors = vec![
843            RowSelector::select(3),
844            RowSelector::skip(33),
845            RowSelector::select(4),
846        ];
847
848        let round_tripped: Vec<RowSelector> = RowSelection::from(selectors.clone())
849            .iter()
850            .copied()
851            .collect();
852        assert_eq!(selectors, round_tripped);
853    }
854
855    #[test]
856    fn test_row_count() {
857        let selection = RowSelection::from(vec![
858            RowSelector::skip(34),
859            RowSelector::select(12),
860            RowSelector::skip(3),
861            RowSelector::select(35),
862        ]);
863
864        assert_eq!(selection.row_count(), 12 + 35);
865        assert_eq!(selection.skipped_row_count(), 34 + 3);
866
867        let selection = RowSelection::from(vec![RowSelector::select(12), RowSelector::select(35)]);
868
869        assert_eq!(selection.row_count(), 12 + 35);
870        assert_eq!(selection.skipped_row_count(), 0);
871
872        let selection = RowSelection::from(vec![RowSelector::skip(34), RowSelector::skip(3)]);
873
874        assert_eq!(selection.row_count(), 0);
875        assert_eq!(selection.skipped_row_count(), 34 + 3);
876
877        let selection = RowSelection::from(vec![]);
878
879        assert_eq!(selection.row_count(), 0);
880        assert_eq!(selection.skipped_row_count(), 0);
881    }
882
883    #[test]
884    fn test_mixed_backing_equality_mismatches() {
885        let mask =
886            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true, true]));
887
888        // Total row counts differ
889        let longer = RowSelection::from(vec![
890            RowSelector::select(1),
891            RowSelector::skip(1),
892            RowSelector::select(2),
893            RowSelector::skip(1),
894        ]);
895        assert_ne!(mask, longer);
896        assert_ne!(longer, mask);
897
898        // A selected bit falls inside a skip run
899        let skip_overlap = RowSelection::from(vec![RowSelector::skip(2), RowSelector::select(2)]);
900        assert_ne!(mask, skip_overlap);
901
902        // Select run boundaries do not line up
903        let misaligned = RowSelection::from(vec![
904            RowSelector::select(2),
905            RowSelector::skip(1),
906            RowSelector::select(1),
907        ]);
908        assert_ne!(mask, misaligned);
909
910        let equal = RowSelection::from(vec![
911            RowSelector::select(1),
912            RowSelector::skip(1),
913            RowSelector::select(2),
914        ]);
915        assert_eq!(mask, equal);
916        assert_eq!(equal, mask);
917    }
918
919    #[test]
920    fn test_from_iter_all_mask_preserves_mask_backing() {
921        let a_bits = vec![true, false, true, true];
922        let b_bits = vec![false, true, false];
923        let c_bits = vec![true, true, false, false, true];
924
925        let parts = vec![
926            RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())),
927            RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone())),
928            RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())),
929        ];
930        let collected: RowSelection = parts.into_iter().collect();
931
932        let combined = a_bits
933            .iter()
934            .chain(b_bits.iter())
935            .chain(c_bits.iter())
936            .copied()
937            .collect::<Vec<_>>();
938        let expected = RowSelection::from_filters(&[BooleanArray::from(combined)]);
939
940        assert!(collected.as_mask().is_some());
941        assert_eq!(collected, expected);
942    }
943
944    #[test]
945    fn test_from_iter_mixed_backing_falls_back_to_selectors() {
946        let a_bits = vec![true, false, true];
947        let b_selectors = vec![RowSelector::skip(2), RowSelector::select(3)];
948        let c_bits = vec![false, true];
949
950        let parts = vec![
951            RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone())),
952            RowSelection::from(b_selectors),
953            RowSelection::from_boolean_buffer(BooleanBuffer::from(c_bits.clone())),
954        ];
955        let collected: RowSelection = parts.into_iter().collect();
956
957        assert!(collected.as_mask().is_none());
958
959        let combined_bits = vec![
960            true, false, true, false, false, true, true, true, false, true,
961        ];
962        let expected = RowSelection::from_filters(&[BooleanArray::from(combined_bits)]);
963        assert_eq!(collected, expected);
964    }
965
966    #[test]
967    fn test_from_iter_empty_yields_empty_selection() {
968        let collected: RowSelection = std::iter::empty::<RowSelection>().collect();
969        assert_eq!(collected, RowSelection::default());
970        assert!(collected.as_mask().is_some());
971        assert_eq!(collected.as_mask().unwrap().len(), 0);
972    }
973}
parquet/arrow/arrow_reader/selection/mod.rs

parquet/arrow/arrow_reader/selection/
mod.rs
