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

parquet/arrow/arrow_reader/selection/
mod.rs
