This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11330-934e16ec8513fe8c1c8a7670f3741070a4c4781d in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
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};
