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::boolean::boolean_mask_from_selectors;
26use super::{RowSelection, RowSelector};
27use crate::errors::ParquetError;
28use arrow_array::BooleanArray;
29use arrow_buffer::BooleanBuffer;
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
54/// Fully resolved strategy for materializing [`RowSelection`] during execution.
55///
56/// This is determined by [`RowSelectionPolicy`], including selector density for
57/// [`RowSelectionPolicy::Auto`].
58#[derive(Clone, Copy, Debug, Eq, PartialEq)]
59pub(crate) enum RowSelectionStrategy {
60    /// Use a queue of [`RowSelector`] values
61    Selectors,
62    /// Use a boolean mask to materialize the selection
63    Mask,
64}
65
66/// Cursor for iterating a [`RowSelection`] during execution within a
67/// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan).
68///
69/// This keeps per-reader state such as the current position and delegates the
70/// actual storage strategy to the internal `RowSelectionInner`.
71#[derive(Debug)]
72pub enum RowSelectionCursor {
73    /// Reading all rows
74    All,
75    /// Use a bitmask to back the selection (dense selections)
76    Mask(MaskCursor),
77    /// Use a queue of selectors to back the selection (sparse selections)
78    Selectors(SelectorsCursor),
79}
80
81impl RowSelectionCursor {
82    /// Create a [`MaskCursor`] cursor backed by a bitmask, from an existing set of selectors
83    pub(crate) fn new_mask_from_selectors(
84        selectors: Vec<RowSelector>,
85        loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
86    ) -> Self {
87        debug_assert!(
88            selectors
89                .last()
90                .map(|selector| !selector.skip)
91                .unwrap_or(true),
92            "Mask selectors must not end with a skip"
93        );
94        Self::Mask(MaskCursor {
95            mask: boolean_mask_from_selectors(&selectors),
96            position: 0,
97            loaded_row_ranges,
98        })
99    }
100
101    /// Create a [`MaskCursor`] cursor backed by an existing bitmask.
102    pub(crate) fn new_mask_from_buffer(
103        mask: BooleanBuffer,
104        loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
105    ) -> Self {
106        debug_assert!(
107            mask.is_empty() || mask.value(mask.len() - 1),
108            "Mask selections must not end with a skip"
109        );
110        Self::Mask(MaskCursor {
111            mask,
112            position: 0,
113            loaded_row_ranges,
114        })
115    }
116
117    /// Create a [`RowSelectionCursor::Selectors`] from the provided selectors
118    pub(crate) fn new_selectors(selectors: Vec<RowSelector>) -> Self {
119        Self::Selectors(SelectorsCursor {
120            selectors: selectors.into(),
121            position: 0,
122        })
123    }
124
125    /// Create a cursor that selects all rows
126    pub(crate) fn new_all() -> Self {
127        Self::All
128    }
129}
130
131/// Cursor for iterating a selector-backed [`RowSelection`]
132///
133/// This is best for sparse selections where large contiguous
134/// blocks of rows are selected or skipped.
135#[derive(Debug)]
136pub struct SelectorsCursor {
137    selectors: VecDeque<RowSelector>,
138    /// Current absolute offset into the selection
139    position: usize,
140}
141
142impl SelectorsCursor {
143    /// Returns `true` when no further rows remain
144    pub fn is_empty(&self) -> bool {
145        self.selectors.is_empty()
146    }
147
148    pub(crate) fn selectors_mut(&mut self) -> &mut VecDeque<RowSelector> {
149        &mut self.selectors
150    }
151
152    /// Return the next [`RowSelector`]
153    pub(crate) fn next_selector(&mut self) -> RowSelector {
154        let selector = self.selectors.pop_front().unwrap();
155        self.position += selector.row_count;
156        selector
157    }
158
159    /// Return a selector to the front, rewinding the position
160    pub(crate) fn return_selector(&mut self, selector: RowSelector) {
161        self.position = self.position.saturating_sub(selector.row_count);
162        self.selectors.push_front(selector);
163    }
164}
165
166/// Cursor for iterating a mask-backed [`RowSelection`]
167///
168/// This is best for dense selections where there are many small skips
169/// or selections. For example, selecting every other row.
170///
171/// When page pruning produces sparse column data, `loaded_row_ranges` limits
172/// each decoded chunk to rows whose pages are loaded for every projected leaf.
173/// For example, two projected columns can have different page boundaries:
174///
175/// ```text
176/// Row ranges:       [0, 4) [4, 6) [6, 8) [8, 10) [10, 12)
177/// Selection mask:   1000   00     00     00      01
178/// Column A pages:   loaded | missing [4, 8) | loaded [8, 12)
179/// Column B pages:   loaded [0, 6) | missing [6, 10) | loaded
180/// LoadedRowRanges:  [0, 4)                         [10, 12)
181/// ```
182///
183/// The first chunk decodes `[0, 4)` with mask `1000`. The next chunk skips to
184/// row 11 and decodes `[11, 12)` with mask `1`. The loaded ranges are decode
185/// boundaries, not output batch boundaries: [`ParquetRecordBatchReader`]
186/// accumulates both chunks and applies the combined mask `10001` once.
187///
188/// [`ParquetRecordBatchReader`]: crate::arrow::arrow_reader::ParquetRecordBatchReader
189#[derive(Debug)]
190pub struct MaskCursor {
191    mask: BooleanBuffer,
192    /// Current absolute offset into the selection
193    position: usize,
194    /// Row ranges whose backing pages are loaded for every projected column.
195    loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
196}
197
198impl MaskCursor {
199    /// Returns `true` when no further rows remain
200    pub fn is_empty(&self) -> bool {
201        self.position >= self.mask.len()
202    }
203
204    /// Advance through the mask representation, producing the next chunk summary
205    pub fn next_mask_chunk(&mut self, batch_size: usize) -> Option<MaskChunk> {
206        if self.is_empty() {
207            return None;
208        }
209
210        Some(self.next_mask_chunk_non_empty(batch_size))
211    }
212
213    /// Produces the next chunk for a non-empty, trailing-skip-free mask.
214    fn next_mask_chunk_non_empty(&mut self, batch_size: usize) -> MaskChunk {
215        debug_assert!(!self.is_empty());
216
217        let (initial_skip, chunk_rows, selected_rows, mask_start, end_position) = {
218            let mask = &self.mask;
219            let start_position = self.position;
220            let mut cursor = start_position;
221            let mut initial_skip = 0;
222
223            while cursor < mask.len() && !mask.value(cursor) {
224                initial_skip += 1;
225                cursor += 1;
226            }
227            debug_assert!(
228                cursor < mask.len(),
229                "ReadPlan must remove trailing skips from Mask selections"
230            );
231
232            let mask_start = cursor;
233            let mut chunk_rows = 0;
234            let mut selected_rows = 0;
235
236            // Advance until enough rows have been selected to satisfy the batch size,
237            // or until the mask is exhausted. This mirrors the behaviour of the legacy
238            // `RowSelector` queue-based iteration.
239            while cursor < mask.len() && selected_rows < batch_size {
240                chunk_rows += 1;
241                if mask.value(cursor) {
242                    selected_rows += 1;
243                }
244                cursor += 1;
245            }
246
247            (initial_skip, chunk_rows, selected_rows, mask_start, cursor)
248        };
249
250        self.position = end_position;
251
252        MaskChunk {
253            initial_skip,
254            chunk_rows,
255            selected_rows,
256            mask_start,
257        }
258    }
259
260    /// Returns the next non-empty mask chunk without crossing an unloaded row range.
261    ///
262    /// The [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) removes trailing
263    /// skips before constructing this cursor. Callers therefore only invoke
264    /// this method for a non-empty mask that has another selected row.
265    pub(crate) fn next_chunk(&mut self, batch_size: usize) -> Result<MaskChunk, ParquetError> {
266        debug_assert!(batch_size > 0);
267        debug_assert!(!self.is_empty());
268
269        if self.loaded_row_ranges.is_none() {
270            return Ok(self.next_mask_chunk_non_empty(batch_size));
271        }
272
273        let start_position = self.position;
274        let mut cursor = start_position;
275        while cursor < self.mask.len() && !self.mask.value(cursor) {
276            cursor += 1;
277        }
278
279        debug_assert!(
280            cursor < self.mask.len(),
281            "ReadPlan must remove trailing skips from Mask selections"
282        );
283
284        let loaded_range_end = self
285            .loaded_row_ranges
286            .as_ref()
287            .and_then(|ranges| ranges.end_containing(cursor))
288            .ok_or_else(|| {
289                ParquetError::General(format!(
290                    "Internal Error: selected row {cursor} has no loaded page range"
291                ))
292            })?;
293
294        let mask_start = cursor;
295        let mut selected_rows = 0;
296        while cursor < loaded_range_end && cursor < self.mask.len() && selected_rows < batch_size {
297            if self.mask.value(cursor) {
298                selected_rows += 1;
299            }
300            cursor += 1;
301        }
302
303        self.position = cursor;
304        Ok(MaskChunk {
305            initial_skip: mask_start - start_position,
306            chunk_rows: cursor - mask_start,
307            selected_rows,
308            mask_start,
309        })
310    }
311
312    /// Materialise the boolean values for a mask-backed chunk
313    pub fn mask_values_for(&self, chunk: &MaskChunk) -> Result<BooleanArray, ParquetError> {
314        if chunk.mask_start.saturating_add(chunk.chunk_rows) > self.mask.len() {
315            return Err(ParquetError::General(
316                "Internal Error: MaskChunk exceeds mask length".to_string(),
317            ));
318        }
319        Ok(BooleanArray::from(
320            self.mask.slice(chunk.mask_start, chunk.chunk_rows),
321        ))
322    }
323}
324
325/// Result of computing the next chunk to read when using a [`MaskCursor`]
326#[derive(Debug)]
327pub struct MaskChunk {
328    /// Number of leading rows to skip before reaching selected rows
329    pub initial_skip: usize,
330    /// Total rows covered by this chunk (selected + skipped)
331    pub chunk_rows: usize,
332    /// Rows actually selected within the chunk
333    pub selected_rows: usize,
334    /// Starting offset within the mask where the chunk begins
335    pub mask_start: usize,
336}
337
338/// Row ranges whose backing pages are loaded for every projected column.
339#[derive(Clone, Debug)]
340pub(crate) struct LoadedRowRanges(Vec<Range<usize>>);
341
342impl LoadedRowRanges {
343    pub(crate) fn from_selection(selection: RowSelection) -> Self {
344        let selectors: Vec<RowSelector> = selection.into();
345        let mut position = 0;
346        let ranges = selectors
347            .into_iter()
348            .filter_map(|selector| {
349                let start = position;
350                position += selector.row_count;
351                (!selector.skip).then_some(start..position)
352            })
353            .collect();
354        Self(ranges)
355    }
356
357    fn end_containing(&self, row: usize) -> Option<usize> {
358        let idx = self.0.partition_point(|range| range.end <= row);
359        self.0
360            .get(idx)
361            .filter(|range| range.start <= row)
362            .map(|range| range.end)
363    }
364
365    #[cfg(test)]
366    pub(crate) fn ranges(&self) -> &[Range<usize>] {
367        &self.0
368    }
369}
370
371#[cfg(test)]
372mod tests {
373    use super::*;
374
375    #[test]
376    fn test_loaded_mask_chunk_stops_at_trimmed_mask_end() {
377        let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges(
378            std::iter::once(0..5),
379            10,
380        ));
381        let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors(
382            vec![RowSelector::select(1)],
383            Some(loaded.into()),
384        ) else {
385            unreachable!()
386        };
387
388        let chunk = cursor.next_chunk(10).unwrap();
389        assert_eq!(chunk.chunk_rows, 1);
390        assert!(cursor.is_empty());
391    }
392
393    #[test]
394    fn test_next_mask_chunk_until_cursor_is_empty() {
395        let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors(
396            vec![
397                RowSelector::skip(2),
398                RowSelector::select(2),
399                RowSelector::skip(1),
400                RowSelector::select(1),
401            ],
402            None,
403        ) else {
404            unreachable!()
405        };
406
407        let first = cursor.next_mask_chunk(2).unwrap();
408        assert_eq!(first.initial_skip, 2);
409        assert_eq!(first.chunk_rows, 2);
410        assert_eq!(first.selected_rows, 2);
411
412        let second = cursor.next_mask_chunk(2).unwrap();
413        assert_eq!(second.initial_skip, 1);
414        assert_eq!(second.chunk_rows, 1);
415        assert_eq!(second.selected_rows, 1);
416
417        assert!(cursor.next_mask_chunk(2).is_none());
418    }
419}
parquet/arrow/arrow_reader/selection/cursor.rs

parquet/arrow/arrow_reader/selection/
cursor.rs
