parquet/arrow/push_decoder/remaining.rs
1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::DecodeResult;
6use crate::arrow::arrow_reader::{ParquetRecordBatchReader, RowGroupPlan};
7use crate::arrow::push_decoder::reader_builder::{
8 RowGroupBuildResult, RowGroupReaderBuilder, RowGroupReaderBuilderParts,
9};
10use crate::arrow::push_decoder::scan_plan::{NextRowGroup, RowBudget, RowGroupFrontier, ScanPlan};
11use crate::errors::ParquetError;
12use crate::file::metadata::ParquetMetaData;
13use arrow_schema::SchemaRef;
14use bytes::Bytes;
15use std::ops::Range;
16use std::sync::Arc;
17
18/// State machine that tracks the remaining high level chunks (row groups) of
19/// Parquet data left to read.
20///
21/// [`RowGroupFrontier`] owns cross-row-group scan state and selects the next
22/// work item. [`RowGroupReaderBuilder`] owns decoding for the active row group.
23#[derive(Debug)]
24pub(crate) struct RemainingRowGroups {
25 /// The arrow schema of the decoded output. Carried only so
26 /// [`Self::into_parts`] can hand it to a rebuilt builder; unused while
27 /// decoding.
28 schema: SchemaRef,
29
30 /// Cross-row-group scan state for queued work.
31 frontier: RowGroupFrontier,
32
33 /// State for building the reader for the current row group
34 row_group_reader_builder: RowGroupReaderBuilder,
35
36 /// The row group that the reader builder is fetching, for
37 /// [`Self::scan_plan`]. `None` when the builder requests no more bytes.
38 active_row_group: Option<NextRowGroup>,
39}
40
41/// The state recovered from a [`RemainingRowGroups`] by
42/// [`RemainingRowGroups::into_parts`], describing the row groups *not* yet
43/// decoded so a builder reconstructed from it resumes where the decoder left off.
44#[derive(Debug)]
45pub(crate) struct RemainingRowGroupsParts {
46 /// The arrow schema of the decoded output.
47 pub schema: SchemaRef,
48 /// The Parquet file metadata.
49 pub metadata: Arc<ParquetMetaData>,
50 /// Row groups and selections not yet handed to the reader builder.
51 pub row_group_plan: RowGroupPlan,
52 /// Offset still to be skipped before the next readable row group.
53 pub offset: Option<usize>,
54 /// Output rows still permitted across the remaining row groups.
55 pub limit: Option<usize>,
56 /// Builder-configurable parts of the inner row-group reader builder.
57 pub reader_builder: RowGroupReaderBuilderParts,
58}
59
60impl RemainingRowGroups {
61 pub fn new(
62 schema: SchemaRef,
63 parquet_metadata: Arc<ParquetMetaData>,
64 row_group_plan: RowGroupPlan,
65 budget: RowBudget,
66 has_predicates: bool,
67 row_group_reader_builder: RowGroupReaderBuilder,
68 ) -> Result<Self, ParquetError> {
69 Ok(Self {
70 schema,
71 frontier: RowGroupFrontier::new(
72 parquet_metadata,
73 row_group_plan,
74 budget,
75 has_predicates,
76 )?,
77 row_group_reader_builder,
78 active_row_group: None,
79 })
80 }
81
82 /// Decompose into [`RemainingRowGroupsParts`].
83 ///
84 /// Must be called at a row-group boundary (see
85 /// [`Self::is_at_row_group_boundary`]). The inner reader builder's runtime
86 /// decode state is discarded; its buffered bytes are carried through.
87 pub(crate) fn into_parts(self) -> RemainingRowGroupsParts {
88 let Self {
89 schema,
90 frontier,
91 row_group_reader_builder,
92 active_row_group: _,
93 } = self;
94 // `has_predicates` is recomputed by `build()` from the filter.
95 let (parquet_metadata, row_group_plan, budget) = frontier.into_parts();
96 RemainingRowGroupsParts {
97 schema,
98 metadata: parquet_metadata,
99 row_group_plan,
100 offset: budget.offset(),
101 limit: budget.limit(),
102 reader_builder: row_group_reader_builder.into_parts(),
103 }
104 }
105
106 /// See [`super::ScanPlan`]. Plans the active row group in full, then the
107 /// row groups that the frontier has not handed over.
108 pub(super) fn scan_plan(&self) -> ScanPlan {
109 self.row_group_reader_builder
110 .scan_plan_builder(self.frontier.clone())
111 .with_active_row_group(self.active_row_group.clone())
112 .build()
113 }
114
115 /// Push new data buffers that can be used to satisfy pending requests
116 pub fn push_data(
117 &mut self,
118 ranges: Vec<Range<u64>>,
119 buffers: Vec<Bytes>,
120 ) -> Result<(), ParquetError> {
121 self.row_group_reader_builder.push_data(ranges, buffers)
122 }
123
124 /// Return the total number of bytes buffered so far
125 pub fn buffered_bytes(&self) -> u64 {
126 self.row_group_reader_builder.buffered_bytes()
127 }
128
129 /// Clear any staged ranges currently buffered for future decode work
130 pub fn clear_all_ranges(&mut self) {
131 self.row_group_reader_builder.clear_all_ranges();
132 }
133
134 /// True iff the inner row-group reader is between row groups (state
135 /// `Finished`). Forward to [`RowGroupReaderBuilder::is_finished`].
136 pub fn is_at_row_group_boundary(&self) -> bool {
137 self.row_group_reader_builder.is_finished()
138 }
139
140 /// Number of row groups remaining (not including the one currently
141 /// being decoded).
142 pub fn row_groups_remaining(&self) -> usize {
143 self.frontier.row_groups_remaining()
144 }
145
146 /// Peek at the file-level row-group index that the next call to
147 /// [`Self::try_next_reader`] will produce a reader for, after
148 /// simulating the same skip logic [`Self::try_next_reader`] applies
149 /// internally (row-selection emptiness + offset/limit budget). Does
150 /// not mutate state.
151 ///
152 /// Returns `None` when the active row group is still being decoded,
153 /// when no row groups remain, or when every remaining row group
154 /// would be skipped under the current selection/budget.
155 ///
156 /// Cost: one clone of the queued row-group plan, selections and
157 /// row-group occurrence counts per call (the frontier is cloned so the
158 /// real advance logic can run non-destructively). For callers that peek once per row-group boundary
159 /// this is O(remaining row groups + selectors) per boundary.
160 pub fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
161 if self.row_group_reader_builder.has_active_row_group() {
162 return Ok(None);
163 }
164 self.frontier.peek_next_row_group()
165 }
166
167 /// Release the buffered bytes that are outside the read column chunks of
168 /// the queued row groups. The decoder does not read these bytes.
169 pub fn release_unread_bytes(&mut self) {
170 self.row_group_reader_builder
171 .release_unread_bytes(self.frontier.queued_row_groups());
172 }
173
174 /// Release the buffered bytes of a row group that is done, unless the
175 /// queue reads it again. The reader of the row group holds its own
176 /// copies of the bytes that it reads.
177 fn release_row_group(&mut self, row_group_idx: Option<usize>) {
178 if let Some(row_group_idx) = row_group_idx
179 && !self.frontier.is_queued(row_group_idx)
180 {
181 self.row_group_reader_builder
182 .release_row_group(row_group_idx);
183 }
184 }
185
186 /// returns [`ParquetRecordBatchReader`] suitable for reading the next
187 /// group of rows from the Parquet data, or the list of data ranges still
188 /// needed to proceed
189 pub fn try_next_reader(
190 &mut self,
191 ) -> Result<DecodeResult<ParquetRecordBatchReader>, ParquetError> {
192 loop {
193 if !self.row_group_reader_builder.has_active_row_group() {
194 // We are done with the previous row group, seek to the next one
195 // from the frontier, if any.
196
197 match self.frontier.next_readable_row_group()? {
198 Some(next_row_group) => {
199 self.active_row_group = Some(next_row_group.clone());
200 let NextRowGroup {
201 row_group_idx,
202 row_count,
203 selection,
204 budget,
205 } = next_row_group;
206 self.row_group_reader_builder.next_row_group(
207 row_group_idx,
208 row_count,
209 selection,
210 budget,
211 )?;
212 }
213 None => return Ok(DecodeResult::Finished),
214 }
215 }
216
217 let row_group_idx = self.row_group_reader_builder.active_row_group_idx();
218 match self.row_group_reader_builder.try_build()? {
219 RowGroupBuildResult::Finished { remaining_budget } => {
220 self.active_row_group = None;
221 self.frontier
222 .update_budget_after_row_group(remaining_budget);
223 self.release_row_group(row_group_idx);
224 // reader is done, proceed to the next row group
225 }
226 RowGroupBuildResult::NeedsData(ranges) => {
227 // need more data to proceed
228 return Ok(DecodeResult::NeedsData(ranges));
229 }
230 RowGroupBuildResult::Data {
231 batch_reader,
232 remaining_budget,
233 } => {
234 self.active_row_group = None;
235 self.frontier
236 .update_budget_after_row_group(remaining_budget);
237 self.release_row_group(row_group_idx);
238 // ready to read the row group
239 return Ok(DecodeResult::Data(batch_reader));
240 }
241 }
242 }
243 }
244}