parquet/arrow/arrow_reader/selection/
cursor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Execution time iteration over a [`RowSelection`].
6//!
7//! A [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) resolves a
8//! [`RowSelectionPolicy`] into a [`RowSelectionStrategy`] and builds the
9//! matching [`RowSelectionCursor`], which keeps the per-reader position while
10//! the selection itself stays immutable.
11
12use super::{RowSelection, RowSelector};
13use crate::errors::ParquetError;
14use arrow_array::BooleanArray;
15use arrow_buffer::BooleanBuffer;
16use arrow_buffer::bit_chunk_iterator::BitChunks;
17use std::collections::VecDeque;
18use std::ops::Range;
19use std::sync::Arc;
20
21/// Policy for picking a strategy to materialize [`RowSelection`] during execution.
22#[derive(Clone, Copy, Debug, Eq, PartialEq)]
23pub enum RowSelectionPolicy {
24    /// Use a queue of [`RowSelector`] values
25    Selectors,
26    /// Use a boolean mask to materialize the selection
27    Mask,
28    /// Choose between [`Self::Mask`] and [`Self::Selectors`] based on selector density
29    Auto {
30        /// Average selector length below which masks are preferred
31        threshold: usize,
32    },
33}
34
35impl Default for RowSelectionPolicy {
36    fn default() -> Self {
37        Self::Auto { threshold: 32 }
38    }
39}
40
41impl RowSelectionPolicy {
42    /// Resolve this policy for a selection without changing its backing.
43    pub(crate) fn resolve(self, selection: Option<&RowSelection>) -> RowSelectionStrategy {
44        match self {
45            Self::Selectors => RowSelectionStrategy::Selectors,
46            Self::Mask => RowSelectionStrategy::Mask,
47            Self::Auto { threshold } => selection
48                .map(|selection| selection.auto_selection_strategy(threshold))
49                .unwrap_or(RowSelectionStrategy::Selectors),
50        }
51    }
52
53    /// Materialize a selection according to this policy, preserving its full
54    /// length, including trailing skips needed for selection composition.
55    pub(crate) fn apply(self, selection: RowSelection) -> RowSelection {
56        match (
57            self.resolve(Some(&selection)),
58            selection.as_mask().is_some(),
59        ) {
60            (RowSelectionStrategy::Mask, true) | (RowSelectionStrategy::Selectors, false) => {
61                selection
62            }
63            (RowSelectionStrategy::Mask, false) => selection.force_mask(),
64            (RowSelectionStrategy::Selectors, true) => {
65                RowSelection::from(Vec::<RowSelector>::from(selection))
66            }
67        }
68    }
69}
70
71/// Fully resolved strategy for materializing [`RowSelection`] during execution.
72///
73/// This is determined by [`RowSelectionPolicy`], including selector density for
74/// [`RowSelectionPolicy::Auto`].
75#[derive(Clone, Copy, Debug, Eq, PartialEq)]
76pub(crate) enum RowSelectionStrategy {
77    /// Use a queue of [`RowSelector`] values
78    Selectors,
79    /// Use a boolean mask to materialize the selection
80    Mask,
81}
82
83impl RowSelectionStrategy {
84    /// Build a cursor using this strategy. The selection must be trimmed of
85    /// trailing skips before constructing a mask cursor.
86    pub(crate) fn build_cursor(
87        self,
88        selection: RowSelection,
89        loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
90    ) -> RowSelectionCursor {
91        match self {
92            Self::Mask => RowSelectionCursor::new_mask_from_buffer(
93                selection.into_boolean_buffer(),
94                loaded_row_ranges,
95            ),
96            Self::Selectors => RowSelectionCursor::new_selectors(selection.into()),
97        }
98    }
99}
100
101/// Cursor for iterating a [`RowSelection`] during execution within a
102/// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan).
103///
104/// This keeps per-reader state such as the current position and delegates the
105/// actual storage strategy to the internal `RowSelectionInner`.
106#[derive(Debug)]
107pub enum RowSelectionCursor {
108    /// Reading all rows
109    All,
110    /// Use a bitmask to back the selection (dense selections)
111    Mask(MaskCursor),
112    /// Use a queue of selectors to back the selection (sparse selections)
113    Selectors(SelectorsCursor),
114}
115
116impl RowSelectionCursor {
117    /// Create a [`MaskCursor`] cursor backed by an existing bitmask.
118    pub(crate) fn new_mask_from_buffer(
119        mask: BooleanBuffer,
120        loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
121    ) -> Self {
122        debug_assert!(
123            mask.is_empty() || mask.value(mask.len() - 1),
124            "Mask selections must not end with a skip"
125        );
126        Self::Mask(MaskCursor {
127            mask,
128            position: 0,
129            loaded_row_ranges,
130        })
131    }
132
133    /// Create a [`RowSelectionCursor::Selectors`] from the provided selectors
134    pub(crate) fn new_selectors(selectors: Vec<RowSelector>) -> Self {
135        Self::Selectors(SelectorsCursor {
136            selectors: selectors.into(),
137            position: 0,
138        })
139    }
140
141    /// Create a cursor that selects all rows
142    pub(crate) fn new_all() -> Self {
143        Self::All
144    }
145}
146
147/// Cursor for iterating a selector-backed [`RowSelection`]
148///
149/// This is best for sparse selections where large contiguous
150/// blocks of rows are selected or skipped.
151#[derive(Debug)]
152pub struct SelectorsCursor {
153    selectors: VecDeque<RowSelector>,
154    /// Current absolute offset into the selection
155    position: usize,
156}
157
158impl SelectorsCursor {
159    /// Returns `true` when no further rows remain
160    pub fn is_empty(&self) -> bool {
161        self.selectors.is_empty()
162    }
163
164    /// Return the next [`RowSelector`]
165    pub(crate) fn next_selector(&mut self) -> RowSelector {
166        let selector = self.selectors.pop_front().unwrap();
167        self.position += selector.row_count;
168        selector
169    }
170
171    /// Return a selector to the front, rewinding the position
172    pub(crate) fn return_selector(&mut self, selector: RowSelector) {
173        self.position = self.position.saturating_sub(selector.row_count);
174        self.selectors.push_front(selector);
175    }
176}
177
178/// Cursor for iterating a mask-backed [`RowSelection`]
179///
180/// This is best for dense selections where there are many small skips
181/// or selections. For example, selecting every other row.
182///
183/// When page pruning produces sparse column data, `loaded_row_ranges` limits
184/// each decoded chunk to rows whose pages are loaded for every projected leaf.
185/// For example, two projected columns can have different page boundaries:
186///
187/// ```text
188/// Row ranges:       [0, 4) [4, 6) [6, 8) [8, 10) [10, 12)
189/// Selection mask:   1000   00     00     00      01
190/// Column A pages:   loaded | missing [4, 8) | loaded [8, 12)
191/// Column B pages:   loaded [0, 6) | missing [6, 10) | loaded
192/// LoadedRowRanges:  [0, 4)                         [10, 12)
193/// ```
194///
195/// The first chunk decodes `[0, 1)` with mask `1`. The next chunk skips to row
196/// 11 and decodes `[11, 12)` with mask `1`. When loaded ranges are present,
197/// every returned chunk ends at a selected row and never includes trailing
198/// unselected rows. [`ParquetRecordBatchReader`] still accumulates both chunks
199/// and applies the combined mask `11` once.
200///
201/// [`ParquetRecordBatchReader`]: crate::arrow::arrow_reader::ParquetRecordBatchReader
202#[derive(Debug)]
203pub struct MaskCursor {
204    mask: BooleanBuffer,
205    /// Current absolute offset into the selection
206    position: usize,
207    /// Row ranges whose backing pages are loaded for every projected column.
208    loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
209}
210
211impl MaskCursor {
212    /// Returns `true` when no further rows remain
213    pub fn is_empty(&self) -> bool {
214        self.position >= self.mask.len()
215    }
216
217    /// Advance through the mask representation, producing the next chunk summary
218    pub fn next_mask_chunk(&mut self, batch_size: usize) -> Option<MaskChunk> {
219        if self.is_empty() {
220            return None;
221        }
222
223        Some(self.next_mask_chunk_non_empty(batch_size))
224    }
225
226    /// Produces the next chunk for a non-empty, trailing-skip-free mask.
227    fn next_mask_chunk_non_empty(&mut self, batch_size: usize) -> MaskChunk {
228        debug_assert!(!self.is_empty());
229        let mask_start = self.next_selected_row().unwrap_or(self.mask.len());
230        debug_assert!(
231            mask_start < self.mask.len(),
232            "ReadPlan must remove trailing skips from Mask selections"
233        );
234        self.scan_chunk(mask_start, self.mask.len(), batch_size)
235    }
236
237    /// Find the first selected row without inspecting every skipped bit.
238    fn next_selected_row(&self) -> Option<usize> {
239        let chunks = BitChunks::new(
240            self.mask.values(),
241            self.mask.offset() + self.position,
242            self.mask.len() - self.position,
243        );
244        chunks.iter_padded().enumerate().find_map(|(idx, word)| {
245            (word != 0).then(|| self.position + idx * 64 + word.trailing_zeros() as usize)
246        })
247    }
248
249    /// Count whole words until the batch is full or the loaded range ends.
250    /// `BitChunks` normalizes words to little-endian and masks the remainder,
251    /// so padding and bits in an unloaded page cannot contribute to the count.
252    // Keep the scanning loop out of the reader's selection dispatch.
253    #[inline(never)]
254    fn scan_chunk(&mut self, mask_start: usize, end: usize, batch_size: usize) -> MaskChunk {
255        let chunks = BitChunks::new(
256            self.mask.values(),
257            self.mask.offset() + mask_start,
258            end - mask_start,
259        );
260        let mut selected_rows = 0;
261        let mut chunk_end = mask_start;
262
263        if batch_size > 0 {
264            for (idx, mut word) in chunks.iter_padded().enumerate() {
265                let count = word.count_ones() as usize;
266                if count == 0 {
267                    continue;
268                }
269
270                let remaining = batch_size - selected_rows;
271                let last_bit = if count <= remaining {
272                    selected_rows += count;
273                    63 - word.leading_zeros() as usize
274                } else {
275                    // Only the final word needs individual set-bit selection.
276                    for _ in 1..remaining {
277                        word &= word - 1;
278                    }
279                    selected_rows = batch_size;
280                    word.trailing_zeros() as usize
281                };
282                chunk_end = mask_start + idx * 64 + last_bit + 1;
283                if selected_rows == batch_size {
284                    break;
285                }
286            }
287        }
288
289        let initial_skip = mask_start - self.position;
290        // Stop immediately after the last selected row. Unset bits at the
291        // end of a loaded range belong to the next chunk's initial skip.
292        self.position = chunk_end;
293        MaskChunk {
294            initial_skip,
295            chunk_rows: chunk_end - mask_start,
296            selected_rows,
297            mask_start,
298        }
299    }
300
301    /// Returns the next non-empty mask chunk without crossing an unloaded row range.
302    /// When loaded ranges are present, every returned chunk ends immediately after a
303    /// selected row and therefore never contains trailing unselected rows. Those rows
304    /// remain for the next call's initial skip.
305    ///
306    /// The [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) removes trailing
307    /// skips before constructing this cursor. Callers therefore only invoke
308    /// this method for a non-empty mask that has another selected row.
309    pub(crate) fn next_chunk(&mut self, batch_size: usize) -> Result<MaskChunk, ParquetError> {
310        debug_assert!(batch_size > 0);
311        debug_assert!(!self.is_empty());
312
313        if self.loaded_row_ranges.is_none() {
314            return Ok(self.next_mask_chunk_non_empty(batch_size));
315        }
316
317        let mask_start = self.next_selected_row().ok_or_else(|| {
318            ParquetError::General(
319                "Internal Error: Mask cursor reached the end without finding a selected row; \
320                 ReadPlan must remove trailing skips"
321                    .to_string(),
322            )
323        })?;
324
325        let loaded_range_end = self
326            .loaded_row_ranges
327            .as_ref()
328            .and_then(|ranges| ranges.end_containing(mask_start))
329            .ok_or_else(|| {
330                ParquetError::General(format!(
331                    "Internal Error: selected row {mask_start} has no loaded page range"
332                ))
333            })?;
334
335        Ok(self.scan_chunk(
336            mask_start,
337            loaded_range_end.min(self.mask.len()),
338            batch_size,
339        ))
340    }
341
342    /// Materialise the boolean values for a mask-backed chunk
343    pub fn mask_values_for(&self, chunk: &MaskChunk) -> Result<BooleanArray, ParquetError> {
344        if chunk.mask_start.saturating_add(chunk.chunk_rows) > self.mask.len() {
345            return Err(ParquetError::General(
346                "Internal Error: MaskChunk exceeds mask length".to_string(),
347            ));
348        }
349        Ok(BooleanArray::from(
350            self.mask.slice(chunk.mask_start, chunk.chunk_rows),
351        ))
352    }
353}
354
355/// Result of computing the next chunk to read when using a [`MaskCursor`]
356#[derive(Debug)]
357pub struct MaskChunk {
358    /// Number of leading rows to skip before reaching selected rows
359    pub initial_skip: usize,
360    /// Total rows covered by this chunk (selected + skipped)
361    pub chunk_rows: usize,
362    /// Rows actually selected within the chunk
363    pub selected_rows: usize,
364    /// Starting offset within the mask where the chunk begins
365    pub mask_start: usize,
366}
367
368/// Row ranges whose backing pages are loaded for every projected column.
369#[derive(Clone, Debug)]
370pub(crate) struct LoadedRowRanges(Vec<Range<usize>>);
371
372impl LoadedRowRanges {
373    pub(crate) fn from_selection(selection: RowSelection) -> Self {
374        let selectors: Vec<RowSelector> = selection.into();
375        let mut position = 0;
376        let ranges = selectors
377            .into_iter()
378            .filter_map(|selector| {
379                let start = position;
380                position += selector.row_count;
381                (!selector.skip).then_some(start..position)
382            })
383            .collect();
384        Self(ranges)
385    }
386
387    fn end_containing(&self, row: usize) -> Option<usize> {
388        let idx = self.0.partition_point(|range| range.end <= row);
389        self.0
390            .get(idx)
391            .filter(|range| range.start <= row)
392            .map(|range| range.end)
393    }
394
395    #[cfg(test)]
396    pub(crate) fn ranges(&self) -> &[Range<usize>] {
397        &self.0
398    }
399}
400
401#[cfg(test)]
402mod tests {
403    use super::*;
404
405    /// Compare chunks with the selected row positions, independently of how
406    /// the cursor scans the bitmap. A chunk may consume at most `batch_size`
407    /// selected rows and may not leave the loaded range containing its first.
408    fn assert_mask_chunks(mask: BooleanBuffer, loaded: Option<LoadedRowRanges>, batch_size: usize) {
409        let selected: Vec<_> = (0..mask.len()).filter(|&row| mask.value(row)).collect();
410        let loaded = loaded.map(Arc::new);
411        let selection = RowSelection::from_boolean_buffer(mask).trim();
412        let RowSelectionCursor::Mask(mut cursor) =
413            RowSelectionStrategy::Mask.build_cursor(selection, loaded.clone())
414        else {
415            unreachable!()
416        };
417        let mut next = 0;
418        let mut position = 0;
419        while next < selected.len() {
420            let first = selected[next];
421            let end = loaded
422                .as_ref()
423                .map_or(usize::MAX, |ranges| ranges.end_containing(first).unwrap());
424            let count = selected[next..]
425                .iter()
426                .take(batch_size)
427                .take_while(|&&row| row < end)
428                .count();
429            let last = selected[next + count - 1];
430            let chunk = cursor.next_chunk(batch_size).unwrap();
431            assert_eq!(chunk.initial_skip, first - position);
432            assert_eq!(chunk.mask_start, first);
433            assert_eq!(chunk.chunk_rows, last + 1 - first);
434            assert_eq!(chunk.selected_rows, count);
435            assert_eq!(cursor.position, last + 1);
436            let values = cursor.mask_values_for(&chunk).unwrap();
437            assert_eq!(values.true_count(), count);
438            assert!(values.value(0));
439            assert!(values.value(values.len() - 1));
440            next += count;
441            position = last + 1;
442        }
443        assert!(cursor.is_empty());
444        assert!(cursor.next_mask_chunk(batch_size).is_none());
445    }
446
447    #[test]
448    fn test_mask_chunks_sparse_aligned_words() {
449        // Distinct byte positions expose native-endian word interpretation
450        // on big-endian platforms. Both selections must form one-row chunks.
451        let mask = BooleanBuffer::from_iter((0..256).map(|row| row == 0 || row == 255));
452        assert_mask_chunks(mask, None, 1);
453    }
454
455    #[test]
456    fn test_mask_chunks_word_boundaries_and_offsets() {
457        use rand::{RngExt, SeedableRng, rngs::StdRng};
458
459        for len in [0, 1, 7, 8, 63, 64, 65, 127, 128, 129, 511, 2049] {
460            for offset in [0, 1, 3, 7, 8, 31, 63, 64, 65] {
461                for pattern in 0..7 {
462                    let mut rng = StdRng::seed_from_u64(42);
463                    let mask = BooleanBuffer::from_iter((0..offset + len + 64).map(|row| {
464                        if !(offset..offset + len).contains(&row) {
465                            // Set padding on both sides to catch reads outside
466                            // the logical slice, including its final word.
467                            return true;
468                        }
469                        let row = row - offset;
470                        match pattern {
471                            0 => false,
472                            1 => true,
473                            2 => row % 2 == 0,
474                            3 => row % 129 == 128,
475                            4 => rng.random_bool(0.1),
476                            5 => rng.random_bool(0.5),
477                            _ => rng.random_bool(0.9),
478                        }
479                    }))
480                    .slice(offset, len);
481                    for batch_size in [1, 2, 3, 63, 64, 65, 1024] {
482                        assert_mask_chunks(mask.clone(), None, batch_size);
483                    }
484                }
485            }
486        }
487    }
488
489    #[test]
490    fn test_mask_chunks_loaded_ranges_word_boundaries() {
491        for page_rows in [1, 31, 63, 64, 65, 127, 257] {
492            let len = page_rows * 7 + 13;
493            let ranges: Vec<_> = (0..len)
494                .step_by(page_rows)
495                .enumerate()
496                .filter(|(page, _)| page % 3 != 1)
497                .map(|(_, start)| start..(start + page_rows).min(len))
498                .collect();
499            let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges(
500                ranges.iter().cloned(),
501                len,
502            ));
503            for offset in [0, 1, 3, 7, 31, 63, 64, 65] {
504                for pattern in 0..4 {
505                    let mask = BooleanBuffer::from_iter((0..offset + len + 64).map(|row| {
506                        if !(offset..offset + len).contains(&row) {
507                            return true;
508                        }
509                        let row = row - offset;
510                        ranges.iter().any(|range| range.contains(&row))
511                            && match pattern {
512                                0 => true,
513                                1 => row % 2 == 0,
514                                // Leave an unset tail in every loaded range.
515                                2 => row % page_rows == 0,
516                                // Select directly on either side of page and
517                                // word boundaries, including the final page.
518                                _ => {
519                                    row % page_rows == page_rows - 1
520                                        || row % 64 == 0
521                                        || row % 64 == 63
522                                }
523                            }
524                    }))
525                    .slice(offset, len);
526                    for batch_size in [1, 2, 3, 63, 64, 65, 1024] {
527                        assert_mask_chunks(mask.clone(), Some(loaded.clone()), batch_size);
528                    }
529                }
530            }
531        }
532    }
533
534    #[test]
535    fn test_mask_chunk_does_not_read_selected_rows_outside_loaded_range() {
536        let mut cursor = MaskCursor {
537            mask: BooleanBuffer::new_set(130),
538            position: 0,
539            loaded_row_ranges: Some(Arc::new(LoadedRowRanges(vec![0..65, 66..130]))),
540        };
541        let chunk = cursor.next_chunk(1024).unwrap();
542        assert_eq!(chunk.chunk_rows, 65);
543        assert_eq!(chunk.selected_rows, 65);
544        assert_eq!(cursor.position, 65);
545        let err = cursor.next_chunk(1024).unwrap_err();
546        assert!(
547            err.to_string()
548                .contains("selected row 65 has no loaded page range")
549        );
550        assert_eq!(cursor.position, 65);
551    }
552
553    #[test]
554    fn selection_respects_policy() {
555        let mask = BooleanBuffer::from(vec![true, true, false, false, true, true, false, false]);
556        for (policy, expect_mask) in [
557            (RowSelectionPolicy::Mask, true),
558            (RowSelectionPolicy::Selectors, false),
559            (RowSelectionPolicy::Auto { threshold: 2 }, false),
560            (RowSelectionPolicy::Auto { threshold: 3 }, true),
561        ] {
562            for selection in [
563                RowSelection::from_boolean_buffer(mask.clone()),
564                RowSelection::from(vec![
565                    RowSelector::select(2),
566                    RowSelector::skip(2),
567                    RowSelector::select(2),
568                    RowSelector::skip(2),
569                ]),
570            ] {
571                let selection = policy.apply(selection);
572                assert_eq!(selection.as_mask().is_some(), expect_mask, "{policy:?}");
573                assert_eq!(selection.into_boolean_buffer(), mask);
574            }
575        }
576    }
577
578    #[test]
579    fn adaptive_selection_uses_selectors_for_long_runs() {
580        let mask = BooleanBuffer::from_iter((0..1_024).map(|idx| (256..768).contains(&idx)));
581        let selection =
582            RowSelectionPolicy::default().apply(RowSelection::from_boolean_buffer(mask));
583        assert!(selection.as_mask().is_none());
584    }
585
586    #[test]
587    fn adaptive_selection_keeps_fragmented_masks() {
588        let mask = BooleanBuffer::from_iter((0..1_024).map(|idx| idx % 2 == 0));
589        let selection =
590            RowSelectionPolicy::default().apply(RowSelection::from_boolean_buffer(mask));
591        assert!(selection.as_mask().is_some());
592    }
593
594    #[test]
595    fn test_loaded_mask_chunk_stops_at_trimmed_mask_end() {
596        let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges(
597            std::iter::once(0..5),
598            10,
599        ));
600        let RowSelectionCursor::Mask(mut cursor) = RowSelectionStrategy::Mask.build_cursor(
601            RowSelection::from(vec![RowSelector::select(1)]),
602            Some(loaded.into()),
603        ) else {
604            unreachable!()
605        };
606
607        let chunk = cursor.next_chunk(10).unwrap();
608        assert_eq!(chunk.chunk_rows, 1);
609        assert!(cursor.is_empty());
610    }
611
612    #[test]
613    fn test_next_mask_chunk_until_cursor_is_empty() {
614        let RowSelectionCursor::Mask(mut cursor) = RowSelectionStrategy::Mask.build_cursor(
615            RowSelection::from(vec![
616                RowSelector::skip(2),
617                RowSelector::select(2),
618                RowSelector::skip(1),
619                RowSelector::select(1),
620            ]),
621            None,
622        ) else {
623            unreachable!()
624        };
625
626        let first = cursor.next_mask_chunk(2).unwrap();
627        assert_eq!(first.initial_skip, 2);
628        assert_eq!(first.chunk_rows, 2);
629        assert_eq!(first.selected_rows, 2);
630
631        let second = cursor.next_mask_chunk(2).unwrap();
632        assert_eq!(second.initial_skip, 1);
633        assert_eq!(second.chunk_rows, 1);
634        assert_eq!(second.selected_rows, 1);
635
636        assert!(cursor.next_mask_chunk(2).is_none());
637    }
638}
parquet/arrow/arrow_reader/selection/cursor.rs

parquet/arrow/arrow_reader/selection/
cursor.rs
