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    /// Return the next [`RowSelector`]
149    pub(crate) fn next_selector(&mut self) -> RowSelector {
150        let selector = self.selectors.pop_front().unwrap();
151        self.position += selector.row_count;
152        selector
153    }
154
155    /// Return a selector to the front, rewinding the position
156    pub(crate) fn return_selector(&mut self, selector: RowSelector) {
157        self.position = self.position.saturating_sub(selector.row_count);
158        self.selectors.push_front(selector);
159    }
160}
161
162/// Cursor for iterating a mask-backed [`RowSelection`]
163///
164/// This is best for dense selections where there are many small skips
165/// or selections. For example, selecting every other row.
166///
167/// When page pruning produces sparse column data, `loaded_row_ranges` limits
168/// each decoded chunk to rows whose pages are loaded for every projected leaf.
169/// For example, two projected columns can have different page boundaries:
170///
171/// ```text
172/// Row ranges:       [0, 4) [4, 6) [6, 8) [8, 10) [10, 12)
173/// Selection mask:   1000   00     00     00      01
174/// Column A pages:   loaded | missing [4, 8) | loaded [8, 12)
175/// Column B pages:   loaded [0, 6) | missing [6, 10) | loaded
176/// LoadedRowRanges:  [0, 4)                         [10, 12)
177/// ```
178///
179/// The first chunk decodes `[0, 4)` with mask `1000`. The next chunk skips to
180/// row 11 and decodes `[11, 12)` with mask `1`. The loaded ranges are decode
181/// boundaries, not output batch boundaries: [`ParquetRecordBatchReader`]
182/// accumulates both chunks and applies the combined mask `10001` once.
183///
184/// [`ParquetRecordBatchReader`]: crate::arrow::arrow_reader::ParquetRecordBatchReader
185#[derive(Debug)]
186pub struct MaskCursor {
187    mask: BooleanBuffer,
188    /// Current absolute offset into the selection
189    position: usize,
190    /// Row ranges whose backing pages are loaded for every projected column.
191    loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
192}
193
194impl MaskCursor {
195    /// Returns `true` when no further rows remain
196    pub fn is_empty(&self) -> bool {
197        self.position >= self.mask.len()
198    }
199
200    /// Advance through the mask representation, producing the next chunk summary
201    pub fn next_mask_chunk(&mut self, batch_size: usize) -> Option<MaskChunk> {
202        if self.is_empty() {
203            return None;
204        }
205
206        Some(self.next_mask_chunk_non_empty(batch_size))
207    }
208
209    /// Produces the next chunk for a non-empty, trailing-skip-free mask.
210    fn next_mask_chunk_non_empty(&mut self, batch_size: usize) -> MaskChunk {
211        debug_assert!(!self.is_empty());
212
213        let (initial_skip, chunk_rows, selected_rows, mask_start, end_position) = {
214            let mask = &self.mask;
215            let start_position = self.position;
216            let mut cursor = start_position;
217            let mut initial_skip = 0;
218
219            while cursor < mask.len() && !mask.value(cursor) {
220                initial_skip += 1;
221                cursor += 1;
222            }
223            debug_assert!(
224                cursor < mask.len(),
225                "ReadPlan must remove trailing skips from Mask selections"
226            );
227
228            let mask_start = cursor;
229            let mut chunk_rows = 0;
230            let mut selected_rows = 0;
231
232            // Advance until enough rows have been selected to satisfy the batch size,
233            // or until the mask is exhausted. This mirrors the behaviour of the legacy
234            // `RowSelector` queue-based iteration.
235            while cursor < mask.len() && selected_rows < batch_size {
236                chunk_rows += 1;
237                if mask.value(cursor) {
238                    selected_rows += 1;
239                }
240                cursor += 1;
241            }
242
243            (initial_skip, chunk_rows, selected_rows, mask_start, cursor)
244        };
245
246        self.position = end_position;
247
248        MaskChunk {
249            initial_skip,
250            chunk_rows,
251            selected_rows,
252            mask_start,
253        }
254    }
255
256    /// Returns the next non-empty mask chunk without crossing an unloaded row range.
257    ///
258    /// The [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) removes trailing
259    /// skips before constructing this cursor. Callers therefore only invoke
260    /// this method for a non-empty mask that has another selected row.
261    pub(crate) fn next_chunk(&mut self, batch_size: usize) -> Result<MaskChunk, ParquetError> {
262        debug_assert!(batch_size > 0);
263        debug_assert!(!self.is_empty());
264
265        if self.loaded_row_ranges.is_none() {
266            return Ok(self.next_mask_chunk_non_empty(batch_size));
267        }
268
269        let start_position = self.position;
270        let mut cursor = start_position;
271        while cursor < self.mask.len() && !self.mask.value(cursor) {
272            cursor += 1;
273        }
274
275        debug_assert!(
276            cursor < self.mask.len(),
277            "ReadPlan must remove trailing skips from Mask selections"
278        );
279
280        let loaded_range_end = self
281            .loaded_row_ranges
282            .as_ref()
283            .and_then(|ranges| ranges.end_containing(cursor))
284            .ok_or_else(|| {
285                ParquetError::General(format!(
286                    "Internal Error: selected row {cursor} has no loaded page range"
287                ))
288            })?;
289
290        let mask_start = cursor;
291        let mut selected_rows = 0;
292        while cursor < loaded_range_end && cursor < self.mask.len() && selected_rows < batch_size {
293            if self.mask.value(cursor) {
294                selected_rows += 1;
295            }
296            cursor += 1;
297        }
298
299        self.position = cursor;
300        Ok(MaskChunk {
301            initial_skip: mask_start - start_position,
302            chunk_rows: cursor - mask_start,
303            selected_rows,
304            mask_start,
305        })
306    }
307
308    /// Materialise the boolean values for a mask-backed chunk
309    pub fn mask_values_for(&self, chunk: &MaskChunk) -> Result<BooleanArray, ParquetError> {
310        if chunk.mask_start.saturating_add(chunk.chunk_rows) > self.mask.len() {
311            return Err(ParquetError::General(
312                "Internal Error: MaskChunk exceeds mask length".to_string(),
313            ));
314        }
315        Ok(BooleanArray::from(
316            self.mask.slice(chunk.mask_start, chunk.chunk_rows),
317        ))
318    }
319}
320
321/// Result of computing the next chunk to read when using a [`MaskCursor`]
322#[derive(Debug)]
323pub struct MaskChunk {
324    /// Number of leading rows to skip before reaching selected rows
325    pub initial_skip: usize,
326    /// Total rows covered by this chunk (selected + skipped)
327    pub chunk_rows: usize,
328    /// Rows actually selected within the chunk
329    pub selected_rows: usize,
330    /// Starting offset within the mask where the chunk begins
331    pub mask_start: usize,
332}
333
334/// Row ranges whose backing pages are loaded for every projected column.
335#[derive(Clone, Debug)]
336pub(crate) struct LoadedRowRanges(Vec<Range<usize>>);
337
338impl LoadedRowRanges {
339    pub(crate) fn from_selection(selection: RowSelection) -> Self {
340        let selectors: Vec<RowSelector> = selection.into();
341        let mut position = 0;
342        let ranges = selectors
343            .into_iter()
344            .filter_map(|selector| {
345                let start = position;
346                position += selector.row_count;
347                (!selector.skip).then_some(start..position)
348            })
349            .collect();
350        Self(ranges)
351    }
352
353    fn end_containing(&self, row: usize) -> Option<usize> {
354        let idx = self.0.partition_point(|range| range.end <= row);
355        self.0
356            .get(idx)
357            .filter(|range| range.start <= row)
358            .map(|range| range.end)
359    }
360
361    #[cfg(test)]
362    pub(crate) fn ranges(&self) -> &[Range<usize>] {
363        &self.0
364    }
365}
366
367#[cfg(test)]
368mod tests {
369    use super::*;
370
371    #[test]
372    fn test_loaded_mask_chunk_stops_at_trimmed_mask_end() {
373        let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges(
374            std::iter::once(0..5),
375            10,
376        ));
377        let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors(
378            vec![RowSelector::select(1)],
379            Some(loaded.into()),
380        ) else {
381            unreachable!()
382        };
383
384        let chunk = cursor.next_chunk(10).unwrap();
385        assert_eq!(chunk.chunk_rows, 1);
386        assert!(cursor.is_empty());
387    }
388
389    #[test]
390    fn test_next_mask_chunk_until_cursor_is_empty() {
391        let RowSelectionCursor::Mask(mut cursor) = RowSelectionCursor::new_mask_from_selectors(
392            vec![
393                RowSelector::skip(2),
394                RowSelector::select(2),
395                RowSelector::skip(1),
396                RowSelector::select(1),
397            ],
398            None,
399        ) else {
400            unreachable!()
401        };
402
403        let first = cursor.next_mask_chunk(2).unwrap();
404        assert_eq!(first.initial_skip, 2);
405        assert_eq!(first.chunk_rows, 2);
406        assert_eq!(first.selected_rows, 2);
407
408        let second = cursor.next_mask_chunk(2).unwrap();
409        assert_eq!(second.initial_skip, 1);
410        assert_eq!(second.chunk_rows, 1);
411        assert_eq!(second.selected_rows, 1);
412
413        assert!(cursor.next_mask_chunk(2).is_none());
414    }
415}
parquet/arrow/arrow_reader/selection/cursor.rs

parquet/arrow/arrow_reader/selection/
cursor.rs
