parquet/arrow/push_decoder/scan_plan/
frontier.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//! Which row groups a scan reads next, with their selections and budget.
19
20use crate::arrow::arrow_reader::{RowGroupPlan, RowGroupSelection, RowSelection};
21use crate::errors::ParquetError;
22use crate::file::metadata::ParquetMetaData;
23use std::collections::{HashMap, VecDeque};
24use std::sync::Arc;
25
26use super::budget::RowBudget;
27
28/// Plan for the next queued row group after row-selection slicing.
29#[derive(Debug)]
30enum QueuedRowGroupDecision {
31    /// Hand this row group to the builder.
32    Read(NextRowGroup),
33    /// Skip this row group, and keep scanning with the updated budget.
34    Skip { remaining_budget: RowBudget },
35}
36
37/// Work item handed from [`RowGroupFrontier`] to [`RowGroupReaderBuilder`](crate::arrow::push_decoder::reader_builder::RowGroupReaderBuilder).
38#[derive(Debug, Clone)]
39pub(crate) struct NextRowGroup {
40    pub(crate) row_group_idx: usize,
41    pub(crate) row_count: usize,
42    /// This row group's selection, or `None` when all rows are selected.
43    pub(crate) selection: Option<RowSelection>,
44    /// Budget snapshot to apply while decoding this row group.
45    pub(crate) budget: RowBudget,
46}
47
48/// Row groups and selections that have not yet been handed to the row-group
49/// reader builder.
50#[derive(Debug, Clone)]
51struct QueuedRowGroups {
52    queue: Queue,
53    /// The number of times each row group is in `queue`, so that
54    /// [`Self::contains`] does not scan the queue.
55    counts: HashMap<usize, usize>,
56}
57
58impl QueuedRowGroups {
59    /// Validate and queue a row-group plan for `parquet_metadata`.
60    fn try_new(
61        parquet_metadata: &ParquetMetaData,
62        row_group_plan: RowGroupPlan,
63    ) -> Result<Self, ParquetError> {
64        let queue = Queue::try_new(parquet_metadata, row_group_plan)?;
65        let mut counts = HashMap::new();
66        for row_group_idx in queue.row_group_indices() {
67            *counts.entry(row_group_idx).or_default() += 1;
68        }
69        Ok(Self { queue, counts })
70    }
71
72    /// Convert the remaining queue back into a builder configuration.
73    fn into_plan(self) -> RowGroupPlan {
74        self.queue.into_plan()
75    }
76
77    fn front(&self) -> Option<usize> {
78        self.queue.front()
79    }
80
81    /// The row groups in the queue, in no specific order and without
82    /// duplicates.
83    fn row_groups(&self) -> impl Iterator<Item = usize> + '_ {
84        self.counts.keys().copied()
85    }
86
87    /// Returns `true` if `row_group_idx` is in the queue.
88    fn contains(&self, row_group_idx: usize) -> bool {
89        self.counts.contains_key(&row_group_idx)
90    }
91
92    fn len(&self) -> usize {
93        self.queue.len()
94    }
95
96    fn clear(&mut self) {
97        self.queue.clear();
98        self.counts.clear();
99    }
100
101    /// See [`Queue::global_selection_is_exhausted`].
102    fn global_selection_is_exhausted(&self) -> bool {
103        self.queue.global_selection_is_exhausted()
104    }
105
106    /// Remove the front row group and return its local selection.
107    fn pop_front_selection(&mut self, row_count: usize) -> Option<RowSelection> {
108        if let Some(row_group_idx) = self.queue.front() {
109            match self.counts.get_mut(&row_group_idx) {
110                Some(1) => {
111                    self.counts.remove(&row_group_idx);
112                }
113                Some(count) => *count -= 1,
114                None => debug_assert!(false, "row group {row_group_idx} is queued but not counted"),
115            }
116        }
117        self.queue.pop_front_selection(row_count)
118    }
119}
120
121/// The queue of [`QueuedRowGroups`].
122#[derive(Debug, Clone)]
123enum Queue {
124    /// One selection cursor spans all queued row groups.
125    Global {
126        row_groups: VecDeque<usize>,
127        selection: Option<RowSelection>,
128    },
129    /// Selections are already relative to their respective row groups.
130    PerRowGroup(VecDeque<RowGroupSelection>),
131}
132
133impl Queue {
134    /// Validate and queue a row-group plan for `parquet_metadata`.
135    fn try_new(
136        parquet_metadata: &ParquetMetaData,
137        row_group_plan: RowGroupPlan,
138    ) -> Result<Self, ParquetError> {
139        match row_group_plan {
140            RowGroupPlan::Global {
141                row_groups,
142                selection,
143            } => Ok(Self::Global {
144                row_groups: row_groups
145                    .unwrap_or_else(|| (0..parquet_metadata.num_row_groups()).collect())
146                    .into(),
147                selection,
148            }),
149            RowGroupPlan::PerRowGroup(row_groups) => {
150                for row_group in &row_groups {
151                    let row_count =
152                        parquet_metadata.row_group_num_rows(row_group.row_group_index)?;
153                    if let Some(selection) = &row_group.selection {
154                        let selection_rows = selection.total_row_count();
155                        if selection_rows > row_count {
156                            return Err(ParquetError::General(format!(
157                                "Row selection for row group {} contains {selection_rows} rows, but the row group has {row_count}",
158                                row_group.row_group_index
159                            )));
160                        }
161                    }
162                }
163                Ok(Self::PerRowGroup(row_groups.into()))
164            }
165            RowGroupPlan::Conflicting => Err(RowGroupPlan::conflict_error()),
166        }
167    }
168
169    /// Convert the remaining queue back into a builder configuration.
170    fn into_plan(self) -> RowGroupPlan {
171        match self {
172            Self::Global {
173                row_groups,
174                selection,
175            } => RowGroupPlan::Global {
176                row_groups: Some(Vec::from(row_groups)),
177                selection,
178            },
179            Self::PerRowGroup(row_groups) => RowGroupPlan::PerRowGroup(Vec::from(row_groups)),
180        }
181    }
182
183    fn front(&self) -> Option<usize> {
184        match self {
185            Self::Global { row_groups, .. } => row_groups.front().copied(),
186            Self::PerRowGroup(row_groups) => row_groups
187                .front()
188                .map(|row_group| row_group.row_group_index),
189        }
190    }
191
192    /// The queued row group indexes, in order.
193    fn row_group_indices(&self) -> Vec<usize> {
194        match self {
195            Self::Global { row_groups, .. } => row_groups.iter().copied().collect(),
196            Self::PerRowGroup(row_groups) => row_groups
197                .iter()
198                .map(|row_group| row_group.row_group_index)
199                .collect(),
200        }
201    }
202
203    fn len(&self) -> usize {
204        match self {
205            Self::Global { row_groups, .. } => row_groups.len(),
206            Self::PerRowGroup(row_groups) => row_groups.len(),
207        }
208    }
209
210    fn clear(&mut self) {
211        match self {
212            Self::Global {
213                row_groups,
214                selection,
215            } => {
216                row_groups.clear();
217                *selection = None;
218            }
219            Self::PerRowGroup(row_groups) => row_groups.clear(),
220        }
221    }
222
223    /// Returns `true` when a shared global selection has no selected rows left.
224    /// Per-row-group selections are independent and are drained one at a time.
225    fn global_selection_is_exhausted(&self) -> bool {
226        matches!(
227            self,
228            Self::Global {
229                selection: Some(selection),
230                ..
231            } if selection.row_count() == 0
232        )
233    }
234
235    /// Remove the front row group and return its local selection.
236    fn pop_front_selection(&mut self, row_count: usize) -> Option<RowSelection> {
237        match self {
238            Self::Global {
239                row_groups,
240                selection,
241            } => {
242                let popped = row_groups.pop_front();
243                debug_assert!(popped.is_some(), "front row group checked before pop");
244                selection
245                    .as_mut()
246                    .map(|selection| selection.split_off(row_count))
247            }
248            Self::PerRowGroup(row_groups) => {
249                row_groups
250                    .pop_front()
251                    .expect("front row group checked before pop")
252                    .selection
253            }
254        }
255    }
256}
257
258#[derive(Debug, Clone)]
259pub(crate) struct RowGroupFrontier {
260    /// Metadata used to resolve row counts for queued row groups.
261    pub(super) parquet_metadata: Arc<ParquetMetaData>,
262    /// Row groups not yet handed to the builder.
263    queued: QueuedRowGroups,
264    /// Offset/limit budget before the next readable row group is planned.
265    pub(super) budget: RowBudget,
266    /// If predicates are present, row groups with selected rows must be read so
267    /// the predicate can decide whether they are actually needed.
268    pub(super) has_predicates: bool,
269}
270
271impl RowGroupFrontier {
272    pub(crate) fn new(
273        parquet_metadata: Arc<ParquetMetaData>,
274        row_group_plan: RowGroupPlan,
275        budget: RowBudget,
276        has_predicates: bool,
277    ) -> Result<Self, ParquetError> {
278        let queued = QueuedRowGroups::try_new(&parquet_metadata, row_group_plan)?;
279
280        Ok(Self {
281            parquet_metadata,
282            queued,
283            budget,
284            has_predicates,
285        })
286    }
287
288    /// The queued row groups, in no specific order and without duplicates.
289    pub(crate) fn queued_row_groups(&self) -> impl Iterator<Item = usize> + '_ {
290        self.queued.row_groups()
291    }
292
293    /// Returns `true` if `row_group_idx` is still queued to be read.
294    pub(crate) fn is_queued(&self, row_group_idx: usize) -> bool {
295        self.queued.contains(row_group_idx)
296    }
297
298    pub(crate) fn update_budget_after_row_group(&mut self, budget: RowBudget) {
299        self.budget = budget;
300    }
301
302    /// Peek at the next row-group index [`Self::next_readable_row_group`]
303    /// would hand out, without mutating any state. Returns `None` if every
304    /// remaining row group would be skipped under the current
305    /// selection/budget, or if the queue is empty.
306    ///
307    /// Runs the real [`Self::next_readable_row_group`] advance logic on a
308    /// throwaway clone of the frontier, so peek can never drift from the
309    /// read path. The clone copies the queued row-group plan and selections;
310    /// see
311    /// [`RemainingRowGroups::peek_next_row_group`](crate::arrow::push_decoder::remaining::RemainingRowGroups::peek_next_row_group).
312    pub(crate) fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
313        Ok(self
314            .clone()
315            .next_readable_row_group()?
316            .map(|next_row_group| next_row_group.row_group_idx))
317    }
318
319    /// Number of row groups not yet handed to the builder.
320    pub(crate) fn row_groups_remaining(&self) -> usize {
321        self.queued.len()
322    }
323
324    /// Decompose into the metadata, the row groups not yet handed to the
325    /// builder, and the remaining offset/limit budget.
326    pub(crate) fn into_parts(self) -> (Arc<ParquetMetaData>, RowGroupPlan, RowBudget) {
327        let Self {
328            parquet_metadata,
329            queued,
330            budget,
331            has_predicates: _,
332        } = self;
333        (parquet_metadata, queued.into_plan(), budget)
334    }
335
336    fn clear_remaining(&mut self) {
337        self.queued.clear();
338    }
339
340    /// Plan whether a selected row group should be read or skipped.
341    ///
342    /// Selection-only skips are handled before this method is called. This
343    /// method applies the remaining offset/limit budget and predicate
344    /// conservatism.
345    fn plan_selected_row_group(
346        &self,
347        next_row_group: NextRowGroup,
348        selected_rows: usize,
349    ) -> QueuedRowGroupDecision {
350        if self.has_predicates {
351            return QueuedRowGroupDecision::Read(next_row_group);
352        }
353
354        let rows_after_budget = self.budget.rows_after(selected_rows);
355        if rows_after_budget != 0 {
356            return QueuedRowGroupDecision::Read(next_row_group);
357        }
358
359        QueuedRowGroupDecision::Skip {
360            remaining_budget: self.budget.advance(selected_rows, rows_after_budget),
361        }
362    }
363
364    /// Advance queued row groups until one should be handed to the builder.
365    pub(crate) fn next_readable_row_group(&mut self) -> Result<Option<NextRowGroup>, ParquetError> {
366        loop {
367            let Some(row_group_idx) = self.queued.front() else {
368                return Ok(None);
369            };
370            // A global selection can be exhausted before its row-group queue.
371            // Per-row-group selections have no shared cursor to exhaust; empty
372            // local selections are discarded by the `selected_rows == 0` path below.
373            if self.budget.is_exhausted() || self.queued.global_selection_is_exhausted() {
374                self.clear_remaining();
375                return Ok(None);
376            }
377
378            let row_count = self.parquet_metadata.row_group_num_rows(row_group_idx)?;
379            let selection = self.queued.pop_front_selection(row_count);
380            let (selection, selected_rows) = match selection {
381                Some(selection) => {
382                    let selected_rows = selection.row_count();
383                    if selected_rows == 0 {
384                        continue;
385                    }
386                    // An all-rows selection is equivalent to no selection
387                    (
388                        (selected_rows != row_count).then_some(selection),
389                        selected_rows,
390                    )
391                }
392                None => (None, row_count),
393            };
394
395            let next_row_group = NextRowGroup {
396                row_group_idx,
397                row_count,
398                selection,
399                budget: self.budget,
400            };
401
402            match self.plan_selected_row_group(next_row_group, selected_rows) {
403                QueuedRowGroupDecision::Read(next_row_group) => {
404                    return Ok(Some(next_row_group));
405                }
406                QueuedRowGroupDecision::Skip { remaining_budget } => {
407                    self.budget = remaining_budget;
408                }
409            }
410        }
411    }
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417    use crate::arrow::arrow_reader::RowSelector;
418    use crate::arrow::push_decoder::test::test_file_parquet_metadata;
419
420    fn global_plan(
421        row_groups: Option<Vec<usize>>,
422        selection: Option<RowSelection>,
423    ) -> RowGroupPlan {
424        RowGroupPlan::Global {
425            row_groups,
426            selection,
427        }
428    }
429
430    #[test]
431    fn queued_row_groups_encapsulates_plan_transitions() {
432        let metadata = test_file_parquet_metadata();
433
434        let mut all_row_groups =
435            QueuedRowGroups::try_new(&metadata, global_plan(None, None)).unwrap();
436        assert_eq!(all_row_groups.len(), 2);
437        assert_eq!(all_row_groups.front(), Some(0));
438        assert!(!all_row_groups.global_selection_is_exhausted());
439        assert!(all_row_groups.pop_front_selection(200).is_none());
440        assert_eq!(all_row_groups.front(), Some(1));
441        all_row_groups.clear();
442        assert_eq!(all_row_groups.len(), 0);
443        assert!(matches!(
444            all_row_groups.into_plan(),
445            RowGroupPlan::Global {
446                row_groups: Some(row_groups),
447                selection: None,
448            } if row_groups.is_empty()
449        ));
450
451        let global_selection = RowSelection::from(vec![
452            RowSelector::skip(10),
453            RowSelector::select(5),
454            RowSelector::skip(185),
455            RowSelector::select(200),
456        ]);
457        let mut global = QueuedRowGroups::try_new(
458            &metadata,
459            global_plan(Some(vec![0, 1]), Some(global_selection)),
460        )
461        .unwrap();
462        let first = global.pop_front_selection(200).unwrap();
463        assert_eq!(first.row_count(), 5);
464        assert!(!global.global_selection_is_exhausted());
465        assert!(matches!(
466            global.into_plan(),
467            RowGroupPlan::Global {
468                row_groups: Some(row_groups),
469                selection: Some(selection),
470            } if row_groups == vec![1] && selection.row_count() == 200
471        ));
472
473        let local_selection =
474            RowSelection::from(vec![RowSelector::skip(5), RowSelector::select(3)]);
475        let mut local = QueuedRowGroups::try_new(
476            &metadata,
477            RowGroupPlan::PerRowGroup(vec![
478                RowGroupSelection::new(1, Some(local_selection)),
479                RowGroupSelection::new(0, None),
480            ]),
481        )
482        .unwrap();
483        assert_eq!(local.front(), Some(1));
484        assert_eq!(local.pop_front_selection(200).unwrap().row_count(), 3);
485        assert!(!local.global_selection_is_exhausted());
486        assert!(matches!(
487            local.into_plan(),
488            RowGroupPlan::PerRowGroup(row_groups)
489                if row_groups == vec![RowGroupSelection::new(0, None)]
490        ));
491
492        let exhausted = QueuedRowGroups::try_new(
493            &metadata,
494            global_plan(
495                Some(vec![0]),
496                Some(RowSelection::from(vec![RowSelector::skip(200)])),
497            ),
498        )
499        .unwrap();
500        assert!(exhausted.global_selection_is_exhausted());
501    }
502
503    #[test]
504    fn frontier_handles_global_and_local_exhaustion() {
505        let metadata = test_file_parquet_metadata();
506        let budget = RowBudget::new(None, None);
507
508        let mut global = RowGroupFrontier::new(
509            Arc::clone(&metadata),
510            global_plan(
511                Some(vec![0, 1]),
512                Some(RowSelection::from(vec![RowSelector::skip(400)])),
513            ),
514            budget,
515            false,
516        )
517        .unwrap();
518        assert!(global.next_readable_row_group().unwrap().is_none());
519        assert_eq!(global.queued.len(), 0);
520
521        let mut local = RowGroupFrontier::new(
522            Arc::clone(&metadata),
523            RowGroupPlan::PerRowGroup(vec![
524                RowGroupSelection::new(0, Some(RowSelection::from(vec![RowSelector::skip(200)]))),
525                RowGroupSelection::new(1, None),
526            ]),
527            budget,
528            false,
529        )
530        .unwrap();
531        let next = local.next_readable_row_group().unwrap().unwrap();
532        assert_eq!(next.row_group_idx, 1);
533        assert_eq!(next.row_count, 200);
534        assert!(next.selection.is_none());
535
536        let mut exhausted_budget = RowGroupFrontier::new(
537            metadata,
538            RowGroupPlan::PerRowGroup(vec![RowGroupSelection::new(0, None)]),
539            RowBudget::new(None, Some(0)),
540            false,
541        )
542        .unwrap();
543        assert!(
544            exhausted_budget
545                .next_readable_row_group()
546                .unwrap()
547                .is_none()
548        );
549        assert_eq!(exhausted_budget.queued.len(), 0);
550    }
551
552    #[test]
553    fn queued_row_groups_count_repeated_row_groups() {
554        let metadata = test_file_parquet_metadata();
555        let budget = RowBudget::new(None, None);
556
557        // Row group 1 is skipped the first time. A row group stays queued
558        // until its last occurrence is popped.
559        let mut local = RowGroupFrontier::new(
560            Arc::clone(&metadata),
561            RowGroupPlan::PerRowGroup(vec![
562                RowGroupSelection::new(0, None),
563                RowGroupSelection::new(1, Some(RowSelection::from(vec![RowSelector::skip(200)]))),
564                RowGroupSelection::new(0, None),
565                RowGroupSelection::new(1, None),
566            ]),
567            budget,
568            false,
569        )
570        .unwrap();
571        let queued = |frontier: &RowGroupFrontier| {
572            (frontier.queued.contains(0), frontier.queued.contains(1))
573        };
574        assert_eq!(queued(&local), (true, true));
575
576        assert_eq!(
577            local
578                .next_readable_row_group()
579                .unwrap()
580                .unwrap()
581                .row_group_idx,
582            0
583        );
584        assert_eq!(queued(&local), (true, true));
585        // Peeking does not change the counts of the frontier.
586        assert_eq!(local.peek_next_row_group().unwrap(), Some(0));
587        assert_eq!(queued(&local), (true, true));
588
589        // Skips the first occurrence of row group 1.
590        assert_eq!(
591            local
592                .next_readable_row_group()
593                .unwrap()
594                .unwrap()
595                .row_group_idx,
596            0
597        );
598        assert_eq!(queued(&local), (false, true));
599
600        assert_eq!(
601            local
602                .next_readable_row_group()
603                .unwrap()
604                .unwrap()
605                .row_group_idx,
606            1
607        );
608        assert_eq!(queued(&local), (false, false));
609        assert!(local.queued.counts.is_empty());
610        assert!(local.next_readable_row_group().unwrap().is_none());
611
612        // An exhausted limit clears the queue and its counts.
613        let mut global = RowGroupFrontier::new(
614            metadata,
615            global_plan(Some(vec![0, 1, 0]), None),
616            RowBudget::new(None, Some(200)),
617            false,
618        )
619        .unwrap();
620        assert_eq!(
621            global
622                .next_readable_row_group()
623                .unwrap()
624                .unwrap()
625                .row_group_idx,
626            0
627        );
628        assert_eq!(queued(&global), (true, true));
629        global.update_budget_after_row_group(RowBudget::new(None, Some(0)));
630        assert!(global.next_readable_row_group().unwrap().is_none());
631        assert_eq!(queued(&global), (false, false));
632        assert!(global.queued.counts.is_empty());
633    }
634
635    #[test]
636    fn frontier_reports_invalid_global_row_group_while_peeking() {
637        let metadata = test_file_parquet_metadata();
638        let frontier = RowGroupFrontier::new(
639            metadata,
640            global_plan(Some(vec![2]), None),
641            RowBudget::new(None, None),
642            false,
643        )
644        .unwrap();
645
646        let error = frontier.peek_next_row_group().unwrap_err();
647        assert!(
648            error
649                .to_string()
650                .contains("Row group index 2 out of bounds for file with 2 row groups")
651        );
652    }
653
654    #[test]
655    fn metadata_row_count_overflow_is_reported() {
656        let metadata = test_file_parquet_metadata();
657        let mut builder = metadata.as_ref().clone().into_builder();
658        let mut row_groups = builder.take_row_groups();
659        let negative_row_group = row_groups
660            .remove(0)
661            .into_builder()
662            .set_num_rows(-1)
663            .build()
664            .unwrap();
665        let metadata = builder.set_row_groups(vec![negative_row_group]).build();
666
667        let error = metadata.row_group_num_rows(0).unwrap_err();
668        assert!(error.to_string().contains("Row count overflow"));
669    }
670}
parquet/arrow/push_decoder/scan_plan/frontier.rs

parquet/arrow/push_decoder/scan_plan/
frontier.rs
