parquet/arrow/arrow_reader/
read_plan.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//! [`ReadPlan`] and [`ReadPlanBuilder`] for determining which rows to read
19//! from a Parquet file
20
21use crate::arrow::array_reader::ArrayReader;
22use crate::arrow::arrow_reader::selection::{
23    LoadedRowRanges, RowSelectionPolicy, RowSelectionStrategy,
24};
25use crate::arrow::arrow_reader::{
26    ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelectionCursor, RowSelector,
27};
28use crate::errors::{ParquetError, Result};
29use arrow_array::{Array, BooleanArray};
30use arrow_buffer::BooleanBuffer;
31use arrow_select::filter::prep_null_mask_filter;
32use std::collections::VecDeque;
33use std::sync::Arc;
34
35/// Options for [`ReadPlanBuilder::with_predicate_options`].
36pub struct PredicateOptions<'a> {
37    array_reader: Box<dyn ArrayReader>,
38    predicate: &'a mut dyn ArrowPredicate,
39    limit: Option<usize>,
40    total_rows: usize,
41}
42
43impl<'a> PredicateOptions<'a> {
44    /// Create options for evaluating `predicate` against rows produced by
45    /// `array_reader`.
46    ///
47    /// By default there is no match-count limit; the predicate is evaluated
48    /// over every row the reader yields. Use [`Self::with_limit`] to enable
49    /// early termination.
50    pub fn new(array_reader: Box<dyn ArrayReader>, predicate: &'a mut dyn ArrowPredicate) -> Self {
51        Self {
52            array_reader,
53            predicate,
54            limit: None,
55            total_rows: 0,
56        }
57    }
58
59    /// Stop scanning `array_reader` once `limit` matches have accumulated.
60    ///
61    /// Performance optimization for `LIMIT` / TopK: when the cumulative
62    /// `true_count` reaches `limit`, the current filter batch is truncated
63    /// at the `limit`-th match and remaining batches are never decoded.
64    ///
65    /// `limit` counts predicate matches, not output rows — callers applying
66    /// an offset must pass `offset + limit`.
67    ///
68    /// `total_rows` is the row count `array_reader` would yield if iterated
69    /// to completion. It is used to pad un-evaluated trailing rows as "not
70    /// selected" so the returned [`RowSelection`] covers the full row group.
71    ///
72    /// Only valid for the *last* predicate in a filter chain: intermediate
73    /// predicates' match counts do not map 1:1 to output rows.
74    pub fn with_limit(mut self, limit: usize, total_rows: usize) -> Self {
75        self.limit = Some(limit);
76        self.total_rows = total_rows;
77        self
78    }
79}
80
81/// A builder for [`ReadPlan`]
82#[derive(Clone, Debug)]
83pub struct ReadPlanBuilder {
84    batch_size: usize,
85    /// Which rows to select. Includes the result of all filters applied so far
86    selection: Option<RowSelection>,
87    /// Policy to use when materializing the row selection
88    row_selection_policy: RowSelectionPolicy,
89    /// Row ranges with page data loaded for the current projection.
90    loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
91}
92
93impl ReadPlanBuilder {
94    /// Create a `ReadPlanBuilder` with the given batch size
95    pub fn new(batch_size: usize) -> Self {
96        Self {
97            batch_size,
98            selection: None,
99            row_selection_policy: RowSelectionPolicy::default(),
100            loaded_row_ranges: None,
101        }
102    }
103
104    /// Set the current selection to the given value
105    pub fn with_selection(mut self, selection: Option<RowSelection>) -> Self {
106        self.selection = selection;
107        self
108    }
109
110    /// Configure the policy to use when materialising the [`RowSelection`]
111    ///
112    /// Defaults to [`RowSelectionPolicy::Auto`]
113    pub fn with_row_selection_policy(mut self, policy: RowSelectionPolicy) -> Self {
114        self.row_selection_policy = policy;
115        self
116    }
117
118    pub(crate) fn with_loaded_row_ranges(mut self, ranges: Option<LoadedRowRanges>) -> Self {
119        self.loaded_row_ranges = ranges.map(Arc::new);
120        self
121    }
122
123    /// Returns the current row selection policy
124    pub fn row_selection_policy(&self) -> &RowSelectionPolicy {
125        &self.row_selection_policy
126    }
127
128    /// Returns the current selection, if any
129    pub fn selection(&self) -> Option<&RowSelection> {
130        self.selection.as_ref()
131    }
132
133    /// Specifies the number of rows in the row group, before filtering is applied.
134    ///
135    /// Returns a [`LimitedReadPlanBuilder`] that can apply
136    /// offset and limit.
137    ///
138    /// Call [`LimitedReadPlanBuilder::build_limited`] to apply the limits to this
139    /// selection.
140    pub(crate) fn limited(self, row_count: usize) -> LimitedReadPlanBuilder {
141        LimitedReadPlanBuilder::new(self, row_count)
142    }
143
144    /// Returns true if the current plan selects any rows
145    pub fn selects_any(&self) -> bool {
146        self.selection
147            .as_ref()
148            .map(|s| s.selects_any())
149            .unwrap_or(true)
150    }
151
152    /// Returns the number of rows selected, or `None` if all rows are selected.
153    pub fn num_rows_selected(&self) -> Option<usize> {
154        self.selection.as_ref().map(|s| s.row_count())
155    }
156
157    /// Returns the [`RowSelectionStrategy`] for this plan.
158    ///
159    /// Guarantees to return either `Selectors` or `Mask`, never `Auto`.
160    pub(crate) fn resolve_selection_strategy(&self) -> RowSelectionStrategy {
161        match self.row_selection_policy {
162            RowSelectionPolicy::Selectors => RowSelectionStrategy::Selectors,
163            RowSelectionPolicy::Mask => RowSelectionStrategy::Mask,
164            RowSelectionPolicy::Auto { threshold, .. } => {
165                let selection = match self.selection.as_ref() {
166                    Some(selection) => selection,
167                    None => return RowSelectionStrategy::Selectors,
168                };
169
170                // total_rows: total number of rows selected / skipped
171                // effective_count: number of non-empty selectors
172                let (total_rows, effective_count) =
173                    selection.iter().fold((0usize, 0usize), |(rows, count), s| {
174                        if s.row_count > 0 {
175                            (rows + s.row_count, count + 1)
176                        } else {
177                            (rows, count)
178                        }
179                    });
180
181                if effective_count == 0 {
182                    return RowSelectionStrategy::Mask;
183                }
184
185                if total_rows < effective_count.saturating_mul(threshold) {
186                    RowSelectionStrategy::Mask
187                } else {
188                    RowSelectionStrategy::Selectors
189                }
190            }
191        }
192    }
193
194    /// Evaluates an [`ArrowPredicate`], updating this plan's `selection`
195    ///
196    /// If the current `selection` is `Some`, the resulting [`RowSelection`]
197    /// will be the conjunction of the existing selection and the rows selected
198    /// by `predicate`.
199    ///
200    /// Note: pre-existing selections may come from evaluating a previous predicate
201    /// or if the [`ParquetRecordBatchReader`] specified an explicit
202    /// [`RowSelection`] in addition to one or more predicates.
203    pub fn with_predicate(
204        self,
205        array_reader: Box<dyn ArrayReader>,
206        predicate: &mut dyn ArrowPredicate,
207    ) -> Result<Self> {
208        self.with_predicate_options(PredicateOptions::new(array_reader, predicate))
209    }
210
211    /// Evaluates an [`ArrowPredicate`] with the given [`PredicateOptions`],
212    /// updating this plan's `selection`.
213    ///
214    /// Like [`Self::with_predicate`], but allows additional options such as a
215    /// match-count limit for early termination (see
216    /// [`PredicateOptions::with_limit`]).
217    pub fn with_predicate_options(mut self, options: PredicateOptions<'_>) -> Result<Self> {
218        let PredicateOptions {
219            array_reader,
220            predicate,
221            limit,
222            total_rows,
223        } = options;
224
225        // Target length for the concatenated filter output:
226        // - Prior selection ⇒ the reader yields that many rows; `and_then`
227        //   below requires the filter output to match.
228        // - No prior selection ⇒ the reader yields `total_rows`. We only
229        //   need to pad when `limit` may short-circuit the loop; otherwise
230        //   iteration naturally exhausts.
231        let expected_rows = match self.selection.as_ref() {
232            Some(s) => Some(s.row_count()),
233            None => limit.map(|_| total_rows),
234        };
235
236        let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build());
237        let mut filters = vec![];
238        let mut processed_rows: usize = 0;
239        let mut matched_rows: usize = 0;
240        for maybe_batch in reader {
241            let maybe_batch = maybe_batch?;
242            let input_rows = maybe_batch.num_rows();
243            let filter = predicate.evaluate(maybe_batch)?;
244            // Since user supplied predicate, check error here to catch bugs quickly
245            if filter.len() != input_rows {
246                return Err(arrow_err!(
247                    "ArrowPredicate predicate returned {} rows, expected {input_rows}",
248                    filter.len()
249                ));
250            }
251            let filter = match filter.null_count() {
252                0 => filter,
253                // RowSelection::from_filters expects non-null filters. Convert
254                // NULL predicate results to false so they are not selected.
255                _ => prep_null_mask_filter(&filter),
256            };
257
258            processed_rows += input_rows;
259
260            match limit {
261                Some(limit) if limit - matched_rows <= filter.len() => {
262                    let truncated = filter.take_n_true(limit - matched_rows);
263                    matched_rows += truncated.true_count();
264                    filters.push(truncated);
265                    if matched_rows >= limit {
266                        break;
267                    }
268                }
269                _ => {
270                    matched_rows += filter.true_count();
271                    filters.push(filter);
272                }
273            }
274        }
275
276        // Pad the tail so the filters cover `expected_rows` total. This keeps
277        // the invariant that the resulting `RowSelection` spans every row the
278        // reader would have produced — rows past the early break are marked
279        // "not selected". When no limit is set the loop always exhausts and
280        // no padding is needed.
281        if let Some(expected) = expected_rows {
282            if processed_rows < expected {
283                let pad_len = expected - processed_rows;
284                filters.push(BooleanArray::new(BooleanBuffer::new_unset(pad_len), None));
285            }
286        }
287
288        // If the predicate selected all rows and there is no prior selection,
289        // skip creating a RowSelection entirely — this avoids the allocation
290        // and keeps selection as None which enables coalesced page fetches.
291        let all_selected = filters.iter().all(|f| f.true_count() == f.len());
292        if all_selected && self.selection.is_none() {
293            return Ok(self);
294        }
295        let raw = RowSelection::from_filters(&filters);
296        self.selection = match self.selection.take() {
297            Some(selection) => Some(selection.and_then(&raw)),
298            None => Some(raw),
299        };
300        Ok(self)
301    }
302
303    /// Create a final `ReadPlan` the read plan for the scan
304    pub fn build(mut self) -> ReadPlan {
305        // If selection is empty, truncate
306        if !self.selects_any() {
307            self.selection = Some(RowSelection::from(vec![]));
308        }
309
310        // Preferred strategy must not be Auto
311        let selection_strategy = self.resolve_selection_strategy();
312
313        let Self {
314            batch_size,
315            selection,
316            row_selection_policy: _,
317            loaded_row_ranges,
318        } = self;
319
320        let selection = selection.map(|s| s.trim());
321
322        let row_selection_cursor = selection
323            .map(|s| {
324                let selectors: Vec<RowSelector> = s.into();
325                match selection_strategy {
326                    RowSelectionStrategy::Mask => {
327                        RowSelectionCursor::new_mask_from_selectors(selectors, loaded_row_ranges)
328                    }
329                    RowSelectionStrategy::Selectors => RowSelectionCursor::new_selectors(selectors),
330                }
331            })
332            .unwrap_or(RowSelectionCursor::new_all());
333
334        ReadPlan {
335            batch_size,
336            row_selection_cursor,
337        }
338    }
339}
340
341/// Builder for [`ReadPlan`] that applies a limit and offset to the read plan
342///
343/// See [`ReadPlanBuilder::limited`] to create this builder.
344pub(crate) struct LimitedReadPlanBuilder {
345    /// The underlying builder
346    inner: ReadPlanBuilder,
347    /// Total number of rows in the row group before the selection, limit or
348    /// offset are applied
349    row_count: usize,
350    /// The offset to apply, if any
351    offset: Option<usize>,
352    /// The limit to apply, if any
353    limit: Option<usize>,
354}
355
356impl LimitedReadPlanBuilder {
357    /// Create a new `LimitedReadPlanBuilder` from the existing builder and number of rows
358    fn new(inner: ReadPlanBuilder, row_count: usize) -> Self {
359        Self {
360            inner,
361            row_count,
362            offset: None,
363            limit: None,
364        }
365    }
366
367    /// Set the offset to apply to the read plan
368    pub(crate) fn with_offset(mut self, offset: Option<usize>) -> Self {
369        self.offset = offset;
370        self
371    }
372
373    /// Set the limit to apply to the read plan
374    pub(crate) fn with_limit(mut self, limit: Option<usize>) -> Self {
375        self.limit = limit;
376        self
377    }
378
379    /// Apply offset and limit, updating the selection on the underlying builder
380    /// and returning it.
381    pub(crate) fn build_limited(self) -> ReadPlanBuilder {
382        let Self {
383            mut inner,
384            row_count,
385            offset,
386            limit,
387        } = self;
388
389        // If the selection is empty, truncate
390        if !inner.selects_any() {
391            inner.selection = Some(RowSelection::from(vec![]));
392        }
393
394        // If an offset is defined, apply it to the `selection`
395        if let Some(offset) = offset {
396            inner.selection = Some(match row_count.checked_sub(offset) {
397                None => RowSelection::from(vec![]),
398                Some(remaining) => inner
399                    .selection
400                    .map(|selection| selection.offset(offset))
401                    .unwrap_or_else(|| {
402                        RowSelection::from(vec![
403                            RowSelector::skip(offset),
404                            RowSelector::select(remaining),
405                        ])
406                    }),
407            });
408        }
409
410        // If a limit is defined, apply it to the final `selection`
411        if let Some(limit) = limit {
412            inner.selection = Some(
413                inner
414                    .selection
415                    .map(|selection| selection.limit(limit))
416                    .unwrap_or_else(|| {
417                        RowSelection::from(vec![RowSelector::select(limit.min(row_count))])
418                    }),
419            );
420        }
421
422        inner
423    }
424}
425
426/// A plan reading specific rows from a Parquet Row Group.
427///
428/// See [`ReadPlanBuilder`] to create `ReadPlan`s
429#[derive(Debug)]
430pub struct ReadPlan {
431    /// The number of rows to read in each batch
432    batch_size: usize,
433    /// Row ranges to be selected from the data source
434    row_selection_cursor: RowSelectionCursor,
435}
436
437impl ReadPlan {
438    /// Returns a mutable reference to the selection selectors, if any
439    #[deprecated(since = "57.1.0", note = "Use `row_selection_cursor_mut` instead")]
440    pub fn selection_mut(&mut self) -> Option<&mut VecDeque<RowSelector>> {
441        if let RowSelectionCursor::Selectors(selectors_cursor) = &mut self.row_selection_cursor {
442            Some(selectors_cursor.selectors_mut())
443        } else {
444            None
445        }
446    }
447
448    /// Returns a mutable reference to the row selection cursor
449    pub fn row_selection_cursor_mut(&mut self) -> &mut RowSelectionCursor {
450        &mut self.row_selection_cursor
451    }
452
453    /// Return the number of rows to read in each output batch
454    #[inline(always)]
455    pub fn batch_size(&self) -> usize {
456        self.batch_size
457    }
458}
459
460#[cfg(test)]
461mod tests {
462    use super::*;
463
464    fn builder_with_selection(selection: RowSelection) -> ReadPlanBuilder {
465        ReadPlanBuilder::new(1024).with_selection(Some(selection))
466    }
467
468    #[test]
469    fn preferred_selection_strategy_prefers_mask_by_default() {
470        let selection = RowSelection::from(vec![RowSelector::select(8)]);
471        let builder = builder_with_selection(selection);
472        assert_eq!(
473            builder.resolve_selection_strategy(),
474            RowSelectionStrategy::Mask
475        );
476    }
477
478    #[test]
479    fn preferred_selection_strategy_prefers_selectors_when_threshold_small() {
480        let selection = RowSelection::from(vec![RowSelector::select(8)]);
481        let builder = builder_with_selection(selection)
482            .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 1 });
483        assert_eq!(
484            builder.resolve_selection_strategy(),
485            RowSelectionStrategy::Selectors
486        );
487    }
488
489    #[test]
490    fn mask_plan_trims_trailing_skips_before_chunking() {
491        let mut plan = ReadPlanBuilder::new(8)
492            .with_selection(Some(RowSelection::from(vec![
493                RowSelector::select(1),
494                RowSelector::skip(7),
495            ])))
496            .with_row_selection_policy(RowSelectionPolicy::Mask)
497            .build();
498        let RowSelectionCursor::Mask(cursor) = plan.row_selection_cursor_mut() else {
499            panic!("expected a Mask cursor");
500        };
501
502        let chunk = cursor.next_chunk(8).unwrap();
503        assert_eq!(chunk.chunk_rows, 1);
504        assert_eq!(chunk.selected_rows, 1);
505        assert!(cursor.is_empty());
506    }
507
508    #[test]
509    fn with_predicate_options_limit_pads_tail_when_no_prior_selection() {
510        use crate::arrow::ProjectionMask;
511        use crate::arrow::array_reader::StructArrayReader;
512        use crate::arrow::array_reader::test_util::make_int32_page_reader;
513        use crate::arrow::arrow_reader::ArrowPredicateFn;
514        use arrow_schema::{DataType as ArrowType, Field, Fields};
515
516        // 100 rows, all match the predicate. Limit stops the loop after 10
517        // matches — but the resulting RowSelection must still describe the
518        // full 100-row row group (90 trailing rows as "not selected"), not
519        // only the 10 rows we happened to evaluate before breaking.
520        const TOTAL_ROWS: usize = 100;
521        const LIMIT: usize = 10;
522
523        let data: Vec<i32> = (0..TOTAL_ROWS as i32).collect();
524        let levels = vec![0; TOTAL_ROWS];
525        let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None);
526        let struct_type = ArrowType::Struct(Fields::from(vec![Field::new(
527            "c0",
528            ArrowType::Int32,
529            false,
530        )]));
531        let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false, None);
532
533        let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), |batch| {
534            Ok(BooleanArray::from(vec![true; batch.num_rows()]))
535        });
536
537        let builder = ReadPlanBuilder::new(16)
538            .with_predicate_options(
539                PredicateOptions::new(Box::new(struct_reader), &mut predicate)
540                    .with_limit(LIMIT, TOTAL_ROWS),
541            )
542            .unwrap();
543
544        let selection = builder
545            .selection()
546            .expect("limit-driven early break must produce a selection");
547
548        // `row_count` counts selected rows — must equal the limit.
549        assert_eq!(selection.row_count(), LIMIT);
550
551        // Total rows covered (selects + skips) must equal the full row group
552        // so downstream offset/limit math stays in absolute-row space.
553        let total: usize = selection.iter().map(|s| s.row_count).sum();
554        assert_eq!(
555            total, TOTAL_ROWS,
556            "selection must span the full row group, not only the prefix evaluated before the limit"
557        );
558    }
559
560    #[test]
561    fn with_predicate_options_limit_handles_null_filters() {
562        use crate::arrow::ProjectionMask;
563        use crate::arrow::array_reader::StructArrayReader;
564        use crate::arrow::array_reader::test_util::make_int32_page_reader;
565        use crate::arrow::arrow_reader::ArrowPredicateFn;
566        use arrow_schema::{DataType as ArrowType, Field, Fields};
567
568        const TOTAL_ROWS: usize = 100;
569        const LIMIT: usize = 10;
570
571        let data: Vec<i32> = (0..TOTAL_ROWS as i32).collect();
572        let levels = vec![0; TOTAL_ROWS];
573        let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None);
574        let struct_type = ArrowType::Struct(Fields::from(vec![Field::new(
575            "c0",
576            ArrowType::Int32,
577            false,
578        )]));
579        let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false, None);
580
581        let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), |batch| {
582            Ok((0..batch.num_rows())
583                .map(|i| match i % 4 {
584                    0 | 2 => Some(true),
585                    1 => None,
586                    _ => Some(false),
587                })
588                .collect::<BooleanArray>())
589        });
590
591        let builder = ReadPlanBuilder::new(16)
592            .with_predicate_options(
593                PredicateOptions::new(Box::new(struct_reader), &mut predicate)
594                    .with_limit(LIMIT, TOTAL_ROWS),
595            )
596            .unwrap();
597
598        let selection = builder
599            .selection()
600            .expect("limit-driven early break must produce a selection");
601
602        assert_eq!(selection.row_count(), LIMIT);
603
604        let total: usize = selection.iter().map(|s| s.row_count).sum();
605        assert_eq!(total, TOTAL_ROWS);
606    }
607}
parquet/arrow/arrow_reader/read_plan.rs

parquet/arrow/arrow_reader/
read_plan.rs
