parquet/arrow/arrow_reader/
selection.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
18use crate::errors::ParquetError;
19use crate::file::page_index::offset_index::PageLocation;
20use arrow_array::{Array, BooleanArray};
21use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
22use arrow_select::filter::SlicesIterator;
23use std::cmp::Ordering;
24use std::collections::VecDeque;
25use std::ops::Range;
26use std::sync::Arc;
27
28/// Policy for picking a strategy to materialise [`RowSelection`] during execution.
29#[derive(Clone, Copy, Debug, Eq, PartialEq)]
30pub enum RowSelectionPolicy {
31    /// Use a queue of [`RowSelector`] values
32    Selectors,
33    /// Use a boolean mask to materialise the selection
34    Mask,
35    /// Choose between [`Self::Mask`] and [`Self::Selectors`] based on selector density
36    Auto {
37        /// Average selector length below which masks are preferred
38        threshold: usize,
39    },
40}
41
42impl Default for RowSelectionPolicy {
43    fn default() -> Self {
44        Self::Auto { threshold: 32 }
45    }
46}
47
48/// Fully resolved strategy for materializing [`RowSelection`] during execution.
49///
50/// This is determined by [`RowSelectionPolicy`], including selector density for
51/// [`RowSelectionPolicy::Auto`].
52#[derive(Clone, Copy, Debug, Eq, PartialEq)]
53pub(crate) enum RowSelectionStrategy {
54    /// Use a queue of [`RowSelector`] values
55    Selectors,
56    /// Use a boolean mask to materialise the selection
57    Mask,
58}
59
60/// [`RowSelection`] is a collection of [`RowSelector`] used to skip rows when
61/// scanning a parquet file
62#[derive(Debug, Clone, Copy, Eq, PartialEq)]
63pub struct RowSelector {
64    /// The number of rows
65    pub row_count: usize,
66
67    /// If true, skip `row_count` rows
68    pub skip: bool,
69}
70
71impl RowSelector {
72    /// Select `row_count` rows
73    pub fn select(row_count: usize) -> Self {
74        Self {
75            row_count,
76            skip: false,
77        }
78    }
79
80    /// Skip `row_count` rows
81    pub fn skip(row_count: usize) -> Self {
82        Self {
83            row_count,
84            skip: true,
85        }
86    }
87}
88
89/// [`RowSelection`] allows selecting or skipping a provided number of rows
90/// when scanning the parquet file.
91///
92/// This is applied prior to reading column data, and can therefore
93/// be used to skip IO to fetch data into memory
94///
95/// A typical use-case would be using the [`PageIndex`] to filter out rows
96/// that don't satisfy a predicate
97///
98/// # Example
99/// ```
100/// use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
101///
102/// let selectors = vec![
103///     RowSelector::skip(5),
104///     RowSelector::select(5),
105///     RowSelector::select(5),
106///     RowSelector::skip(5),
107/// ];
108///
109/// // Creating a selection will combine adjacent selectors
110/// let selection: RowSelection = selectors.into();
111///
112/// let expected = vec![
113///     RowSelector::skip(5),
114///     RowSelector::select(10),
115///     RowSelector::skip(5),
116/// ];
117///
118/// let actual: Vec<RowSelector> = selection.into();
119/// assert_eq!(actual, expected);
120///
121/// // you can also create a selection from consecutive ranges
122/// let ranges = vec![5..10, 10..15];
123/// let selection =
124///   RowSelection::from_consecutive_ranges(ranges.into_iter(), 20);
125/// let actual: Vec<RowSelector> = selection.into();
126/// assert_eq!(actual, expected);
127/// ```
128///
129/// A [`RowSelection`] maintains the following invariants:
130///
131/// * It contains no [`RowSelector`] of 0 rows
132/// * Consecutive [`RowSelector`]s alternate skipping or selecting rows
133///
134/// [`PageIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData
135#[derive(Debug, Clone, Default, Eq, PartialEq)]
136pub struct RowSelection {
137    selectors: Vec<RowSelector>,
138}
139
140impl RowSelection {
141    /// Creates a [`RowSelection`] from a slice of [`BooleanArray`]
142    ///
143    /// # Panic
144    ///
145    /// Panics if any of the [`BooleanArray`] contain nulls
146    pub fn from_filters(filters: &[BooleanArray]) -> Self {
147        let mut next_offset = 0;
148        let total_rows = filters.iter().map(|x| x.len()).sum();
149
150        let iter = filters.iter().flat_map(|filter| {
151            let offset = next_offset;
152            next_offset += filter.len();
153            assert_eq!(filter.null_count(), 0);
154            SlicesIterator::new(filter).map(move |(start, end)| start + offset..end + offset)
155        });
156
157        Self::from_consecutive_ranges(iter, total_rows)
158    }
159
160    /// Creates a [`RowSelection`] from an iterator of consecutive ranges to keep
161    pub fn from_consecutive_ranges<I: Iterator<Item = Range<usize>>>(
162        ranges: I,
163        total_rows: usize,
164    ) -> Self {
165        let mut selectors: Vec<RowSelector> = Vec::with_capacity(ranges.size_hint().0);
166        let mut last_end = 0;
167        for range in ranges {
168            let len = range.end - range.start;
169            if len == 0 {
170                continue;
171            }
172
173            match range.start.cmp(&last_end) {
174                Ordering::Equal => match selectors.last_mut() {
175                    Some(last) => last.row_count = last.row_count.checked_add(len).unwrap(),
176                    None => selectors.push(RowSelector::select(len)),
177                },
178                Ordering::Greater => {
179                    selectors.push(RowSelector::skip(range.start - last_end));
180                    selectors.push(RowSelector::select(len))
181                }
182                Ordering::Less => panic!("out of order"),
183            }
184            last_end = range.end;
185        }
186
187        if last_end != total_rows {
188            selectors.push(RowSelector::skip(total_rows - last_end))
189        }
190
191        Self { selectors }
192    }
193
194    /// Given an offset index, return the byte ranges for all data pages selected by `self`
195    ///
196    /// This is useful for determining what byte ranges to fetch from underlying storage
197    ///
198    /// Note: this method does not make any effort to combine consecutive ranges, nor coalesce
199    /// ranges that are close together. This is instead delegated to the IO subsystem to optimise,
200    /// e.g. [`ObjectStore::get_ranges`](object_store::ObjectStore::get_ranges)
201    pub fn scan_ranges(&self, page_locations: &[PageLocation]) -> Vec<Range<u64>> {
202        let mut ranges: Vec<Range<u64>> = vec![];
203        let mut row_offset = 0;
204
205        let mut pages = page_locations.iter().peekable();
206        let mut selectors = self.selectors.iter().cloned();
207        let mut current_selector = selectors.next();
208        let mut current_page = pages.next();
209
210        let mut current_page_included = false;
211
212        while let Some((selector, page)) = current_selector.as_mut().zip(current_page) {
213            if !(selector.skip || current_page_included) {
214                let start = page.offset as u64;
215                let end = start + page.compressed_page_size as u64;
216                ranges.push(start..end);
217                current_page_included = true;
218            }
219
220            if let Some(next_page) = pages.peek() {
221                if row_offset + selector.row_count > next_page.first_row_index as usize {
222                    let remaining_in_page = next_page.first_row_index as usize - row_offset;
223                    selector.row_count -= remaining_in_page;
224                    row_offset += remaining_in_page;
225                    current_page = pages.next();
226                    current_page_included = false;
227
228                    continue;
229                } else {
230                    if row_offset + selector.row_count == next_page.first_row_index as usize {
231                        current_page = pages.next();
232                        current_page_included = false;
233                    }
234                    row_offset += selector.row_count;
235                    current_selector = selectors.next();
236                }
237            } else {
238                if !(selector.skip || current_page_included) {
239                    let start = page.offset as u64;
240                    let end = start + page.compressed_page_size as u64;
241                    ranges.push(start..end);
242                }
243                current_selector = selectors.next()
244            }
245        }
246
247        ranges
248    }
249
250    /// Returns the complete row ranges of the pages selected by [`Self::scan_ranges`].
251    pub(crate) fn row_ranges_for_selected_pages(
252        &self,
253        page_locations: &[PageLocation],
254        total_rows: usize,
255    ) -> Vec<Range<usize>> {
256        let mut selected_pages = self.scan_ranges(page_locations).into_iter().peekable();
257        let mut row_ranges = Vec::new();
258
259        for (idx, page) in page_locations.iter().enumerate() {
260            let Some(selected_page) = selected_pages.peek() else {
261                break;
262            };
263            if selected_page.start != page.offset as u64 {
264                continue;
265            }
266            selected_pages.next();
267
268            let end = page_locations
269                .get(idx + 1)
270                .map(|next| next.first_row_index as usize)
271                .unwrap_or(total_rows);
272            row_ranges.push(page.first_row_index as usize..end);
273        }
274
275        row_ranges
276    }
277
278    /// Splits off the first `row_count` from this [`RowSelection`]
279    pub fn split_off(&mut self, row_count: usize) -> Self {
280        let mut total_count = 0;
281
282        // Find the index where the selector exceeds the row count
283        let find = self.selectors.iter().position(|selector| {
284            total_count += selector.row_count;
285            total_count > row_count
286        });
287
288        let split_idx = match find {
289            Some(idx) => idx,
290            None => {
291                let selectors = std::mem::take(&mut self.selectors);
292                return Self { selectors };
293            }
294        };
295
296        let mut remaining = self.selectors.split_off(split_idx);
297
298        // Always present as `split_idx < self.selectors.len`
299        let next = remaining.first_mut().unwrap();
300        let overflow = total_count - row_count;
301
302        if next.row_count != overflow {
303            self.selectors.push(RowSelector {
304                row_count: next.row_count - overflow,
305                skip: next.skip,
306            })
307        }
308        next.row_count = overflow;
309
310        std::mem::swap(&mut remaining, &mut self.selectors);
311        Self {
312            selectors: remaining,
313        }
314    }
315    /// returns a [`RowSelection`] representing rows that are selected in both
316    /// input [`RowSelection`]s.
317    ///
318    /// This is equivalent to the logical `AND` / conjunction of the two
319    /// selections.
320    ///
321    /// # Example
322    /// If `N` means the row is not selected, and `Y` means it is
323    /// selected:
324    ///
325    /// ```text
326    /// self:     NNNNNNNNNNNNYYYYYYYYYYYYYYYYYYYYYYNNNYYYYY
327    /// other:                YYYYYNNNNYYYYYYYYYYYYY   YYNNN
328    ///
329    /// returned: NNNNNNNNNNNNYYYYYNNNNYYYYYYYYYYYYYNNNYYNNN
330    /// ```
331    ///
332    /// # Panics
333    ///
334    /// Panics if `other` does not have a length equal to the number of rows selected
335    /// by this RowSelection
336    ///
337    pub fn and_then(&self, other: &Self) -> Self {
338        let mut selectors = vec![];
339        let mut first = self.selectors.iter().cloned().peekable();
340        let mut second = other.selectors.iter().cloned().peekable();
341
342        let mut to_skip = 0;
343        while let Some(b) = second.peek_mut() {
344            let a = first
345                .peek_mut()
346                .expect("selection exceeds the number of selected rows");
347
348            if b.row_count == 0 {
349                second.next().unwrap();
350                continue;
351            }
352
353            if a.row_count == 0 {
354                first.next().unwrap();
355                continue;
356            }
357
358            if a.skip {
359                // Records were skipped when producing second
360                to_skip += a.row_count;
361                first.next().unwrap();
362                continue;
363            }
364
365            let skip = b.skip;
366            let to_process = a.row_count.min(b.row_count);
367
368            a.row_count -= to_process;
369            b.row_count -= to_process;
370
371            match skip {
372                true => to_skip += to_process,
373                false => {
374                    if to_skip != 0 {
375                        selectors.push(RowSelector::skip(to_skip));
376                        to_skip = 0;
377                    }
378                    selectors.push(RowSelector::select(to_process))
379                }
380            }
381        }
382
383        for v in first {
384            if v.row_count != 0 {
385                assert!(
386                    v.skip,
387                    "selection contains less than the number of selected rows"
388                );
389                to_skip += v.row_count
390            }
391        }
392
393        if to_skip != 0 {
394            selectors.push(RowSelector::skip(to_skip));
395        }
396
397        Self { selectors }
398    }
399
400    /// Compute the intersection of two [`RowSelection`]
401    /// For example:
402    /// self:      NNYYYYNNYYNYN
403    /// other:     NYNNNNNNY
404    ///
405    /// returned:  NNNNNNNNYYNYN
406    pub fn intersection(&self, other: &Self) -> Self {
407        intersect_row_selections(&self.selectors, &other.selectors)
408    }
409
410    /// Compute the union of two [`RowSelection`]
411    /// For example:
412    /// self:      NNYYYYNNYYNYN
413    /// other:     NYNNNNNNN
414    ///
415    /// returned:  NYYYYYNNYYNYN
416    pub fn union(&self, other: &Self) -> Self {
417        union_row_selections(&self.selectors, &other.selectors)
418    }
419
420    /// Returns `true` if this [`RowSelection`] selects any rows
421    pub fn selects_any(&self) -> bool {
422        self.selectors.iter().any(|x| !x.skip)
423    }
424
425    /// Trims this [`RowSelection`] removing any trailing skips
426    pub(crate) fn trim(mut self) -> Self {
427        while self.selectors.last().map(|x| x.skip).unwrap_or(false) {
428            self.selectors.pop();
429        }
430        self
431    }
432
433    /// Applies an offset to this [`RowSelection`], skipping the first `offset` selected rows
434    pub(crate) fn offset(mut self, offset: usize) -> Self {
435        if offset == 0 {
436            return self;
437        }
438
439        let mut selected_count = 0;
440        let mut skipped_count = 0;
441
442        // Find the index where the selector exceeds the row count
443        let find = self
444            .selectors
445            .iter()
446            .position(|selector| match selector.skip {
447                true => {
448                    skipped_count += selector.row_count;
449                    false
450                }
451                false => {
452                    selected_count += selector.row_count;
453                    selected_count > offset
454                }
455            });
456
457        let split_idx = match find {
458            Some(idx) => idx,
459            None => {
460                self.selectors.clear();
461                return self;
462            }
463        };
464
465        let mut selectors = Vec::with_capacity(self.selectors.len() - split_idx + 1);
466        selectors.push(RowSelector::skip(skipped_count + offset));
467        selectors.push(RowSelector::select(selected_count - offset));
468        selectors.extend_from_slice(&self.selectors[split_idx + 1..]);
469
470        Self { selectors }
471    }
472
473    /// Limit this [`RowSelection`] to only select `limit` rows
474    pub(crate) fn limit(mut self, mut limit: usize) -> Self {
475        if limit == 0 {
476            self.selectors.clear();
477        }
478
479        for (idx, selection) in self.selectors.iter_mut().enumerate() {
480            if !selection.skip {
481                if selection.row_count >= limit {
482                    selection.row_count = limit;
483                    self.selectors.truncate(idx + 1);
484                    break;
485                } else {
486                    limit -= selection.row_count;
487                }
488            }
489        }
490        self
491    }
492
493    /// Returns an iterator over the [`RowSelector`]s for this
494    /// [`RowSelection`].
495    pub fn iter(&self) -> impl Iterator<Item = &RowSelector> {
496        self.selectors.iter()
497    }
498
499    /// Returns the number of selected rows
500    pub fn row_count(&self) -> usize {
501        self.iter().filter(|s| !s.skip).map(|s| s.row_count).sum()
502    }
503
504    /// Returns the number of de-selected rows
505    pub fn skipped_row_count(&self) -> usize {
506        self.iter().filter(|s| s.skip).map(|s| s.row_count).sum()
507    }
508
509    /// Expands the selection to align with batch boundaries.
510    /// This is needed when using cached array readers to ensure that
511    /// the cached data covers full batches.
512    pub(crate) fn expand_to_batch_boundaries(&self, batch_size: usize, total_rows: usize) -> Self {
513        if batch_size == 0 {
514            return self.clone();
515        }
516
517        let mut expanded_ranges = Vec::new();
518        let mut row_offset = 0;
519
520        for selector in &self.selectors {
521            if selector.skip {
522                row_offset += selector.row_count;
523            } else {
524                let start = row_offset;
525                let end = row_offset + selector.row_count;
526
527                // Expand start to batch boundary
528                let expanded_start = (start / batch_size) * batch_size;
529                // Expand end to batch boundary
530                let expanded_end = end.div_ceil(batch_size) * batch_size;
531                let expanded_end = expanded_end.min(total_rows);
532
533                expanded_ranges.push(expanded_start..expanded_end);
534                row_offset += selector.row_count;
535            }
536        }
537
538        // Sort ranges by start position
539        expanded_ranges.sort_by_key(|range| range.start);
540
541        // Merge overlapping or consecutive ranges
542        let mut merged_ranges: Vec<Range<usize>> = Vec::new();
543        for range in expanded_ranges {
544            if let Some(last) = merged_ranges.last_mut() {
545                if range.start <= last.end {
546                    // Overlapping or consecutive - merge them
547                    last.end = last.end.max(range.end);
548                } else {
549                    // No overlap - add new range
550                    merged_ranges.push(range);
551                }
552            } else {
553                // First range
554                merged_ranges.push(range);
555            }
556        }
557
558        Self::from_consecutive_ranges(merged_ranges.into_iter(), total_rows)
559    }
560}
561
562impl From<Vec<RowSelector>> for RowSelection {
563    fn from(selectors: Vec<RowSelector>) -> Self {
564        selectors.into_iter().collect()
565    }
566}
567
568impl FromIterator<RowSelector> for RowSelection {
569    fn from_iter<T: IntoIterator<Item = RowSelector>>(iter: T) -> Self {
570        let iter = iter.into_iter();
571
572        // Capacity before filter
573        let mut selectors = Vec::with_capacity(iter.size_hint().0);
574
575        let mut filtered = iter.filter(|x| x.row_count != 0);
576        if let Some(x) = filtered.next() {
577            selectors.push(x);
578        }
579
580        for s in filtered {
581            if s.row_count == 0 {
582                continue;
583            }
584
585            // Combine consecutive selectors
586            let last = selectors.last_mut().unwrap();
587            if last.skip == s.skip {
588                last.row_count = last.row_count.checked_add(s.row_count).unwrap();
589            } else {
590                selectors.push(s)
591            }
592        }
593
594        Self { selectors }
595    }
596}
597
598impl From<RowSelection> for Vec<RowSelector> {
599    fn from(r: RowSelection) -> Self {
600        r.selectors
601    }
602}
603
604impl From<RowSelection> for VecDeque<RowSelector> {
605    fn from(r: RowSelection) -> Self {
606        r.selectors.into()
607    }
608}
609
610/// Combine two lists of `RowSelection` return the intersection of them
611/// For example:
612/// self:      NNYYYYNNYYNYN
613/// other:     NYNNNNNNY
614///
615/// returned:  NNNNNNNNYYNYN
616fn intersect_row_selections(left: &[RowSelector], right: &[RowSelector]) -> RowSelection {
617    let mut l_iter = left.iter().copied().peekable();
618    let mut r_iter = right.iter().copied().peekable();
619
620    let iter = std::iter::from_fn(move || {
621        loop {
622            let l = l_iter.peek_mut();
623            let r = r_iter.peek_mut();
624
625            match (l, r) {
626                (Some(a), _) if a.row_count == 0 => {
627                    l_iter.next().unwrap();
628                }
629                (_, Some(b)) if b.row_count == 0 => {
630                    r_iter.next().unwrap();
631                }
632                (Some(l), Some(r)) => {
633                    return match (l.skip, r.skip) {
634                        // Keep both ranges
635                        (false, false) => {
636                            if l.row_count < r.row_count {
637                                r.row_count -= l.row_count;
638                                l_iter.next()
639                            } else {
640                                l.row_count -= r.row_count;
641                                r_iter.next()
642                            }
643                        }
644                        // skip at least one
645                        _ => {
646                            if l.row_count < r.row_count {
647                                let skip = l.row_count;
648                                r.row_count -= l.row_count;
649                                l_iter.next();
650                                Some(RowSelector::skip(skip))
651                            } else {
652                                let skip = r.row_count;
653                                l.row_count -= skip;
654                                r_iter.next();
655                                Some(RowSelector::skip(skip))
656                            }
657                        }
658                    };
659                }
660                (Some(_), None) => return l_iter.next(),
661                (None, Some(_)) => return r_iter.next(),
662                (None, None) => return None,
663            }
664        }
665    });
666
667    iter.collect()
668}
669
670/// Combine two lists of `RowSelector` return the union of them
671/// For example:
672/// self:      NNYYYYNNYYNYN
673/// other:     NYNNNNNNY
674///
675/// returned:  NYYYYYNNYYNYN
676///
677/// This can be removed from here once RowSelection::union is in parquet::arrow
678fn union_row_selections(left: &[RowSelector], right: &[RowSelector]) -> RowSelection {
679    let mut l_iter = left.iter().copied().peekable();
680    let mut r_iter = right.iter().copied().peekable();
681
682    let iter = std::iter::from_fn(move || {
683        loop {
684            let l = l_iter.peek_mut();
685            let r = r_iter.peek_mut();
686
687            match (l, r) {
688                (Some(a), _) if a.row_count == 0 => {
689                    l_iter.next().unwrap();
690                }
691                (_, Some(b)) if b.row_count == 0 => {
692                    r_iter.next().unwrap();
693                }
694                (Some(l), Some(r)) => {
695                    return match (l.skip, r.skip) {
696                        // Skip both ranges
697                        (true, true) => {
698                            if l.row_count < r.row_count {
699                                let skip = l.row_count;
700                                r.row_count -= l.row_count;
701                                l_iter.next();
702                                Some(RowSelector::skip(skip))
703                            } else {
704                                let skip = r.row_count;
705                                l.row_count -= skip;
706                                r_iter.next();
707                                Some(RowSelector::skip(skip))
708                            }
709                        }
710                        // Keep rows from left
711                        (false, true) => {
712                            if l.row_count < r.row_count {
713                                r.row_count -= l.row_count;
714                                l_iter.next()
715                            } else {
716                                let r_row_count = r.row_count;
717                                l.row_count -= r_row_count;
718                                r_iter.next();
719                                Some(RowSelector::select(r_row_count))
720                            }
721                        }
722                        // Keep rows from right
723                        (true, false) => {
724                            if l.row_count < r.row_count {
725                                let l_row_count = l.row_count;
726                                r.row_count -= l_row_count;
727                                l_iter.next();
728                                Some(RowSelector::select(l_row_count))
729                            } else {
730                                l.row_count -= r.row_count;
731                                r_iter.next()
732                            }
733                        }
734                        // Keep at least one
735                        _ => {
736                            if l.row_count < r.row_count {
737                                r.row_count -= l.row_count;
738                                l_iter.next()
739                            } else {
740                                l.row_count -= r.row_count;
741                                r_iter.next()
742                            }
743                        }
744                    };
745                }
746                (Some(_), None) => return l_iter.next(),
747                (None, Some(_)) => return r_iter.next(),
748                (None, None) => return None,
749            }
750        }
751    });
752
753    iter.collect()
754}
755
756/// Cursor for iterating a mask-backed [`RowSelection`]
757///
758/// This is best for dense selections where there are many small skips
759/// or selections. For example, selecting every other row.
760///
761/// When page pruning produces sparse column data, `loaded_row_ranges` limits
762/// each decoded chunk to rows whose pages are loaded for every projected leaf.
763/// For example, two projected columns can have different page boundaries:
764///
765/// ```text
766/// Row ranges:       [0, 4) [4, 6) [6, 8) [8, 10) [10, 12)
767/// Selection mask:   1000   00     00     00      01
768/// Column A pages:   loaded | missing [4, 8) | loaded [8, 12)
769/// Column B pages:   loaded [0, 6) | missing [6, 10) | loaded
770/// LoadedRowRanges:  [0, 4)                         [10, 12)
771/// ```
772///
773/// The first chunk decodes `[0, 4)` with mask `1000`. The next chunk skips to
774/// row 11 and decodes `[11, 12)` with mask `1`. The loaded ranges are decode
775/// boundaries, not output batch boundaries: [`ParquetRecordBatchReader`]
776/// accumulates both chunks and applies the combined mask `10001` once.
777///
778/// [`ParquetRecordBatchReader`]: super::ParquetRecordBatchReader
779#[derive(Debug)]
780pub struct MaskCursor {
781    mask: BooleanBuffer,
782    /// Current absolute offset into the selection
783    position: usize,
784    /// Row ranges whose backing pages are loaded for every projected column.
785    loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
786}
787
788impl MaskCursor {
789    /// Returns `true` when no further rows remain
790    pub fn is_empty(&self) -> bool {
791        self.position >= self.mask.len()
792    }
793
794    /// Advance through the mask representation, producing the next chunk summary
795    pub fn next_mask_chunk(&mut self, batch_size: usize) -> Option<MaskChunk> {
796        if self.is_empty() {
797            return None;
798        }
799
800        Some(self.next_mask_chunk_non_empty(batch_size))
801    }
802
803    /// Produces the next chunk for a non-empty, trailing-skip-free mask.
804    fn next_mask_chunk_non_empty(&mut self, batch_size: usize) -> MaskChunk {
805        debug_assert!(!self.is_empty());
806
807        let (initial_skip, chunk_rows, selected_rows, mask_start, end_position) = {
808            let mask = &self.mask;
809            let start_position = self.position;
810            let mut cursor = start_position;
811            let mut initial_skip = 0;
812
813            while cursor < mask.len() && !mask.value(cursor) {
814                initial_skip += 1;
815                cursor += 1;
816            }
817            debug_assert!(
818                cursor < mask.len(),
819                "ReadPlan must remove trailing skips from Mask selections"
820            );
821
822            let mask_start = cursor;
823            let mut chunk_rows = 0;
824            let mut selected_rows = 0;
825
826            // Advance until enough rows have been selected to satisfy the batch size,
827            // or until the mask is exhausted. This mirrors the behaviour of the legacy
828            // `RowSelector` queue-based iteration.
829            while cursor < mask.len() && selected_rows < batch_size {
830                chunk_rows += 1;
831                if mask.value(cursor) {
832                    selected_rows += 1;
833                }
834                cursor += 1;
835            }
836
837            (initial_skip, chunk_rows, selected_rows, mask_start, cursor)
838        };
839
840        self.position = end_position;
841
842        MaskChunk {
843            initial_skip,
844            chunk_rows,
845            selected_rows,
846            mask_start,
847        }
848    }
849
850    /// Returns the next non-empty mask chunk without crossing an unloaded row range.
851    ///
852    /// The [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) removes trailing
853    /// skips before constructing this cursor. Callers therefore only invoke
854    /// this method for a non-empty mask that has another selected row.
855    pub(crate) fn next_chunk(&mut self, batch_size: usize) -> Result<MaskChunk, ParquetError> {
856        debug_assert!(batch_size > 0);
857        debug_assert!(!self.is_empty());
858
859        if self.loaded_row_ranges.is_none() {
860            return Ok(self.next_mask_chunk_non_empty(batch_size));
861        }
862
863        let start_position = self.position;
864        let mut cursor = start_position;
865        while cursor < self.mask.len() && !self.mask.value(cursor) {
866            cursor += 1;
867        }
868
869        debug_assert!(
870            cursor < self.mask.len(),
871            "ReadPlan must remove trailing skips from Mask selections"
872        );
873
874        let loaded_range_end = self
875            .loaded_row_ranges
876            .as_ref()
877            .and_then(|ranges| ranges.end_containing(cursor))
878            .ok_or_else(|| {
879                ParquetError::General(format!(
880                    "Internal Error: selected row {cursor} has no loaded page range"
881                ))
882            })?;
883
884        let mask_start = cursor;
885        let mut selected_rows = 0;
886        while cursor < loaded_range_end && cursor < self.mask.len() && selected_rows < batch_size {
887            if self.mask.value(cursor) {
888                selected_rows += 1;
889            }
890            cursor += 1;
891        }
892
893        self.position = cursor;
894        Ok(MaskChunk {
895            initial_skip: mask_start - start_position,
896            chunk_rows: cursor - mask_start,
897            selected_rows,
898            mask_start,
899        })
900    }
901
902    /// Materialise the boolean values for a mask-backed chunk
903    pub fn mask_values_for(&self, chunk: &MaskChunk) -> Result<BooleanArray, ParquetError> {
904        if chunk.mask_start.saturating_add(chunk.chunk_rows) > self.mask.len() {
905            return Err(ParquetError::General(
906                "Internal Error: MaskChunk exceeds mask length".to_string(),
907            ));
908        }
909        Ok(BooleanArray::from(
910            self.mask.slice(chunk.mask_start, chunk.chunk_rows),
911        ))
912    }
913}
914
915/// Cursor for iterating a selector-backed [`RowSelection`]
916///
917/// This is best for sparse selections where large contiguous
918/// blocks of rows are selected or skipped.
919#[derive(Debug)]
920pub struct SelectorsCursor {
921    selectors: VecDeque<RowSelector>,
922    /// Current absolute offset into the selection
923    position: usize,
924}
925
926impl SelectorsCursor {
927    /// Returns `true` when no further rows remain
928    pub fn is_empty(&self) -> bool {
929        self.selectors.is_empty()
930    }
931
932    pub(crate) fn selectors_mut(&mut self) -> &mut VecDeque<RowSelector> {
933        &mut self.selectors
934    }
935
936    /// Return the next [`RowSelector`]
937    pub(crate) fn next_selector(&mut self) -> RowSelector {
938        let selector = self.selectors.pop_front().unwrap();
939        self.position += selector.row_count;
940        selector
941    }
942
943    /// Return a selector to the front, rewinding the position
944    pub(crate) fn return_selector(&mut self, selector: RowSelector) {
945        self.position = self.position.saturating_sub(selector.row_count);
946        self.selectors.push_front(selector);
947    }
948}
949
950/// Result of computing the next chunk to read when using a [`MaskCursor`]
951#[derive(Debug)]
952pub struct MaskChunk {
953    /// Number of leading rows to skip before reaching selected rows
954    pub initial_skip: usize,
955    /// Total rows covered by this chunk (selected + skipped)
956    pub chunk_rows: usize,
957    /// Rows actually selected within the chunk
958    pub selected_rows: usize,
959    /// Starting offset within the mask where the chunk begins
960    pub mask_start: usize,
961}
962
963/// Row ranges whose backing pages are loaded for every projected column.
964#[derive(Clone, Debug)]
965pub(crate) struct LoadedRowRanges(Vec<Range<usize>>);
966
967impl LoadedRowRanges {
968    pub(crate) fn from_selection(selection: RowSelection) -> Self {
969        let selectors: Vec<RowSelector> = selection.into();
970        let mut position = 0;
971        let ranges = selectors
972            .into_iter()
973            .filter_map(|selector| {
974                let start = position;
975                position += selector.row_count;
976                (!selector.skip).then_some(start..position)
977            })
978            .collect();
979        Self(ranges)
980    }
981
982    fn end_containing(&self, row: usize) -> Option<usize> {
983        let idx = self.0.partition_point(|range| range.end <= row);
984        self.0
985            .get(idx)
986            .filter(|range| range.start <= row)
987            .map(|range| range.end)
988    }
989
990    #[cfg(test)]
991    pub(crate) fn ranges(&self) -> &[Range<usize>] {
992        &self.0
993    }
994}
995
996/// Cursor for iterating a [`RowSelection`] during execution within a
997/// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan).
998///
999/// This keeps per-reader state such as the current position and delegates the
1000/// actual storage strategy to the internal `RowSelectionBacking`.
1001#[derive(Debug)]
1002pub enum RowSelectionCursor {
1003    /// Reading all rows
1004    All,
1005    /// Use a bitmask to back the selection (dense selections)
1006    Mask(MaskCursor),
1007    /// Use a queue of selectors to back the selection (sparse selections)
1008    Selectors(SelectorsCursor),
1009}
1010
1011impl RowSelectionCursor {
1012    /// Create a [`MaskCursor`] cursor backed by a bitmask, from an existing set of selectors
1013    pub(crate) fn new_mask_from_selectors(
1014        selectors: Vec<RowSelector>,
1015        loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
1016    ) -> Self {
1017        debug_assert!(
1018            selectors
1019                .last()
1020                .map(|selector| !selector.skip)
1021                .unwrap_or(true),
1022            "Mask selectors must not end with a skip"
1023        );
1024        Self::Mask(MaskCursor {
1025            mask: boolean_mask_from_selectors(&selectors),
1026            position: 0,
1027            loaded_row_ranges,
1028        })
1029    }
1030
1031    /// Create a [`RowSelectionCursor::Selectors`] from the provided selectors
1032    pub(crate) fn new_selectors(selectors: Vec<RowSelector>) -> Self {
1033        Self::Selectors(SelectorsCursor {
1034            selectors: selectors.into(),
1035            position: 0,
1036        })
1037    }
1038
1039    /// Create a cursor that selects all rows
1040    pub(crate) fn new_all() -> Self {
1041        Self::All
1042    }
1043}
1044
1045fn boolean_mask_from_selectors(selectors: &[RowSelector]) -> BooleanBuffer {
1046    let total_rows: usize = selectors.iter().map(|s| s.row_count).sum();
1047    let mut builder = BooleanBufferBuilder::new(total_rows);
1048    for selector in selectors {
1049        builder.append_n(selector.row_count, !selector.skip);
1050    }
1051    builder.finish()
1052}
1053
1054#[cfg(test)]
1055mod tests {
1056    use super::*;
1057    use rand::{Rng, rng};
1058
1059    #[test]
1060    fn test_from_filters() {
1061        let filters = vec![
1062            BooleanArray::from(vec![false, false, false, true, true, true, true]),
1063            BooleanArray::from(vec![true, true, false, false, true, true, true]),
1064            BooleanArray::from(vec![false, false, false, false]),
1065            BooleanArray::from(Vec::<bool>::new()),
1066        ];
1067
1068        let selection = RowSelection::from_filters(&filters[..1]);
1069        assert!(selection.selects_any());
1070        assert_eq!(
1071            selection.selectors,
1072            vec![RowSelector::skip(3), RowSelector::select(4)]
1073        );
1074
1075        let selection = RowSelection::from_filters(&filters[..2]);
1076        assert!(selection.selects_any());
1077        assert_eq!(
1078            selection.selectors,
1079            vec![
1080                RowSelector::skip(3),
1081                RowSelector::select(6),
1082                RowSelector::skip(2),
1083                RowSelector::select(3)
1084            ]
1085        );
1086
1087        let selection = RowSelection::from_filters(&filters);
1088        assert!(selection.selects_any());
1089        assert_eq!(
1090            selection.selectors,
1091            vec![
1092                RowSelector::skip(3),
1093                RowSelector::select(6),
1094                RowSelector::skip(2),
1095                RowSelector::select(3),
1096                RowSelector::skip(4)
1097            ]
1098        );
1099
1100        let selection = RowSelection::from_filters(&filters[2..3]);
1101        assert!(!selection.selects_any());
1102        assert_eq!(selection.selectors, vec![RowSelector::skip(4)]);
1103    }
1104
1105    #[test]
1106    fn test_split_off() {
1107        let mut selection = RowSelection::from(vec![
1108            RowSelector::skip(34),
1109            RowSelector::select(12),
1110            RowSelector::skip(3),
1111            RowSelector::select(35),
1112        ]);
1113
1114        let split = selection.split_off(34);
1115        assert_eq!(split.selectors, vec![RowSelector::skip(34)]);
1116        assert_eq!(
1117            selection.selectors,
1118            vec![
1119                RowSelector::select(12),
1120                RowSelector::skip(3),
1121                RowSelector::select(35)
1122            ]
1123        );
1124
1125        let split = selection.split_off(5);
1126        assert_eq!(split.selectors, vec![RowSelector::select(5)]);
1127        assert_eq!(
1128            selection.selectors,
1129            vec![
1130                RowSelector::select(7),
1131                RowSelector::skip(3),
1132                RowSelector::select(35)
1133            ]
1134        );
1135
1136        let split = selection.split_off(8);
1137        assert_eq!(
1138            split.selectors,
1139            vec![RowSelector::select(7), RowSelector::skip(1)]
1140        );
1141        assert_eq!(
1142            selection.selectors,
1143            vec![RowSelector::skip(2), RowSelector::select(35)]
1144        );
1145
1146        let split = selection.split_off(200);
1147        assert_eq!(
1148            split.selectors,
1149            vec![RowSelector::skip(2), RowSelector::select(35)]
1150        );
1151        assert!(selection.selectors.is_empty());
1152    }
1153
1154    #[test]
1155    fn test_offset() {
1156        let selection = RowSelection::from(vec![
1157            RowSelector::select(5),
1158            RowSelector::skip(23),
1159            RowSelector::select(7),
1160            RowSelector::skip(33),
1161            RowSelector::select(6),
1162        ]);
1163
1164        let selection = selection.offset(2);
1165        assert_eq!(
1166            selection.selectors,
1167            vec![
1168                RowSelector::skip(2),
1169                RowSelector::select(3),
1170                RowSelector::skip(23),
1171                RowSelector::select(7),
1172                RowSelector::skip(33),
1173                RowSelector::select(6),
1174            ]
1175        );
1176
1177        let selection = selection.offset(5);
1178        assert_eq!(
1179            selection.selectors,
1180            vec![
1181                RowSelector::skip(30),
1182                RowSelector::select(5),
1183                RowSelector::skip(33),
1184                RowSelector::select(6),
1185            ]
1186        );
1187
1188        let selection = selection.offset(3);
1189        assert_eq!(
1190            selection.selectors,
1191            vec![
1192                RowSelector::skip(33),
1193                RowSelector::select(2),
1194                RowSelector::skip(33),
1195                RowSelector::select(6),
1196            ]
1197        );
1198
1199        let selection = selection.offset(2);
1200        assert_eq!(
1201            selection.selectors,
1202            vec![RowSelector::skip(68), RowSelector::select(6),]
1203        );
1204
1205        let selection = selection.offset(3);
1206        assert_eq!(
1207            selection.selectors,
1208            vec![RowSelector::skip(71), RowSelector::select(3),]
1209        );
1210    }
1211
1212    #[test]
1213    fn test_and() {
1214        let mut a = RowSelection::from(vec![
1215            RowSelector::skip(12),
1216            RowSelector::select(23),
1217            RowSelector::skip(3),
1218            RowSelector::select(5),
1219        ]);
1220
1221        let b = RowSelection::from(vec![
1222            RowSelector::select(5),
1223            RowSelector::skip(4),
1224            RowSelector::select(15),
1225            RowSelector::skip(4),
1226        ]);
1227
1228        let mut expected = RowSelection::from(vec![
1229            RowSelector::skip(12),
1230            RowSelector::select(5),
1231            RowSelector::skip(4),
1232            RowSelector::select(14),
1233            RowSelector::skip(3),
1234            RowSelector::select(1),
1235            RowSelector::skip(4),
1236        ]);
1237
1238        assert_eq!(a.and_then(&b), expected);
1239
1240        a.split_off(7);
1241        expected.split_off(7);
1242        assert_eq!(a.and_then(&b), expected);
1243
1244        let a = RowSelection::from(vec![RowSelector::select(5), RowSelector::skip(3)]);
1245
1246        let b = RowSelection::from(vec![
1247            RowSelector::select(2),
1248            RowSelector::skip(1),
1249            RowSelector::select(1),
1250            RowSelector::skip(1),
1251        ]);
1252
1253        assert_eq!(
1254            a.and_then(&b).selectors,
1255            vec![
1256                RowSelector::select(2),
1257                RowSelector::skip(1),
1258                RowSelector::select(1),
1259                RowSelector::skip(4)
1260            ]
1261        );
1262    }
1263
1264    #[test]
1265    fn test_combine() {
1266        let a = vec![
1267            RowSelector::skip(3),
1268            RowSelector::skip(3),
1269            RowSelector::select(10),
1270            RowSelector::skip(4),
1271        ];
1272
1273        let b = vec![
1274            RowSelector::skip(3),
1275            RowSelector::skip(3),
1276            RowSelector::select(10),
1277            RowSelector::skip(4),
1278            RowSelector::skip(0),
1279        ];
1280
1281        let c = vec![
1282            RowSelector::skip(2),
1283            RowSelector::skip(4),
1284            RowSelector::select(3),
1285            RowSelector::select(3),
1286            RowSelector::select(4),
1287            RowSelector::skip(3),
1288            RowSelector::skip(1),
1289            RowSelector::skip(0),
1290        ];
1291
1292        let expected = RowSelection::from(vec![
1293            RowSelector::skip(6),
1294            RowSelector::select(10),
1295            RowSelector::skip(4),
1296        ]);
1297
1298        assert_eq!(RowSelection::from_iter(a), expected);
1299        assert_eq!(RowSelection::from_iter(b), expected);
1300        assert_eq!(RowSelection::from_iter(c), expected);
1301    }
1302
1303    #[test]
1304    fn test_combine_2elements() {
1305        let a = vec![RowSelector::select(10), RowSelector::select(5)];
1306        let a_expect = vec![RowSelector::select(15)];
1307        assert_eq!(RowSelection::from_iter(a).selectors, a_expect);
1308
1309        let b = vec![RowSelector::select(10), RowSelector::skip(5)];
1310        let b_expect = vec![RowSelector::select(10), RowSelector::skip(5)];
1311        assert_eq!(RowSelection::from_iter(b).selectors, b_expect);
1312
1313        let c = vec![RowSelector::skip(10), RowSelector::select(5)];
1314        let c_expect = vec![RowSelector::skip(10), RowSelector::select(5)];
1315        assert_eq!(RowSelection::from_iter(c).selectors, c_expect);
1316
1317        let d = vec![RowSelector::skip(10), RowSelector::skip(5)];
1318        let d_expect = vec![RowSelector::skip(15)];
1319        assert_eq!(RowSelection::from_iter(d).selectors, d_expect);
1320    }
1321
1322    #[test]
1323    fn test_from_one_and_empty() {
1324        let a = vec![RowSelector::select(10)];
1325        let selection1 = RowSelection::from(a.clone());
1326        assert_eq!(selection1.selectors, a);
1327
1328        let b = vec![];
1329        let selection1 = RowSelection::from(b.clone());
1330        assert_eq!(selection1.selectors, b)
1331    }
1332
1333    #[test]
1334    #[should_panic(expected = "selection exceeds the number of selected rows")]
1335    fn test_and_longer() {
1336        let a = RowSelection::from(vec![
1337            RowSelector::select(3),
1338            RowSelector::skip(33),
1339            RowSelector::select(3),
1340            RowSelector::skip(33),
1341        ]);
1342        let b = RowSelection::from(vec![RowSelector::select(36)]);
1343        a.and_then(&b);
1344    }
1345
1346    #[test]
1347    #[should_panic(expected = "selection contains less than the number of selected rows")]
1348    fn test_and_shorter() {
1349        let a = RowSelection::from(vec![
1350            RowSelector::select(3),
1351            RowSelector::skip(33),
1352            RowSelector::select(3),
1353            RowSelector::skip(33),
1354        ]);
1355        let b = RowSelection::from(vec![RowSelector::select(3)]);
1356        a.and_then(&b);
1357    }
1358
1359    #[test]
1360    fn test_intersect_row_selection_and_combine() {
1361        // a size equal b size
1362        let a = vec![
1363            RowSelector::select(5),
1364            RowSelector::skip(4),
1365            RowSelector::select(1),
1366        ];
1367        let b = vec![
1368            RowSelector::select(8),
1369            RowSelector::skip(1),
1370            RowSelector::select(1),
1371        ];
1372
1373        let res = intersect_row_selections(&a, &b);
1374        assert_eq!(
1375            res.selectors,
1376            vec![
1377                RowSelector::select(5),
1378                RowSelector::skip(4),
1379                RowSelector::select(1),
1380            ],
1381        );
1382
1383        // a size larger than b size
1384        let a = vec![
1385            RowSelector::select(3),
1386            RowSelector::skip(33),
1387            RowSelector::select(3),
1388            RowSelector::skip(33),
1389        ];
1390        let b = vec![RowSelector::select(36), RowSelector::skip(36)];
1391        let res = intersect_row_selections(&a, &b);
1392        assert_eq!(
1393            res.selectors,
1394            vec![RowSelector::select(3), RowSelector::skip(69)]
1395        );
1396
1397        // a size less than b size
1398        let a = vec![RowSelector::select(3), RowSelector::skip(7)];
1399        let b = vec![
1400            RowSelector::select(2),
1401            RowSelector::skip(2),
1402            RowSelector::select(2),
1403            RowSelector::skip(2),
1404            RowSelector::select(2),
1405        ];
1406        let res = intersect_row_selections(&a, &b);
1407        assert_eq!(
1408            res.selectors,
1409            vec![RowSelector::select(2), RowSelector::skip(8)]
1410        );
1411
1412        let a = vec![RowSelector::select(3), RowSelector::skip(7)];
1413        let b = vec![
1414            RowSelector::select(2),
1415            RowSelector::skip(2),
1416            RowSelector::select(2),
1417            RowSelector::skip(2),
1418            RowSelector::select(2),
1419        ];
1420        let res = intersect_row_selections(&a, &b);
1421        assert_eq!(
1422            res.selectors,
1423            vec![RowSelector::select(2), RowSelector::skip(8)]
1424        );
1425    }
1426
1427    #[test]
1428    fn test_and_fuzz() {
1429        let mut rand = rng();
1430        for _ in 0..100 {
1431            let a_len = rand.random_range(10..100);
1432            let a_bools: Vec<_> = (0..a_len).map(|_| rand.random_bool(0.2)).collect();
1433            let a = RowSelection::from_filters(&[BooleanArray::from(a_bools.clone())]);
1434
1435            let b_len: usize = a_bools.iter().map(|x| *x as usize).sum();
1436            let b_bools: Vec<_> = (0..b_len).map(|_| rand.random_bool(0.8)).collect();
1437            let b = RowSelection::from_filters(&[BooleanArray::from(b_bools.clone())]);
1438
1439            let mut expected_bools = vec![false; a_len];
1440
1441            let mut iter_b = b_bools.iter();
1442            for (idx, b) in a_bools.iter().enumerate() {
1443                if *b && *iter_b.next().unwrap() {
1444                    expected_bools[idx] = true;
1445                }
1446            }
1447
1448            let expected = RowSelection::from_filters(&[BooleanArray::from(expected_bools)]);
1449
1450            let total_rows: usize = expected.selectors.iter().map(|s| s.row_count).sum();
1451            assert_eq!(a_len, total_rows);
1452
1453            assert_eq!(a.and_then(&b), expected);
1454        }
1455    }
1456
1457    #[test]
1458    fn test_iter() {
1459        // use the iter() API to show it does what is expected and
1460        // avoid accidental deletion
1461        let selectors = vec![
1462            RowSelector::select(3),
1463            RowSelector::skip(33),
1464            RowSelector::select(4),
1465        ];
1466
1467        let round_tripped = RowSelection::from(selectors.clone())
1468            .iter()
1469            .cloned()
1470            .collect::<Vec<_>>();
1471        assert_eq!(selectors, round_tripped);
1472    }
1473
1474    #[test]
1475    fn test_limit() {
1476        // Limit to existing limit should no-op
1477        let selection = RowSelection::from(vec![RowSelector::select(10), RowSelector::skip(90)]);
1478        let limited = selection.limit(10);
1479        assert_eq!(RowSelection::from(vec![RowSelector::select(10)]), limited);
1480
1481        let selection = RowSelection::from(vec![
1482            RowSelector::select(10),
1483            RowSelector::skip(10),
1484            RowSelector::select(10),
1485            RowSelector::skip(10),
1486            RowSelector::select(10),
1487        ]);
1488
1489        let limited = selection.clone().limit(5);
1490        let expected = vec![RowSelector::select(5)];
1491        assert_eq!(limited.selectors, expected);
1492
1493        let limited = selection.clone().limit(15);
1494        let expected = vec![
1495            RowSelector::select(10),
1496            RowSelector::skip(10),
1497            RowSelector::select(5),
1498        ];
1499        assert_eq!(limited.selectors, expected);
1500
1501        let limited = selection.clone().limit(0);
1502        let expected = vec![];
1503        assert_eq!(limited.selectors, expected);
1504
1505        let limited = selection.clone().limit(30);
1506        let expected = vec![
1507            RowSelector::select(10),
1508            RowSelector::skip(10),
1509            RowSelector::select(10),
1510            RowSelector::skip(10),
1511            RowSelector::select(10),
1512        ];
1513        assert_eq!(limited.selectors, expected);
1514
1515        let limited = selection.limit(100);
1516        let expected = vec![
1517            RowSelector::select(10),
1518            RowSelector::skip(10),
1519            RowSelector::select(10),
1520            RowSelector::skip(10),
1521            RowSelector::select(10),
1522        ];
1523        assert_eq!(limited.selectors, expected);
1524    }
1525
1526    #[test]
1527    fn test_scan_ranges() {
1528        let index = vec![
1529            PageLocation {
1530                offset: 0,
1531                compressed_page_size: 10,
1532                first_row_index: 0,
1533            },
1534            PageLocation {
1535                offset: 10,
1536                compressed_page_size: 10,
1537                first_row_index: 10,
1538            },
1539            PageLocation {
1540                offset: 20,
1541                compressed_page_size: 10,
1542                first_row_index: 20,
1543            },
1544            PageLocation {
1545                offset: 30,
1546                compressed_page_size: 10,
1547                first_row_index: 30,
1548            },
1549            PageLocation {
1550                offset: 40,
1551                compressed_page_size: 10,
1552                first_row_index: 40,
1553            },
1554            PageLocation {
1555                offset: 50,
1556                compressed_page_size: 10,
1557                first_row_index: 50,
1558            },
1559            PageLocation {
1560                offset: 60,
1561                compressed_page_size: 10,
1562                first_row_index: 60,
1563            },
1564        ];
1565
1566        let selection = RowSelection::from(vec![
1567            // Skip first page
1568            RowSelector::skip(10),
1569            // Multiple selects in same page
1570            RowSelector::select(3),
1571            RowSelector::skip(3),
1572            RowSelector::select(4),
1573            // Select to page boundary
1574            RowSelector::skip(5),
1575            RowSelector::select(5),
1576            // Skip full page past page boundary
1577            RowSelector::skip(12),
1578            // Select across page boundaries
1579            RowSelector::select(12),
1580            // Skip final page
1581            RowSelector::skip(12),
1582        ]);
1583
1584        let ranges = selection.scan_ranges(&index);
1585
1586        // assert_eq!(mask, vec![false, true, true, false, true, true, false]);
1587        assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60]);
1588        assert_eq!(
1589            selection.row_ranges_for_selected_pages(&index, 70),
1590            vec![10..20, 20..30, 40..50, 50..60]
1591        );
1592
1593        let selection = RowSelection::from(vec![
1594            // Skip first page
1595            RowSelector::skip(10),
1596            // Multiple selects in same page
1597            RowSelector::select(3),
1598            RowSelector::skip(3),
1599            RowSelector::select(4),
1600            // Select to page boundary
1601            RowSelector::skip(5),
1602            RowSelector::select(5),
1603            // Skip full page past page boundary
1604            RowSelector::skip(12),
1605            // Select across page boundaries
1606            RowSelector::select(12),
1607            RowSelector::skip(1),
1608            // Select across page boundaries including final page
1609            RowSelector::select(8),
1610        ]);
1611
1612        let ranges = selection.scan_ranges(&index);
1613
1614        // assert_eq!(mask, vec![false, true, true, false, true, true, true]);
1615        assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]);
1616
1617        let selection = RowSelection::from(vec![
1618            // Skip first page
1619            RowSelector::skip(10),
1620            // Multiple selects in same page
1621            RowSelector::select(3),
1622            RowSelector::skip(3),
1623            RowSelector::select(4),
1624            // Select to page boundary
1625            RowSelector::skip(5),
1626            RowSelector::select(5),
1627            // Skip full page past page boundary
1628            RowSelector::skip(12),
1629            // Select to final page boundary
1630            RowSelector::select(12),
1631            RowSelector::skip(1),
1632            // Skip across final page boundary
1633            RowSelector::skip(8),
1634            // Select from final page
1635            RowSelector::select(4),
1636        ]);
1637
1638        let ranges = selection.scan_ranges(&index);
1639
1640        // assert_eq!(mask, vec![false, true, true, false, true, true, true]);
1641        assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]);
1642
1643        let selection = RowSelection::from(vec![
1644            // Skip first page
1645            RowSelector::skip(10),
1646            // Multiple selects in same page
1647            RowSelector::select(3),
1648            RowSelector::skip(3),
1649            RowSelector::select(4),
1650            // Select to remaining in page and first row of next page
1651            RowSelector::skip(5),
1652            RowSelector::select(6),
1653            // Skip remaining
1654            RowSelector::skip(50),
1655        ]);
1656
1657        let ranges = selection.scan_ranges(&index);
1658
1659        // assert_eq!(mask, vec![false, true, true, false, true, true, true]);
1660        assert_eq!(ranges, vec![10..20, 20..30, 30..40]);
1661    }
1662
1663    #[test]
1664    fn test_loaded_mask_chunk_stops_at_trimmed_mask_end() {
1665        let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges(
1666            std::iter::once(0..5),
1667            10,
1668        ));
1669        let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors(
1670            vec![RowSelector::select(1)],
1671            Some(loaded.into()),
1672        ) else {
1673            unreachable!()
1674        };
1675
1676        let chunk = cursor.next_chunk(10).unwrap();
1677        assert_eq!(chunk.chunk_rows, 1);
1678        assert!(cursor.is_empty());
1679    }
1680
1681    #[test]
1682    fn test_next_mask_chunk_until_cursor_is_empty() {
1683        let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors(
1684            vec![
1685                RowSelector::skip(2),
1686                RowSelector::select(2),
1687                RowSelector::skip(1),
1688                RowSelector::select(1),
1689            ],
1690            None,
1691        ) else {
1692            unreachable!()
1693        };
1694
1695        let first = cursor.next_mask_chunk(2).unwrap();
1696        assert_eq!(first.initial_skip, 2);
1697        assert_eq!(first.chunk_rows, 2);
1698        assert_eq!(first.selected_rows, 2);
1699
1700        let second = cursor.next_mask_chunk(2).unwrap();
1701        assert_eq!(second.initial_skip, 1);
1702        assert_eq!(second.chunk_rows, 1);
1703        assert_eq!(second.selected_rows, 1);
1704
1705        assert!(cursor.next_mask_chunk(2).is_none());
1706    }
1707
1708    #[test]
1709    fn test_from_ranges() {
1710        let ranges = [1..3, 4..6, 6..6, 8..8, 9..10];
1711        let selection = RowSelection::from_consecutive_ranges(ranges.into_iter(), 10);
1712        assert_eq!(
1713            selection.selectors,
1714            vec![
1715                RowSelector::skip(1),
1716                RowSelector::select(2),
1717                RowSelector::skip(1),
1718                RowSelector::select(2),
1719                RowSelector::skip(3),
1720                RowSelector::select(1)
1721            ]
1722        );
1723
1724        let out_of_order_ranges = [1..3, 8..10, 4..7];
1725        let result = std::panic::catch_unwind(|| {
1726            RowSelection::from_consecutive_ranges(out_of_order_ranges.into_iter(), 10)
1727        });
1728        assert!(result.is_err());
1729    }
1730
1731    #[test]
1732    fn test_empty_selector() {
1733        let selection = RowSelection::from(vec![
1734            RowSelector::skip(0),
1735            RowSelector::select(2),
1736            RowSelector::skip(0),
1737            RowSelector::select(2),
1738        ]);
1739        assert_eq!(selection.selectors, vec![RowSelector::select(4)]);
1740
1741        let selection = RowSelection::from(vec![
1742            RowSelector::select(0),
1743            RowSelector::skip(2),
1744            RowSelector::select(0),
1745            RowSelector::skip(2),
1746        ]);
1747        assert_eq!(selection.selectors, vec![RowSelector::skip(4)]);
1748    }
1749
1750    #[test]
1751    fn test_intersection() {
1752        let selection = RowSelection::from(vec![RowSelector::select(1048576)]);
1753        let result = selection.intersection(&selection);
1754        assert_eq!(result, selection);
1755
1756        let a = RowSelection::from(vec![
1757            RowSelector::skip(10),
1758            RowSelector::select(10),
1759            RowSelector::skip(10),
1760            RowSelector::select(20),
1761        ]);
1762
1763        let b = RowSelection::from(vec![
1764            RowSelector::skip(20),
1765            RowSelector::select(20),
1766            RowSelector::skip(10),
1767        ]);
1768
1769        let result = a.intersection(&b);
1770        assert_eq!(
1771            result.selectors,
1772            vec![
1773                RowSelector::skip(30),
1774                RowSelector::select(10),
1775                RowSelector::skip(10)
1776            ]
1777        );
1778    }
1779
1780    #[test]
1781    fn test_union() {
1782        let selection = RowSelection::from(vec![RowSelector::select(1048576)]);
1783        let result = selection.union(&selection);
1784        assert_eq!(result, selection);
1785
1786        // NYNYY
1787        let a = RowSelection::from(vec![
1788            RowSelector::skip(10),
1789            RowSelector::select(10),
1790            RowSelector::skip(10),
1791            RowSelector::select(20),
1792        ]);
1793
1794        // NNYYNYN
1795        let b = RowSelection::from(vec![
1796            RowSelector::skip(20),
1797            RowSelector::select(20),
1798            RowSelector::skip(10),
1799            RowSelector::select(10),
1800            RowSelector::skip(10),
1801        ]);
1802
1803        let result = a.union(&b);
1804
1805        // NYYYYYN
1806        assert_eq!(
1807            result.iter().collect::<Vec<_>>(),
1808            vec![
1809                &RowSelector::skip(10),
1810                &RowSelector::select(50),
1811                &RowSelector::skip(10),
1812            ]
1813        );
1814    }
1815
1816    #[test]
1817    fn test_row_count() {
1818        let selection = RowSelection::from(vec![
1819            RowSelector::skip(34),
1820            RowSelector::select(12),
1821            RowSelector::skip(3),
1822            RowSelector::select(35),
1823        ]);
1824
1825        assert_eq!(selection.row_count(), 12 + 35);
1826        assert_eq!(selection.skipped_row_count(), 34 + 3);
1827
1828        let selection = RowSelection::from(vec![RowSelector::select(12), RowSelector::select(35)]);
1829
1830        assert_eq!(selection.row_count(), 12 + 35);
1831        assert_eq!(selection.skipped_row_count(), 0);
1832
1833        let selection = RowSelection::from(vec![RowSelector::skip(34), RowSelector::skip(3)]);
1834
1835        assert_eq!(selection.row_count(), 0);
1836        assert_eq!(selection.skipped_row_count(), 34 + 3);
1837
1838        let selection = RowSelection::from(vec![]);
1839
1840        assert_eq!(selection.row_count(), 0);
1841        assert_eq!(selection.skipped_row_count(), 0);
1842    }
1843
1844    #[test]
1845    fn test_trim() {
1846        let selection = RowSelection::from(vec![
1847            RowSelector::skip(34),
1848            RowSelector::select(12),
1849            RowSelector::skip(3),
1850            RowSelector::select(35),
1851        ]);
1852
1853        let expected = vec![
1854            RowSelector::skip(34),
1855            RowSelector::select(12),
1856            RowSelector::skip(3),
1857            RowSelector::select(35),
1858        ];
1859
1860        assert_eq!(selection.trim().selectors, expected);
1861
1862        let selection = RowSelection::from(vec![
1863            RowSelector::skip(34),
1864            RowSelector::select(12),
1865            RowSelector::skip(3),
1866        ]);
1867
1868        let expected = vec![RowSelector::skip(34), RowSelector::select(12)];
1869
1870        assert_eq!(selection.trim().selectors, expected);
1871    }
1872}
parquet/arrow/arrow_reader/selection.rs

parquet/arrow/arrow_reader/
selection.rs
