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