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

parquet/arrow/push_decoder/
remaining.rs
