Skip to main content

parquet/arrow/arrow_reader/selection/
cursor.rs

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