This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 02accced63 refactor(parquet): share push decoder fetch-range and stage
logic (prep for #10555) (#11330)
02accced63 is described below
commit 02accced63a992a20c91cbb5ac24261863531751
Author: Adrian Garcia Badaracco <[email protected]>
AuthorDate: Mon Oct 5 23:36:28 2026 +0000
refactor(parquet): share push decoder fetch-range and stage logic (prep for
#10555) (#11330)
# Which issue does this PR close?
- Part of https://github.com/apache/arrow-rs/issues/6946.
- Prepares https://github.com/apache/arrow-rs/pull/10555, as requested
in [this
review](https://github.com/apache/arrow-rs/pull/10555#pullrequestreview-5385246575).
# Rationale for this change
https://github.com/apache/arrow-rs/pull/10555 adds
`ParquetPushDecoder::scan_plan()`, which tells an application which byte
ranges the decoder will request. In review, the planner copied two parts
of the decoder's logic:
| Logic | Where the decoder has it |
|---|---|
| The byte ranges to fetch for one column chunk: the dictionary page,
the data pages that the row selection needs, or the full chunk if there
is no offset index | `InMemoryRowGroup::fetch_ranges` |
| The decoding stages of a row group: each `RowFilter` predicate, then
the output projection, with the columns that each stage fetches |
`RowGroupReaderBuilder::try_transition` |
This PR puts each part in one place that both the decoder and the
planner can use. Then https://github.com/apache/arrow-rs/pull/10555 adds
only the planner. It does not have to keep a copy of the decoder logic
in sync.
Terms:
- *Offset index*: optional Parquet metadata that gives the byte location
of each data page.
- *Predicate cache*: decoded predicate columns that the decoder keeps
and uses again for the output. Its columns are fetched with the
selection expanded to batch boundaries.
# What changes are included in this PR?
Three commits. Each commit compiles and passes the tests. There is no
behavior change.
| Commit | Change |
|---|---|
| 1. Move | `RowBudget` and `BudgetedReadPlan` move to
`push_decoder/scan_plan/budget.rs`. `RowGroupFrontier`,
`QueuedRowGroups` and `NextRowGroup` move to
`push_decoder/scan_plan/frontier.rs`. Their tests move with them. Only
visibility changes, plus two accessors that replace direct field access.
Review it with `git show --color-moved`. |
| 2. `ColumnFetch` | `fetch_ranges` is rewritten as a loop over
`ColumnFetch`, which gives the ranges for one column chunk. The output
is the same. |
| 3. `StageSchedule` | One `StageSchedule` gives the projection and the
cache projection of each stage. `try_transition` builds its data
requests from it. |
# Are these changes tested?
Yes, by the existing tests: `push_decoder`, `arrow::` and `async_reader`
(539 tests). The new code has no new behavior to test.
# Are there any user-facing changes?
No. All changes are `pub(crate)` or private.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---
parquet/src/arrow/in_memory_row_group.rs | 218 +++++++---
parquet/src/arrow/push_decoder/mod.rs | 4 +-
.../arrow/push_decoder/reader_builder/filter.rs | 10 +-
.../src/arrow/push_decoder/reader_builder/mod.rs | 181 ++------
.../arrow/push_decoder/reader_builder/stages.rs | 91 ++++
parquet/src/arrow/push_decoder/remaining.rs | 469 +--------------------
parquet/src/arrow/push_decoder/scan_plan/budget.rs | 168 ++++++++
.../{remaining.rs => scan_plan/frontier.rs} | 255 ++---------
parquet/src/arrow/push_decoder/scan_plan/mod.rs | 25 ++
9 files changed, 521 insertions(+), 900 deletions(-)
diff --git a/parquet/src/arrow/in_memory_row_group.rs
b/parquet/src/arrow/in_memory_row_group.rs
index e10a5e2f2a..4a849a0661 100644
--- a/parquet/src/arrow/in_memory_row_group.rs
+++ b/parquet/src/arrow/in_memory_row_group.rs
@@ -22,6 +22,7 @@ use crate::column::page::{PageIterator, PageReader};
use crate::errors::ParquetError;
use crate::file::metadata::page_index::RowGroupPageIndex;
use crate::file::metadata::{ParquetMetaData, RowGroupMetaData};
+use crate::file::page_index::offset_index::PageLocation;
use crate::file::reader::{ChunkReader, Length, SerializedPageReader};
use bytes::{Buf, Bytes};
use std::ops::Range;
@@ -56,6 +57,8 @@ impl InMemoryRowGroup<'_> {
/// [`RowGroupCache`](crate::arrow::array_reader::RowGroupCache).
/// The `selection` for Cached columns is expanded to batch boundaries to
simplify
/// accounting for what data is cached.
+ ///
+ /// The ranges of each column chunk come from [`ColumnFetch`].
pub(crate) fn fetch_ranges(
&self,
projection: &ProjectionMask,
@@ -64,76 +67,43 @@ impl InMemoryRowGroup<'_> {
cache_mask: Option<&ProjectionMask>,
) -> FetchRanges {
let metadata = self.metadata.row_group(self.row_group_idx);
- if let Some((selection, page_index)) =
selection.zip(self.page_index.as_ref()) {
- let expanded_selection =
- selection.expand_to_batch_boundaries(batch_size,
self.row_count);
-
- // If we have a `RowSelection` and an `OffsetIndex` then only fetch
- // pages required for the `RowSelection`
- // Consider preallocating outer vec:
https://github.com/apache/arrow-rs/issues/8667
- let mut page_start_offsets: Vec<Option<Vec<u64>>> = vec![];
-
- let ranges = self
- .column_chunks
- .iter()
- .zip(metadata.columns())
- .enumerate()
- .filter(|&(idx, (chunk, _chunk_meta))| {
- chunk.is_none() && projection.leaf_included(idx)
- })
- .flat_map(|(idx, (_chunk, chunk_meta))| {
- // If the first page does not start at the beginning of
the column,
- // then we need to also fetch a dictionary page.
- let mut ranges: Vec<Range<u64>> = vec![];
- let (start, len) = chunk_meta.byte_range();
- let Some(offset_idx) = page_index.offset_index(idx) else {
- // No offset index for this column, fetch the entire
column
- ranges.push(start..start + len);
- page_start_offsets.push(None);
- return ranges;
- };
-
- match offset_idx.page_locations().first() {
- Some(first) if first.offset as u64 != start => {
- ranges.push(start..first.offset as u64);
- }
- _ => (),
- }
+ // With a `RowSelection` and an `OffsetIndex`, only fetch the pages
+ // required for the `RowSelection`
+ let page_index = self.page_index.as_ref().filter(|_|
selection.is_some());
+ let expanded_selection = selection
+ .filter(|_| page_index.is_some() && cache_mask.is_some())
+ .map(|selection| selection.expand_to_batch_boundaries(batch_size,
self.row_count));
- // Expand selection to batch boundaries if needed for
caching
- // (see doc comment for this function for details on
`cache_mask`)
- let use_expanded = cache_mask.map(|m|
m.leaf_included(idx)).unwrap_or(false);
- if use_expanded {
-
ranges.extend(expanded_selection.scan_ranges(offset_idx.page_locations()));
- } else {
-
ranges.extend(selection.scan_ranges(offset_idx.page_locations()));
+ // Consider preallocating outer vec:
https://github.com/apache/arrow-rs/issues/8667
+ let mut page_start_offsets: Option<Vec<Option<Vec<u64>>>> =
page_index.map(|_| vec![]);
+ let mut ranges = vec![];
+ let columns = columns_to_fetch(projection, self.column_chunks.len(),
|idx| {
+ self.column_chunks[idx].is_some()
+ });
+ for idx in columns {
+ let locations = page_index
+ .and_then(|page_index| page_index.offset_index(idx))
+ .map(|offset_index| offset_index.page_locations().as_slice());
+ let column_selection =
+ column_selection(selection, expanded_selection.as_ref(),
cache_mask, idx);
+ let (start, len) = metadata.column(idx).byte_range();
+ let fetch = ColumnFetch::new(start..start + len, locations,
column_selection);
+ let first = ranges.len();
+ ranges.extend(fetch.ranges());
+ if let Some(page_start_offsets) = page_start_offsets.as_mut() {
+ page_start_offsets.push(match fetch {
+ // No offset index for this column, fetch the entire column
+ ColumnFetch::Chunk { .. } => None,
+ ColumnFetch::Pages { .. } => {
+ Some(ranges[first..].iter().map(|range|
range.start).collect())
}
- page_start_offsets.push(Some(ranges.iter().map(|range|
range.start).collect()));
-
- ranges
- })
- .collect();
- FetchRanges {
- ranges,
- page_start_offsets: Some(page_start_offsets),
- }
- } else {
- let ranges = self
- .column_chunks
- .iter()
- .enumerate()
- .filter(|&(idx, chunk)| chunk.is_none() &&
projection.leaf_included(idx))
- .map(|(idx, _chunk)| {
- let column = metadata.column(idx);
- let (start, length) = column.byte_range();
- start..(start + length)
- })
- .collect();
- FetchRanges {
- ranges,
- page_start_offsets: None,
+ });
}
}
+ FetchRanges {
+ ranges,
+ page_start_offsets,
+ }
}
/// Fills in `self.column_chunks` with the data fetched from `chunk_data`.
@@ -205,6 +175,122 @@ impl InMemoryRowGroup<'_> {
}
}
+/// The leaf columns in `projection` that are not yet read, in column order.
+///
+/// A decoding stage of a row group does not fetch a column that an earlier
+/// stage of the same row group read.
+pub(crate) fn columns_to_fetch<'a>(
+ projection: &'a ProjectionMask,
+ num_columns: usize,
+ is_read: impl Fn(usize) -> bool + 'a,
+) -> impl Iterator<Item = usize> + 'a {
+ (0..num_columns).filter(move |&idx| projection.leaf_included(idx) &&
!is_read(idx))
+}
+
+/// The selection that [`InMemoryRowGroup::fetch_ranges`] uses to choose the
+/// pages of column `idx`: `expanded_selection` (the selection expanded to
+/// batch boundaries) if `cache_mask` includes the column, else `selection`.
+pub(crate) fn column_selection<'a>(
+ selection: Option<&'a RowSelection>,
+ expanded_selection: Option<&'a RowSelection>,
+ cache_mask: Option<&ProjectionMask>,
+ idx: usize,
+) -> Option<&'a RowSelection> {
+ match expanded_selection {
+ Some(expanded) if cache_mask.is_some_and(|mask|
mask.leaf_included(idx)) => Some(expanded),
+ _ => selection,
+ }
+}
+
+/// What [`InMemoryRowGroup::fetch_ranges`] fetches for one column chunk.
+///
+/// This is the one place that decides which bytes of a column chunk are
+/// requested, so other users of this rule cannot differ from the decoder.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) enum ColumnFetch<'a> {
+ /// The whole column chunk, as one range. This is the case if there is
+ /// no selection or the column has no offset index.
+ Chunk {
+ /// Byte range of the column chunk.
+ range: Range<u64>,
+ },
+ /// The dictionary page, if any, and the data pages that contain a
+ /// selected row.
+ Pages {
+ /// Byte range of the dictionary page, if the column chunk has one.
+ dictionary: Option<Range<u64>>,
+ /// The page locations of the column chunk.
+ locations: &'a [PageLocation],
+ /// Indexes in `locations` of the data pages to fetch, in page order.
+ pages: Vec<usize>,
+ },
+}
+
+impl<'a> ColumnFetch<'a> {
+ /// The fetch of the column chunk at byte range `chunk`, with page
+ /// locations `locations` from its offset index, if any, and row
+ /// selection `selection`, if any.
+ pub(crate) fn new(
+ chunk: Range<u64>,
+ locations: Option<&'a [PageLocation]>,
+ selection: Option<&RowSelection>,
+ ) -> Self {
+ let (Some(selection), Some(locations)) = (selection, locations) else {
+ return Self::Chunk { range: chunk };
+ };
+ // `scan_ranges` returns the ranges of the selected pages in page
+ // order. Map them back to page indexes.
+ let fetched = selection.scan_ranges(locations);
+ let mut fetched = fetched.iter().peekable();
+ let pages = locations
+ .iter()
+ .enumerate()
+ .filter(|(_, location)| {
+ fetched
+ .next_if(|range| range.start == location.offset as u64)
+ .is_some()
+ })
+ .map(|(idx, _)| idx)
+ .collect();
+ Self::Pages {
+ dictionary: dictionary_range(chunk.start, locations),
+ locations,
+ pages,
+ }
+ }
+
+ /// The byte ranges to fetch, in file order.
+ pub(crate) fn ranges(&self) -> impl Iterator<Item = Range<u64>> + '_ {
+ let (chunk, dictionary, pages) = match self {
+ Self::Chunk { range, .. } => (Some(range.clone()), None, None),
+ Self::Pages {
+ dictionary,
+ locations,
+ pages,
+ } => (None, dictionary.clone(), Some((*locations, pages))),
+ };
+ let data_pages = pages
+ .into_iter()
+ .flat_map(|(locations, pages)| pages.iter().map(|&page|
page_range(&locations[page])));
+ chunk.into_iter().chain(dictionary).chain(data_pages)
+ }
+}
+
+/// Byte range of the dictionary page of a column chunk that starts at
+/// `chunk_start`: the bytes before the first data page, if any.
+pub(crate) fn dictionary_range(chunk_start: u64, locations: &[PageLocation])
-> Option<Range<u64>> {
+ match locations.first() {
+ Some(first) if first.offset as u64 != chunk_start =>
Some(chunk_start..first.offset as u64),
+ _ => None,
+ }
+}
+
+/// Byte range of a data page.
+pub(crate) fn page_range(location: &PageLocation) -> Range<u64> {
+ let start = location.offset as u64;
+ start..start + location.compressed_page_size as u64
+}
+
impl RowGroups for InMemoryRowGroup<'_> {
fn num_rows(&self) -> usize {
self.row_count
diff --git a/parquet/src/arrow/push_decoder/mod.rs
b/parquet/src/arrow/push_decoder/mod.rs
index 57cb2b9212..1f5054cc45 100644
--- a/parquet/src/arrow/push_decoder/mod.rs
+++ b/parquet/src/arrow/push_decoder/mod.rs
@@ -20,6 +20,7 @@
mod reader_builder;
mod remaining;
+mod scan_plan;
use crate::DecodeResult;
pub use crate::arrow::arrow_reader::RowGroupSelection;
@@ -31,8 +32,9 @@ use crate::file::metadata::ParquetMetaData;
pub use crate::util::push_buffers::PushBuffers;
use arrow_array::RecordBatch;
use bytes::Bytes;
-use reader_builder::{RowBudget, RowGroupReaderBuilder,
RowGroupReaderBuilderParts};
+use reader_builder::{RowGroupReaderBuilder, RowGroupReaderBuilderParts};
use remaining::{RemainingRowGroups, RemainingRowGroupsParts};
+use scan_plan::RowBudget;
use std::ops::Range;
use std::sync::Arc;
diff --git a/parquet/src/arrow/push_decoder/reader_builder/filter.rs
b/parquet/src/arrow/push_decoder/reader_builder/filter.rs
index 219ab06da9..a2ec35077f 100644
--- a/parquet/src/arrow/push_decoder/reader_builder/filter.rs
+++ b/parquet/src/arrow/push_decoder/reader_builder/filter.rs
@@ -116,6 +116,11 @@ impl FilterInfo {
.as_mut()
}
+ /// Index in the [`RowFilter`] of the current predicate
+ pub(super) fn index(&self) -> usize {
+ self.next_predicate.get() - 1
+ }
+
/// Return the current predicate to evaluate
pub(super) fn current(&self) -> &dyn ArrowPredicate {
self.filter
@@ -133,11 +138,6 @@ impl FilterInfo {
self.next_predicate.get() == self.filter.predicates.len()
}
- /// Return a reference to the cache projection
- pub(super) fn cache_projection(&self) -> &ProjectionMask {
- &self.cache_info.cache_projection
- }
-
/// Return a cache builder to save the results of predicate evaluation
pub(super) fn cache_builder(&self) -> CacheOptionsBuilder<'_> {
self.cache_info.builder()
diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs
b/parquet/src/arrow/push_decoder/reader_builder/mod.rs
index 0e6a0e0a66..074f8fc427 100644
--- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs
+++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs
@@ -17,6 +17,7 @@
mod data;
mod filter;
+mod stages;
use crate::arrow::ProjectionMask;
use crate::arrow::array_reader::{ArrayReaderBuilder, CacheOptions,
RowGroupCache};
@@ -29,6 +30,7 @@ use crate::arrow::arrow_reader::{
use crate::arrow::in_memory_row_group::ColumnChunkData;
use crate::arrow::push_decoder::reader_builder::data::DataRequestBuilder;
use crate::arrow::push_decoder::reader_builder::filter::CacheInfo;
+use crate::arrow::push_decoder::scan_plan::{BudgetedReadPlan, RowBudget};
use crate::arrow::schema::ParquetField;
use crate::errors::ParquetError;
use crate::file::metadata::ParquetMetaData;
@@ -38,6 +40,7 @@ use bytes::Bytes;
use data::DataRequest;
use filter::AdvanceResult;
use filter::FilterInfo;
+pub(crate) use stages::{Stage, StageSchedule};
use std::ops::Range;
use std::sync::{Arc, RwLock};
@@ -88,100 +91,6 @@ enum RowGroupDecoderState {
Finished,
}
-/// Running offset/limit budget shared across row groups.
-#[derive(Debug, Clone, Copy, Eq, PartialEq)]
-pub(crate) struct RowBudget {
- offset: Option<usize>,
- limit: Option<usize>,
-}
-
-impl RowBudget {
- pub(crate) fn new(offset: Option<usize>, limit: Option<usize>) -> Self {
- Self { offset, limit }
- }
-
- pub(crate) fn is_exhausted(self) -> bool {
- matches!(self.limit, Some(0))
- }
-
- /// The offset still to be skipped before the next readable row group.
- pub(crate) fn offset(self) -> Option<usize> {
- self.offset
- }
-
- /// The number of output rows still permitted across the remaining row
groups.
- pub(crate) fn limit(self) -> Option<usize> {
- self.limit
- }
-
- /// Returns how many selected rows remain after applying this budget.
- pub(crate) fn rows_after(self, rows_before_budget: usize) -> usize {
- let rows_after_offset =
rows_before_budget.saturating_sub(self.offset.unwrap_or(0));
- match self.limit {
- Some(limit) => rows_after_offset.min(limit),
- None => rows_after_offset,
- }
- }
-
- /// Returns the number of selected rows needed before applying the offset.
- fn selected_row_limit(self) -> Option<usize> {
- self.limit
- .map(|limit| limit.saturating_add(self.offset.unwrap_or(0)))
- }
-
- fn apply_to_plan(self, plan_builder: ReadPlanBuilder, row_count: usize) ->
BudgetedReadPlan {
- let rows_before_budget =
plan_builder.num_rows_selected().unwrap_or(row_count);
- let plan_builder = plan_builder
- .limited(row_count)
- .with_offset(self.offset)
- .with_limit(self.limit)
- .build_limited();
- let rows_after_budget = self.rows_after(rows_before_budget);
-
- BudgetedReadPlan {
- plan_builder,
- rows_before_budget,
- rows_after_budget,
- remaining_budget: self.advance(rows_before_budget,
rows_after_budget),
- }
- }
-
- /// Advance the budget past one row group.
- ///
- /// `rows_before_budget` is the number of rows selected before applying the
- /// budget, and `rows_after_budget` is the number retained for output from
- /// this row group.
- pub(crate) fn advance(mut self, rows_before_budget: usize,
rows_after_budget: usize) -> Self {
- if let Some(offset) = &mut self.offset {
- // Reduction is either because of offset or limit, as limit is
applied
- // after offset has been "exhausted" can just use saturating sub
here.
- *offset = offset.saturating_sub(rows_before_budget -
rows_after_budget);
- }
-
- if rows_after_budget != 0
- && let Some(limit) = &mut self.limit
- {
- *limit -= rows_after_budget;
- }
-
- self
- }
-}
-
-#[derive(Debug)]
-struct BudgetedReadPlan {
- /// Read plan after applying this row group's share of the offset/limit
budget.
- plan_builder: ReadPlanBuilder,
- /// Number of rows selected by row selection and predicates before applying
- /// this row group's offset/limit budget.
- rows_before_budget: usize,
- /// Number of selected rows that remain to be read after applying this row
- /// group's offset/limit budget.
- rows_after_budget: usize,
- /// Budget remaining for later row groups.
- remaining_budget: RowBudget,
-}
-
#[derive(Debug)]
pub(crate) enum RowGroupBuildResult {
/// The active row group is complete without producing a reader.
@@ -271,6 +180,10 @@ pub(crate) struct RowGroupReaderBuilder {
/// The underlying data store
buffers: PushBuffers,
+
+ /// What each decoding stage fetches. Kept here because the filter is
+ /// moved out of the builder while a row group is decoded.
+ stages: StageSchedule,
}
/// The parts of a [`RowGroupReaderBuilder`] needed to rebuild it, recovered by
@@ -306,7 +219,7 @@ impl RowGroupReaderBuilder {
buffers: PushBuffers,
row_selection_policy: RowSelectionPolicy,
) -> Self {
- Self {
+ let mut builder = Self {
batch_size,
projection,
metadata,
@@ -317,7 +230,25 @@ impl RowGroupReaderBuilder {
row_selection_policy,
state: Some(RowGroupDecoderState::Finished),
buffers,
- }
+ stages: StageSchedule::new(ProjectionMask::all(), vec![], None),
+ };
+ let (predicate_projections, cache_projection) = match &builder.filter {
+ Some(filter) => (
+ filter
+ .predicates
+ .iter()
+ .map(|predicate| predicate.projection().clone())
+ .collect(),
+ builder.compute_cache_projection_inner(filter),
+ ),
+ None => (vec![], None),
+ };
+ builder.stages = StageSchedule::new(
+ builder.projection.clone(),
+ predicate_projections,
+ cache_projection,
+ );
+ builder
}
/// Decompose into [`RowGroupReaderBuilderParts`] so the builder can be
@@ -337,6 +268,8 @@ impl RowGroupReaderBuilder {
row_selection_policy,
state: _,
buffers,
+ // Recomputed from `filter` when the builder is rebuilt.
+ stages: _,
} = self;
RowGroupReaderBuilderParts {
batch_size,
@@ -542,7 +475,7 @@ impl RowGroupReaderBuilder {
}
// Make a request for the data needed to evaluate the current
predicate
- let predicate = filter_info.current();
+ let fetch =
self.stages.fetch(Stage::Predicate(filter_info.index()));
// need to fetch pages the column needs for decoding, figure
// that out based on the current selection and projection
@@ -551,13 +484,13 @@ impl RowGroupReaderBuilder {
row_count,
self.batch_size,
&self.metadata,
- predicate.projection(), // use the predicate's projection
+ fetch.projection, // use the predicate's projection
)
.with_selection(plan_builder.selection())
// Cached output columns reuse these predicate-stage chunks.
Expand their
// selection to cache batch boundaries so a cache miss can
safely fetch a
// complete batch from the retained sparse column data.
- .with_cache_projection(Some(filter_info.cache_projection()))
+ .with_cache_projection(fetch.cache_projection)
.with_column_chunks(column_chunks)
.build();
@@ -717,17 +650,17 @@ impl RowGroupReaderBuilder {
));
}
+ let fetch = self.stages.fetch(Stage::Projection);
let data_request = DataRequestBuilder::new(
row_group_idx,
row_count,
self.batch_size,
&self.metadata,
- &self.projection,
+ fetch.projection,
)
.with_selection(plan_builder.selection())
.with_column_chunks(column_chunks)
- // Final projection fetch shouldn't expand selection for cache
- // so don't call with_cache_projection here
+ .with_cache_projection(fetch.cache_projection)
.build();
plan_builder =
plan_builder.with_row_selection_policy(self.row_selection_policy);
@@ -1088,48 +1021,4 @@ mod tests {
assert_eq!(prepared.row_selection_policy(), &RowSelectionPolicy::Mask);
}
-
- #[test]
- fn test_row_budget_offset_limit_across_row_groups() {
- let first =
- RowBudget::new(Some(225),
Some(20)).apply_to_plan(ReadPlanBuilder::new(1024), 200);
- assert_eq!(first.rows_before_budget, 200);
- assert_eq!(first.rows_after_budget, 0);
- assert_eq!(first.remaining_budget, RowBudget::new(Some(25), Some(20)));
- assert_eq!(first.plan_builder.num_rows_selected(), Some(0));
-
- let second = first
- .remaining_budget
- .apply_to_plan(ReadPlanBuilder::new(1024), 200);
- assert_eq!(second.rows_before_budget, 200);
- assert_eq!(second.rows_after_budget, 20);
- assert_eq!(second.remaining_budget, RowBudget::new(Some(0), Some(0)));
- assert_eq!(second.plan_builder.num_rows_selected(), Some(20));
- }
-
- #[test]
- fn test_row_budget_limit_only() {
- let budgeted =
- RowBudget::new(None,
Some(20)).apply_to_plan(ReadPlanBuilder::new(1024), 200);
- assert_eq!(budgeted.rows_before_budget, 200);
- assert_eq!(budgeted.rows_after_budget, 20);
- assert_eq!(budgeted.remaining_budget, RowBudget::new(None, Some(0)));
- assert_eq!(budgeted.plan_builder.num_rows_selected(), Some(20));
- }
-
- #[test]
- fn test_row_budget_empty_selection() {
- let empty_selection = RowSelection::from(vec![RowSelector::skip(200)]);
- let budgeted = RowBudget::new(Some(10), Some(20)).apply_to_plan(
- ReadPlanBuilder::new(1024).with_selection(Some(empty_selection)),
- 200,
- );
- assert_eq!(budgeted.rows_before_budget, 0);
- assert_eq!(budgeted.rows_after_budget, 0);
- assert_eq!(
- budgeted.remaining_budget,
- RowBudget::new(Some(10), Some(20))
- );
- assert_eq!(budgeted.plan_builder.num_rows_selected(), Some(0));
- }
}
diff --git a/parquet/src/arrow/push_decoder/reader_builder/stages.rs
b/parquet/src/arrow/push_decoder/reader_builder/stages.rs
new file mode 100644
index 0000000000..018adf5760
--- /dev/null
+++ b/parquet/src/arrow/push_decoder/reader_builder/stages.rs
@@ -0,0 +1,91 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`StageSchedule`]: what each decoding stage of a row group fetches.
+
+use crate::arrow::ProjectionMask;
+
+/// A decoding stage of a row group.
+///
+/// The decoder evaluates the [`RowFilter`] predicates in order, then decodes
+/// the output columns.
+///
+/// [`RowFilter`]: crate::arrow::arrow_reader::RowFilter
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
+pub(crate) enum Stage {
+ /// Evaluation of the predicate at this index in the `RowFilter`.
+ Predicate(usize),
+ /// Decoding of the output columns.
+ Projection,
+}
+
+/// The columns that each decoding stage of a row group fetches.
+///
+/// A stage does not fetch a column that an earlier stage of the same row
+/// group read (see [`columns_to_fetch`]).
+///
+/// [`columns_to_fetch`]: crate::arrow::in_memory_row_group::columns_to_fetch
+#[derive(Debug)]
+pub(crate) struct StageSchedule {
+ /// Columns in the output.
+ projection: ProjectionMask,
+ /// Columns each predicate reads, in evaluation order.
+ predicate_projections: Vec<ProjectionMask>,
+ /// Predicate columns whose decoded values are cached for the output, if
+ /// any.
+ cache_projection: Option<ProjectionMask>,
+}
+
+/// What one decoding stage fetches.
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct StageFetch<'a> {
+ /// Columns the stage decodes.
+ pub(crate) projection: &'a ProjectionMask,
+ /// Columns fetched with the selection expanded to batch boundaries, so
+ /// that the predicate cache can serve complete batches. Only predicate
+ /// stages have one.
+ pub(crate) cache_projection: Option<&'a ProjectionMask>,
+}
+
+impl StageSchedule {
+ pub(crate) fn new(
+ projection: ProjectionMask,
+ predicate_projections: Vec<ProjectionMask>,
+ cache_projection: Option<ProjectionMask>,
+ ) -> Self {
+ Self {
+ projection,
+ predicate_projections,
+ cache_projection,
+ }
+ }
+
+ /// What `stage` fetches.
+ pub(crate) fn fetch(&self, stage: Stage) -> StageFetch<'_> {
+ match stage {
+ Stage::Predicate(idx) => StageFetch {
+ projection: &self.predicate_projections[idx],
+ cache_projection: self.cache_projection.as_ref(),
+ },
+ // The final projection fetch does not expand the selection.
+ Stage::Projection => StageFetch {
+ projection: &self.projection,
+ cache_projection: None,
+ },
+ }
+ }
+}
diff --git a/parquet/src/arrow/push_decoder/remaining.rs
b/parquet/src/arrow/push_decoder/remaining.rs
index 6d53045bf9..77f0e2f36f 100644
--- a/parquet/src/arrow/push_decoder/remaining.rs
+++ b/parquet/src/arrow/push_decoder/remaining.rs
@@ -16,296 +16,18 @@
// under the License.
use crate::DecodeResult;
-use crate::arrow::arrow_reader::{
- ParquetRecordBatchReader, RowGroupPlan, RowGroupSelection, RowSelection,
-};
+use crate::arrow::arrow_reader::{ParquetRecordBatchReader, RowGroupPlan};
use crate::arrow::push_decoder::reader_builder::{
- RowBudget, RowGroupBuildResult, RowGroupReaderBuilder,
RowGroupReaderBuilderParts,
+ RowGroupBuildResult, RowGroupReaderBuilder, RowGroupReaderBuilderParts,
};
+use crate::arrow::push_decoder::scan_plan::{NextRowGroup, RowBudget,
RowGroupFrontier};
use crate::errors::ParquetError;
use crate::file::metadata::ParquetMetaData;
use arrow_schema::SchemaRef;
use bytes::Bytes;
-use std::collections::VecDeque;
use std::ops::Range;
use std::sync::Arc;
-/// Plan for the next queued row group after row-selection slicing.
-#[derive(Debug)]
-enum QueuedRowGroupDecision {
- /// Hand this row group to the builder.
- Read(NextRowGroup),
- /// Skip this row group, and keep scanning with the updated budget.
- Skip { remaining_budget: RowBudget },
-}
-
-/// Work item handed from [`RowGroupFrontier`] to [`RowGroupReaderBuilder`].
-#[derive(Debug)]
-struct NextRowGroup {
- row_group_idx: usize,
- row_count: usize,
- /// This row group's selection, or `None` when all rows are selected.
- selection: Option<RowSelection>,
- /// Budget snapshot to apply while decoding this row group.
- budget: RowBudget,
-}
-
-/// Row groups and selections that have not yet been handed to the row-group
-/// reader builder.
-#[derive(Debug, Clone)]
-enum QueuedRowGroups {
- /// One selection cursor spans all queued row groups.
- Global {
- row_groups: VecDeque<usize>,
- selection: Option<RowSelection>,
- },
- /// Selections are already relative to their respective row groups.
- PerRowGroup(VecDeque<RowGroupSelection>),
-}
-
-impl QueuedRowGroups {
- /// Validate and queue a row-group plan for `parquet_metadata`.
- fn try_new(
- parquet_metadata: &ParquetMetaData,
- row_group_plan: RowGroupPlan,
- ) -> Result<Self, ParquetError> {
- match row_group_plan {
- RowGroupPlan::Global {
- row_groups,
- selection,
- } => Ok(Self::Global {
- row_groups: row_groups
- .unwrap_or_else(||
(0..parquet_metadata.num_row_groups()).collect())
- .into(),
- selection,
- }),
- RowGroupPlan::PerRowGroup(row_groups) => {
- for row_group in &row_groups {
- let row_count =
-
parquet_metadata.row_group_num_rows(row_group.row_group_index)?;
- if let Some(selection) = &row_group.selection {
- let selection_rows = selection.total_row_count();
- if selection_rows > row_count {
- return Err(ParquetError::General(format!(
- "Row selection for row group {} contains
{selection_rows} rows, but the row group has {row_count}",
- row_group.row_group_index
- )));
- }
- }
- }
- Ok(Self::PerRowGroup(row_groups.into()))
- }
- RowGroupPlan::Conflicting => Err(RowGroupPlan::conflict_error()),
- }
- }
-
- /// Convert the remaining queue back into a builder configuration.
- fn into_plan(self) -> RowGroupPlan {
- match self {
- Self::Global {
- row_groups,
- selection,
- } => RowGroupPlan::Global {
- row_groups: Some(Vec::from(row_groups)),
- selection,
- },
- Self::PerRowGroup(row_groups) =>
RowGroupPlan::PerRowGroup(Vec::from(row_groups)),
- }
- }
-
- fn front(&self) -> Option<usize> {
- match self {
- Self::Global { row_groups, .. } => row_groups.front().copied(),
- Self::PerRowGroup(row_groups) => row_groups
- .front()
- .map(|row_group| row_group.row_group_index),
- }
- }
-
- fn len(&self) -> usize {
- match self {
- Self::Global { row_groups, .. } => row_groups.len(),
- Self::PerRowGroup(row_groups) => row_groups.len(),
- }
- }
-
- fn clear(&mut self) {
- match self {
- Self::Global {
- row_groups,
- selection,
- } => {
- row_groups.clear();
- *selection = None;
- }
- Self::PerRowGroup(row_groups) => row_groups.clear(),
- }
- }
-
- /// Returns `true` when a shared global selection has no selected rows
left.
- /// Per-row-group selections are independent and are drained one at a time.
- fn global_selection_is_exhausted(&self) -> bool {
- matches!(
- self,
- Self::Global {
- selection: Some(selection),
- ..
- } if selection.row_count() == 0
- )
- }
-
- /// Remove the front row group and return its local selection.
- fn pop_front_selection(&mut self, row_count: usize) ->
Option<RowSelection> {
- match self {
- Self::Global {
- row_groups,
- selection,
- } => {
- let popped = row_groups.pop_front();
- debug_assert!(popped.is_some(), "front row group checked
before pop");
- selection
- .as_mut()
- .map(|selection| selection.split_off(row_count))
- }
- Self::PerRowGroup(row_groups) => {
- row_groups
- .pop_front()
- .expect("front row group checked before pop")
- .selection
- }
- }
- }
-}
-
-#[derive(Debug, Clone)]
-struct RowGroupFrontier {
- /// Metadata used to resolve row counts for queued row groups.
- parquet_metadata: Arc<ParquetMetaData>,
- /// Row groups not yet handed to the builder.
- queued: QueuedRowGroups,
- /// Offset/limit budget before the next readable row group is planned.
- budget: RowBudget,
- /// If predicates are present, row groups with selected rows must be read
so
- /// the predicate can decide whether they are actually needed.
- has_predicates: bool,
-}
-
-impl RowGroupFrontier {
- fn new(
- parquet_metadata: Arc<ParquetMetaData>,
- row_group_plan: RowGroupPlan,
- budget: RowBudget,
- has_predicates: bool,
- ) -> Result<Self, ParquetError> {
- let queued = QueuedRowGroups::try_new(&parquet_metadata,
row_group_plan)?;
-
- Ok(Self {
- parquet_metadata,
- queued,
- budget,
- has_predicates,
- })
- }
-
- fn update_budget_after_row_group(&mut self, budget: RowBudget) {
- self.budget = budget;
- }
-
- /// Peek at the next row-group index [`Self::next_readable_row_group`]
- /// would hand out, without mutating any state. Returns `None` if every
- /// remaining row group would be skipped under the current
- /// selection/budget, or if the queue is empty.
- ///
- /// Runs the real [`Self::next_readable_row_group`] advance logic on a
- /// throwaway clone of the frontier, so peek can never drift from the
- /// read path. The clone copies the queued row-group plan and selections;
- /// see
- /// [`RemainingRowGroups::peek_next_row_group`].
- fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
- Ok(self
- .clone()
- .next_readable_row_group()?
- .map(|next_row_group| next_row_group.row_group_idx))
- }
-
- fn clear_remaining(&mut self) {
- self.queued.clear();
- }
-
- /// Plan whether a selected row group should be read or skipped.
- ///
- /// Selection-only skips are handled before this method is called. This
- /// method applies the remaining offset/limit budget and predicate
- /// conservatism.
- fn plan_selected_row_group(
- &self,
- next_row_group: NextRowGroup,
- selected_rows: usize,
- ) -> QueuedRowGroupDecision {
- if self.has_predicates {
- return QueuedRowGroupDecision::Read(next_row_group);
- }
-
- let rows_after_budget = self.budget.rows_after(selected_rows);
- if rows_after_budget != 0 {
- return QueuedRowGroupDecision::Read(next_row_group);
- }
-
- QueuedRowGroupDecision::Skip {
- remaining_budget: self.budget.advance(selected_rows,
rows_after_budget),
- }
- }
-
- /// Advance queued row groups until one should be handed to the builder.
- fn next_readable_row_group(&mut self) -> Result<Option<NextRowGroup>,
ParquetError> {
- loop {
- let Some(row_group_idx) = self.queued.front() else {
- return Ok(None);
- };
- // A global selection can be exhausted before its row-group queue.
- // Per-row-group selections have no shared cursor to exhaust; empty
- // local selections are discarded by the `selected_rows == 0` path
below.
- if self.budget.is_exhausted() ||
self.queued.global_selection_is_exhausted() {
- self.clear_remaining();
- return Ok(None);
- }
-
- let row_count =
self.parquet_metadata.row_group_num_rows(row_group_idx)?;
- let selection = self.queued.pop_front_selection(row_count);
- let (selection, selected_rows) = match selection {
- Some(selection) => {
- let selected_rows = selection.row_count();
- if selected_rows == 0 {
- continue;
- }
- // An all-rows selection is equivalent to no selection
- (
- (selected_rows != row_count).then_some(selection),
- selected_rows,
- )
- }
- None => (None, row_count),
- };
-
- let next_row_group = NextRowGroup {
- row_group_idx,
- row_count,
- selection,
- budget: self.budget,
- };
-
- match self.plan_selected_row_group(next_row_group, selected_rows) {
- QueuedRowGroupDecision::Read(next_row_group) => {
- return Ok(Some(next_row_group));
- }
- QueuedRowGroupDecision::Skip { remaining_budget } => {
- self.budget = remaining_budget;
- }
- }
- }
- }
-}
-
/// State machine that tracks the remaining high level chunks (row groups) of
/// Parquet data left to read.
///
@@ -377,13 +99,7 @@ impl RemainingRowGroups {
row_group_reader_builder,
} = self;
// `has_predicates` is recomputed by `build()` from the filter.
- let RowGroupFrontier {
- parquet_metadata,
- queued,
- budget,
- has_predicates: _,
- } = frontier;
- let row_group_plan = queued.into_plan();
+ let (parquet_metadata, row_group_plan, budget) = frontier.into_parts();
RemainingRowGroupsParts {
schema,
metadata: parquet_metadata,
@@ -422,7 +138,7 @@ impl RemainingRowGroups {
/// Number of row groups remaining (not including the one currently
/// being decoded).
pub fn row_groups_remaining(&self) -> usize {
- self.frontier.queued.len()
+ self.frontier.row_groups_remaining()
}
/// Peek at the file-level row-group index that the next call to
@@ -498,178 +214,3 @@ impl RemainingRowGroups {
}
}
}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::arrow::arrow_reader::RowSelector;
- use crate::arrow::push_decoder::test::test_file_parquet_metadata;
-
- fn global_plan(
- row_groups: Option<Vec<usize>>,
- selection: Option<RowSelection>,
- ) -> RowGroupPlan {
- RowGroupPlan::Global {
- row_groups,
- selection,
- }
- }
-
- #[test]
- fn queued_row_groups_encapsulates_plan_transitions() {
- let metadata = test_file_parquet_metadata();
-
- let mut all_row_groups =
- QueuedRowGroups::try_new(&metadata, global_plan(None,
None)).unwrap();
- assert_eq!(all_row_groups.len(), 2);
- assert_eq!(all_row_groups.front(), Some(0));
- assert!(!all_row_groups.global_selection_is_exhausted());
- assert!(all_row_groups.pop_front_selection(200).is_none());
- assert_eq!(all_row_groups.front(), Some(1));
- all_row_groups.clear();
- assert_eq!(all_row_groups.len(), 0);
- assert!(matches!(
- all_row_groups.into_plan(),
- RowGroupPlan::Global {
- row_groups: Some(row_groups),
- selection: None,
- } if row_groups.is_empty()
- ));
-
- let global_selection = RowSelection::from(vec![
- RowSelector::skip(10),
- RowSelector::select(5),
- RowSelector::skip(185),
- RowSelector::select(200),
- ]);
- let mut global = QueuedRowGroups::try_new(
- &metadata,
- global_plan(Some(vec![0, 1]), Some(global_selection)),
- )
- .unwrap();
- let first = global.pop_front_selection(200).unwrap();
- assert_eq!(first.row_count(), 5);
- assert!(!global.global_selection_is_exhausted());
- assert!(matches!(
- global.into_plan(),
- RowGroupPlan::Global {
- row_groups: Some(row_groups),
- selection: Some(selection),
- } if row_groups == vec![1] && selection.row_count() == 200
- ));
-
- let local_selection =
- RowSelection::from(vec![RowSelector::skip(5),
RowSelector::select(3)]);
- let mut local = QueuedRowGroups::try_new(
- &metadata,
- RowGroupPlan::PerRowGroup(vec![
- RowGroupSelection::new(1, Some(local_selection)),
- RowGroupSelection::new(0, None),
- ]),
- )
- .unwrap();
- assert_eq!(local.front(), Some(1));
- assert_eq!(local.pop_front_selection(200).unwrap().row_count(), 3);
- assert!(!local.global_selection_is_exhausted());
- assert!(matches!(
- local.into_plan(),
- RowGroupPlan::PerRowGroup(row_groups)
- if row_groups == vec![RowGroupSelection::new(0, None)]
- ));
-
- let exhausted = QueuedRowGroups::try_new(
- &metadata,
- global_plan(
- Some(vec![0]),
- Some(RowSelection::from(vec![RowSelector::skip(200)])),
- ),
- )
- .unwrap();
- assert!(exhausted.global_selection_is_exhausted());
- }
-
- #[test]
- fn frontier_handles_global_and_local_exhaustion() {
- let metadata = test_file_parquet_metadata();
- let budget = RowBudget::new(None, None);
-
- let mut global = RowGroupFrontier::new(
- Arc::clone(&metadata),
- global_plan(
- Some(vec![0, 1]),
- Some(RowSelection::from(vec![RowSelector::skip(400)])),
- ),
- budget,
- false,
- )
- .unwrap();
- assert!(global.next_readable_row_group().unwrap().is_none());
- assert_eq!(global.queued.len(), 0);
-
- let mut local = RowGroupFrontier::new(
- Arc::clone(&metadata),
- RowGroupPlan::PerRowGroup(vec![
- RowGroupSelection::new(0,
Some(RowSelection::from(vec![RowSelector::skip(200)]))),
- RowGroupSelection::new(1, None),
- ]),
- budget,
- false,
- )
- .unwrap();
- let next = local.next_readable_row_group().unwrap().unwrap();
- assert_eq!(next.row_group_idx, 1);
- assert_eq!(next.row_count, 200);
- assert!(next.selection.is_none());
-
- let mut exhausted_budget = RowGroupFrontier::new(
- metadata,
- RowGroupPlan::PerRowGroup(vec![RowGroupSelection::new(0, None)]),
- RowBudget::new(None, Some(0)),
- false,
- )
- .unwrap();
- assert!(
- exhausted_budget
- .next_readable_row_group()
- .unwrap()
- .is_none()
- );
- assert_eq!(exhausted_budget.queued.len(), 0);
- }
-
- #[test]
- fn frontier_reports_invalid_global_row_group_while_peeking() {
- let metadata = test_file_parquet_metadata();
- let frontier = RowGroupFrontier::new(
- metadata,
- global_plan(Some(vec![2]), None),
- RowBudget::new(None, None),
- false,
- )
- .unwrap();
-
- let error = frontier.peek_next_row_group().unwrap_err();
- assert!(
- error
- .to_string()
- .contains("Row group index 2 out of bounds for file with 2 row
groups")
- );
- }
-
- #[test]
- fn metadata_row_count_overflow_is_reported() {
- let metadata = test_file_parquet_metadata();
- let mut builder = metadata.as_ref().clone().into_builder();
- let mut row_groups = builder.take_row_groups();
- let negative_row_group = row_groups
- .remove(0)
- .into_builder()
- .set_num_rows(-1)
- .build()
- .unwrap();
- let metadata =
builder.set_row_groups(vec![negative_row_group]).build();
-
- let error = metadata.row_group_num_rows(0).unwrap_err();
- assert!(error.to_string().contains("Row count overflow"));
- }
-}
diff --git a/parquet/src/arrow/push_decoder/scan_plan/budget.rs
b/parquet/src/arrow/push_decoder/scan_plan/budget.rs
new file mode 100644
index 0000000000..303ce52a30
--- /dev/null
+++ b/parquet/src/arrow/push_decoder/scan_plan/budget.rs
@@ -0,0 +1,168 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Offset/limit budget shared across the row groups of a scan.
+
+use crate::arrow::arrow_reader::ReadPlanBuilder;
+
+/// Running offset/limit budget shared across row groups.
+#[derive(Debug, Clone, Copy, Eq, PartialEq)]
+pub(crate) struct RowBudget {
+ offset: Option<usize>,
+ limit: Option<usize>,
+}
+
+impl RowBudget {
+ pub(crate) fn new(offset: Option<usize>, limit: Option<usize>) -> Self {
+ Self { offset, limit }
+ }
+
+ pub(crate) fn is_exhausted(self) -> bool {
+ matches!(self.limit, Some(0))
+ }
+
+ /// The offset still to be skipped before the next readable row group.
+ pub(crate) fn offset(self) -> Option<usize> {
+ self.offset
+ }
+
+ /// The number of output rows still permitted across the remaining row
groups.
+ pub(crate) fn limit(self) -> Option<usize> {
+ self.limit
+ }
+
+ /// Returns how many selected rows remain after applying this budget.
+ pub(crate) fn rows_after(self, rows_before_budget: usize) -> usize {
+ let rows_after_offset =
rows_before_budget.saturating_sub(self.offset.unwrap_or(0));
+ match self.limit {
+ Some(limit) => rows_after_offset.min(limit),
+ None => rows_after_offset,
+ }
+ }
+
+ /// Returns the number of selected rows needed before applying the offset.
+ pub(crate) fn selected_row_limit(self) -> Option<usize> {
+ self.limit
+ .map(|limit| limit.saturating_add(self.offset.unwrap_or(0)))
+ }
+
+ pub(crate) fn apply_to_plan(
+ self,
+ plan_builder: ReadPlanBuilder,
+ row_count: usize,
+ ) -> BudgetedReadPlan {
+ let rows_before_budget =
plan_builder.num_rows_selected().unwrap_or(row_count);
+ let plan_builder = plan_builder
+ .limited(row_count)
+ .with_offset(self.offset)
+ .with_limit(self.limit)
+ .build_limited();
+ let rows_after_budget = self.rows_after(rows_before_budget);
+
+ BudgetedReadPlan {
+ plan_builder,
+ rows_before_budget,
+ rows_after_budget,
+ remaining_budget: self.advance(rows_before_budget,
rows_after_budget),
+ }
+ }
+
+ /// Advance the budget past one row group.
+ ///
+ /// `rows_before_budget` is the number of rows selected before applying the
+ /// budget, and `rows_after_budget` is the number retained for output from
+ /// this row group.
+ pub(crate) fn advance(mut self, rows_before_budget: usize,
rows_after_budget: usize) -> Self {
+ if let Some(offset) = &mut self.offset {
+ // Reduction is either because of offset or limit, as limit is
applied
+ // after offset has been "exhausted" can just use saturating sub
here.
+ *offset = offset.saturating_sub(rows_before_budget -
rows_after_budget);
+ }
+
+ if rows_after_budget != 0
+ && let Some(limit) = &mut self.limit
+ {
+ *limit -= rows_after_budget;
+ }
+
+ self
+ }
+}
+
+#[derive(Debug)]
+pub(crate) struct BudgetedReadPlan {
+ /// Read plan after applying this row group's share of the offset/limit
budget.
+ pub(crate) plan_builder: ReadPlanBuilder,
+ /// Number of rows selected by row selection and predicates before applying
+ /// this row group's offset/limit budget.
+ pub(crate) rows_before_budget: usize,
+ /// Number of selected rows that remain to be read after applying this row
+ /// group's offset/limit budget.
+ pub(crate) rows_after_budget: usize,
+ /// Budget remaining for later row groups.
+ pub(crate) remaining_budget: RowBudget,
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::arrow::arrow_reader::{RowSelection, RowSelector};
+
+ #[test]
+ fn test_row_budget_offset_limit_across_row_groups() {
+ let first =
+ RowBudget::new(Some(225),
Some(20)).apply_to_plan(ReadPlanBuilder::new(1024), 200);
+ assert_eq!(first.rows_before_budget, 200);
+ assert_eq!(first.rows_after_budget, 0);
+ assert_eq!(first.remaining_budget, RowBudget::new(Some(25), Some(20)));
+ assert_eq!(first.plan_builder.num_rows_selected(), Some(0));
+
+ let second = first
+ .remaining_budget
+ .apply_to_plan(ReadPlanBuilder::new(1024), 200);
+ assert_eq!(second.rows_before_budget, 200);
+ assert_eq!(second.rows_after_budget, 20);
+ assert_eq!(second.remaining_budget, RowBudget::new(Some(0), Some(0)));
+ assert_eq!(second.plan_builder.num_rows_selected(), Some(20));
+ }
+
+ #[test]
+ fn test_row_budget_limit_only() {
+ let budgeted =
+ RowBudget::new(None,
Some(20)).apply_to_plan(ReadPlanBuilder::new(1024), 200);
+ assert_eq!(budgeted.rows_before_budget, 200);
+ assert_eq!(budgeted.rows_after_budget, 20);
+ assert_eq!(budgeted.remaining_budget, RowBudget::new(None, Some(0)));
+ assert_eq!(budgeted.plan_builder.num_rows_selected(), Some(20));
+ }
+
+ #[test]
+ fn test_row_budget_empty_selection() {
+ let empty_selection = RowSelection::from(vec![RowSelector::skip(200)]);
+ let budgeted = RowBudget::new(Some(10), Some(20)).apply_to_plan(
+ ReadPlanBuilder::new(1024).with_selection(Some(empty_selection)),
+ 200,
+ );
+ assert_eq!(budgeted.rows_before_budget, 0);
+ assert_eq!(budgeted.rows_after_budget, 0);
+ assert_eq!(
+ budgeted.remaining_budget,
+ RowBudget::new(Some(10), Some(20))
+ );
+ assert_eq!(budgeted.plan_builder.num_rows_selected(), Some(0));
+ }
+}
diff --git a/parquet/src/arrow/push_decoder/remaining.rs
b/parquet/src/arrow/push_decoder/scan_plan/frontier.rs
similarity index 66%
copy from parquet/src/arrow/push_decoder/remaining.rs
copy to parquet/src/arrow/push_decoder/scan_plan/frontier.rs
index 6d53045bf9..c35d91c1d9 100644
--- a/parquet/src/arrow/push_decoder/remaining.rs
+++ b/parquet/src/arrow/push_decoder/scan_plan/frontier.rs
@@ -15,21 +15,16 @@
// specific language governing permissions and limitations
// under the License.
-use crate::DecodeResult;
-use crate::arrow::arrow_reader::{
- ParquetRecordBatchReader, RowGroupPlan, RowGroupSelection, RowSelection,
-};
-use crate::arrow::push_decoder::reader_builder::{
- RowBudget, RowGroupBuildResult, RowGroupReaderBuilder,
RowGroupReaderBuilderParts,
-};
+//! Which row groups a scan reads next, with their selections and budget.
+
+use crate::arrow::arrow_reader::{RowGroupPlan, RowGroupSelection,
RowSelection};
use crate::errors::ParquetError;
use crate::file::metadata::ParquetMetaData;
-use arrow_schema::SchemaRef;
-use bytes::Bytes;
use std::collections::VecDeque;
-use std::ops::Range;
use std::sync::Arc;
+use super::budget::RowBudget;
+
/// Plan for the next queued row group after row-selection slicing.
#[derive(Debug)]
enum QueuedRowGroupDecision {
@@ -39,15 +34,15 @@ enum QueuedRowGroupDecision {
Skip { remaining_budget: RowBudget },
}
-/// Work item handed from [`RowGroupFrontier`] to [`RowGroupReaderBuilder`].
+/// Work item handed from [`RowGroupFrontier`] to
[`RowGroupReaderBuilder`](crate::arrow::push_decoder::reader_builder::RowGroupReaderBuilder).
#[derive(Debug)]
-struct NextRowGroup {
- row_group_idx: usize,
- row_count: usize,
+pub(crate) struct NextRowGroup {
+ pub(crate) row_group_idx: usize,
+ pub(crate) row_count: usize,
/// This row group's selection, or `None` when all rows are selected.
- selection: Option<RowSelection>,
+ pub(crate) selection: Option<RowSelection>,
/// Budget snapshot to apply while decoding this row group.
- budget: RowBudget,
+ pub(crate) budget: RowBudget,
}
/// Row groups and selections that have not yet been handed to the row-group
@@ -178,20 +173,20 @@ impl QueuedRowGroups {
}
#[derive(Debug, Clone)]
-struct RowGroupFrontier {
+pub(crate) struct RowGroupFrontier {
/// Metadata used to resolve row counts for queued row groups.
- parquet_metadata: Arc<ParquetMetaData>,
+ pub(super) parquet_metadata: Arc<ParquetMetaData>,
/// Row groups not yet handed to the builder.
queued: QueuedRowGroups,
/// Offset/limit budget before the next readable row group is planned.
- budget: RowBudget,
+ pub(super) budget: RowBudget,
/// If predicates are present, row groups with selected rows must be read
so
/// the predicate can decide whether they are actually needed.
- has_predicates: bool,
+ pub(super) has_predicates: bool,
}
impl RowGroupFrontier {
- fn new(
+ pub(crate) fn new(
parquet_metadata: Arc<ParquetMetaData>,
row_group_plan: RowGroupPlan,
budget: RowBudget,
@@ -207,7 +202,7 @@ impl RowGroupFrontier {
})
}
- fn update_budget_after_row_group(&mut self, budget: RowBudget) {
+ pub(crate) fn update_budget_after_row_group(&mut self, budget: RowBudget) {
self.budget = budget;
}
@@ -220,14 +215,31 @@ impl RowGroupFrontier {
/// throwaway clone of the frontier, so peek can never drift from the
/// read path. The clone copies the queued row-group plan and selections;
/// see
- /// [`RemainingRowGroups::peek_next_row_group`].
- fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
+ ///
[`RemainingRowGroups::peek_next_row_group`](crate::arrow::push_decoder::remaining::RemainingRowGroups::peek_next_row_group).
+ pub(crate) fn peek_next_row_group(&self) -> Result<Option<usize>,
ParquetError> {
Ok(self
.clone()
.next_readable_row_group()?
.map(|next_row_group| next_row_group.row_group_idx))
}
+ /// Number of row groups not yet handed to the builder.
+ pub(crate) fn row_groups_remaining(&self) -> usize {
+ self.queued.len()
+ }
+
+ /// Decompose into the metadata, the row groups not yet handed to the
+ /// builder, and the remaining offset/limit budget.
+ pub(crate) fn into_parts(self) -> (Arc<ParquetMetaData>, RowGroupPlan,
RowBudget) {
+ let Self {
+ parquet_metadata,
+ queued,
+ budget,
+ has_predicates: _,
+ } = self;
+ (parquet_metadata, queued.into_plan(), budget)
+ }
+
fn clear_remaining(&mut self) {
self.queued.clear();
}
@@ -257,7 +269,7 @@ impl RowGroupFrontier {
}
/// Advance queued row groups until one should be handed to the builder.
- fn next_readable_row_group(&mut self) -> Result<Option<NextRowGroup>,
ParquetError> {
+ pub(crate) fn next_readable_row_group(&mut self) ->
Result<Option<NextRowGroup>, ParquetError> {
loop {
let Some(row_group_idx) = self.queued.front() else {
return Ok(None);
@@ -306,199 +318,6 @@ impl RowGroupFrontier {
}
}
-/// State machine that tracks the remaining high level chunks (row groups) of
-/// Parquet data left to read.
-///
-/// [`RowGroupFrontier`] owns cross-row-group scan state and selects the next
-/// work item. [`RowGroupReaderBuilder`] owns decoding for the active row
group.
-#[derive(Debug)]
-pub(crate) struct RemainingRowGroups {
- /// The arrow schema of the decoded output. Carried only so
- /// [`Self::into_parts`] can hand it to a rebuilt builder; unused while
- /// decoding.
- schema: SchemaRef,
-
- /// Cross-row-group scan state for queued work.
- frontier: RowGroupFrontier,
-
- /// State for building the reader for the current row group
- row_group_reader_builder: RowGroupReaderBuilder,
-}
-
-/// The state recovered from a [`RemainingRowGroups`] by
-/// [`RemainingRowGroups::into_parts`], describing the row groups *not* yet
-/// decoded so a builder reconstructed from it resumes where the decoder left
off.
-#[derive(Debug)]
-pub(crate) struct RemainingRowGroupsParts {
- /// The arrow schema of the decoded output.
- pub schema: SchemaRef,
- /// The Parquet file metadata.
- pub metadata: Arc<ParquetMetaData>,
- /// Row groups and selections not yet handed to the reader builder.
- pub row_group_plan: RowGroupPlan,
- /// Offset still to be skipped before the next readable row group.
- pub offset: Option<usize>,
- /// Output rows still permitted across the remaining row groups.
- pub limit: Option<usize>,
- /// Builder-configurable parts of the inner row-group reader builder.
- pub reader_builder: RowGroupReaderBuilderParts,
-}
-
-impl RemainingRowGroups {
- pub fn new(
- schema: SchemaRef,
- parquet_metadata: Arc<ParquetMetaData>,
- row_group_plan: RowGroupPlan,
- budget: RowBudget,
- has_predicates: bool,
- row_group_reader_builder: RowGroupReaderBuilder,
- ) -> Result<Self, ParquetError> {
- Ok(Self {
- schema,
- frontier: RowGroupFrontier::new(
- parquet_metadata,
- row_group_plan,
- budget,
- has_predicates,
- )?,
- row_group_reader_builder,
- })
- }
-
- /// Decompose into [`RemainingRowGroupsParts`].
- ///
- /// Must be called at a row-group boundary (see
- /// [`Self::is_at_row_group_boundary`]). The inner reader builder's runtime
- /// decode state is discarded; its buffered bytes are carried through.
- pub(crate) fn into_parts(self) -> RemainingRowGroupsParts {
- let Self {
- schema,
- frontier,
- row_group_reader_builder,
- } = self;
- // `has_predicates` is recomputed by `build()` from the filter.
- let RowGroupFrontier {
- parquet_metadata,
- queued,
- budget,
- has_predicates: _,
- } = frontier;
- let row_group_plan = queued.into_plan();
- RemainingRowGroupsParts {
- schema,
- metadata: parquet_metadata,
- row_group_plan,
- offset: budget.offset(),
- limit: budget.limit(),
- reader_builder: row_group_reader_builder.into_parts(),
- }
- }
-
- /// Push new data buffers that can be used to satisfy pending requests
- pub fn push_data(
- &mut self,
- ranges: Vec<Range<u64>>,
- buffers: Vec<Bytes>,
- ) -> Result<(), ParquetError> {
- self.row_group_reader_builder.push_data(ranges, buffers)
- }
-
- /// Return the total number of bytes buffered so far
- pub fn buffered_bytes(&self) -> u64 {
- self.row_group_reader_builder.buffered_bytes()
- }
-
- /// Clear any staged ranges currently buffered for future decode work
- pub fn clear_all_ranges(&mut self) {
- self.row_group_reader_builder.clear_all_ranges();
- }
-
- /// True iff the inner row-group reader is between row groups (state
- /// `Finished`). Forward to [`RowGroupReaderBuilder::is_finished`].
- pub fn is_at_row_group_boundary(&self) -> bool {
- self.row_group_reader_builder.is_finished()
- }
-
- /// Number of row groups remaining (not including the one currently
- /// being decoded).
- pub fn row_groups_remaining(&self) -> usize {
- self.frontier.queued.len()
- }
-
- /// Peek at the file-level row-group index that the next call to
- /// [`Self::try_next_reader`] will produce a reader for, after
- /// simulating the same skip logic [`Self::try_next_reader`] applies
- /// internally (row-selection emptiness + offset/limit budget). Does
- /// not mutate state.
- ///
- /// Returns `None` when the active row group is still being decoded,
- /// when no row groups remain, or when every remaining row group
- /// would be skipped under the current selection/budget.
- ///
- /// Cost: one clone of the queued row-group plan and selections per call
- /// (the frontier is cloned so the real advance logic can run
- /// non-destructively). For callers that peek once per row-group boundary
- /// this is O(remaining row groups + selectors) per boundary.
- pub fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
- if self.row_group_reader_builder.has_active_row_group() {
- return Ok(None);
- }
- self.frontier.peek_next_row_group()
- }
-
- /// returns [`ParquetRecordBatchReader`] suitable for reading the next
- /// group of rows from the Parquet data, or the list of data ranges still
- /// needed to proceed
- pub fn try_next_reader(
- &mut self,
- ) -> Result<DecodeResult<ParquetRecordBatchReader>, ParquetError> {
- loop {
- if !self.row_group_reader_builder.has_active_row_group() {
- // We are done with the previous row group, seek to the next
one
- // from the frontier, if any.
-
- match self.frontier.next_readable_row_group()? {
- Some(NextRowGroup {
- row_group_idx,
- row_count,
- selection,
- budget,
- }) => {
- self.row_group_reader_builder.next_row_group(
- row_group_idx,
- row_count,
- selection,
- budget,
- )?;
- }
- None => return Ok(DecodeResult::Finished),
- }
- }
-
- match self.row_group_reader_builder.try_build()? {
- RowGroupBuildResult::Finished { remaining_budget } => {
- self.frontier
- .update_budget_after_row_group(remaining_budget);
- // reader is done, proceed to the next row group
- }
- RowGroupBuildResult::NeedsData(ranges) => {
- // need more data to proceed
- return Ok(DecodeResult::NeedsData(ranges));
- }
- RowGroupBuildResult::Data {
- batch_reader,
- remaining_budget,
- } => {
- self.frontier
- .update_budget_after_row_group(remaining_budget);
- // ready to read the row group
- return Ok(DecodeResult::Data(batch_reader));
- }
- }
- }
- }
-}
-
#[cfg(test)]
mod tests {
use super::*;
diff --git a/parquet/src/arrow/push_decoder/scan_plan/mod.rs
b/parquet/src/arrow/push_decoder/scan_plan/mod.rs
new file mode 100644
index 0000000000..f4c8f27d1f
--- /dev/null
+++ b/parquet/src/arrow/push_decoder/scan_plan/mod.rs
@@ -0,0 +1,25 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Row-group scheduling for the push decoder: which row groups to read next,
+//! and how the offset/limit budget carries across them.
+
+mod budget;
+mod frontier;
+
+pub(crate) use budget::{BudgetedReadPlan, RowBudget};
+pub(crate) use frontier::{NextRowGroup, RowGroupFrontier};
