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-11393-20475454f3d84a9dfd24d2a003d681c27734b613 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 4056f9ce95a4e28e9bf2d3b652a9938328a287ab Author: Adrian Garcia Badaracco <[email protected]> AuthorDate: Wed Oct 7 19:43:41 2026 +0000 refactor(parquet): build the push decoder StageSchedule once (#11393) # Which issue does this PR close? - Follow-up to https://github.com/apache/arrow-rs/pull/11330, for [this review comment](https://github.com/apache/arrow-rs/pull/11330#discussion_r4189970839). # Rationale for this change `RowGroupReaderBuilder::new` made a placeholder `StageSchedule`, then replaced it a few lines down: ```rust let mut builder = Self { // ... stages: StageSchedule::new(ProjectionMask::all(), vec![], None), }; // ... compute the projections with `builder.compute_cache_projection_inner(filter)` ... builder.stages = StageSchedule::new(/* the real values */); builder ``` The placeholder was necessary because `compute_cache_projection_inner` was a `&self` method, so the builder had to exist before the schedule could be computed. The placeholder did not allocate on the heap: `ProjectionMask::all()` holds `None` and `vec![]` is an empty `Vec`. But it was a value that was built only to be discarded, and a reader had to check that it was never used. # What changes are included in this PR? - `compute_cache_projection_inner` becomes an associated function. It takes the filter, the projection, the metadata and the predicate cache size as arguments. - `RowGroupReaderBuilder::new` computes the `StageSchedule` first, then builds the struct one time. The `let mut builder` and the placeholder are gone. - `exclude_nested_columns_from_cache` had one caller and one line. It is now inline in `compute_cache_projection_inner`. There is no behavior change. # Are these changes tested? Yes, by the existing tests. These pass locally: - `cargo test -p parquet --all-features --lib -- push_decoder arrow:: async_reader` (555 tests) - `cargo test -p parquet --all-features --test arrow_reader` (150 tests) - `cargo clippy -p parquet --all-features --all-targets -- -D warnings` # Are there any user-facing changes? No. All changed items are private. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 5.5 <[email protected]> --- .../src/arrow/push_decoder/reader_builder/mod.rs | 70 ++++++++++++---------- 1 file changed, 40 insertions(+), 30 deletions(-) diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs index 074f8fc427..8f122863be 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs @@ -219,7 +219,25 @@ impl RowGroupReaderBuilder { buffers: PushBuffers, row_selection_policy: RowSelectionPolicy, ) -> Self { - let mut builder = Self { + let (predicate_projections, cache_projection) = match &filter { + Some(filter) => ( + filter + .predicates + .iter() + .map(|predicate| predicate.projection().clone()) + .collect(), + Self::compute_cache_projection_inner( + filter, + &projection, + &metadata, + max_predicate_cache_size, + ), + ), + None => (vec![], None), + }; + let stages = + StageSchedule::new(projection.clone(), predicate_projections, cache_projection); + Self { batch_size, projection, metadata, @@ -230,25 +248,8 @@ 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 + stages, + } } /// Decompose into [`RowGroupReaderBuilderParts`] so the builder can be @@ -761,28 +762,37 @@ impl RowGroupReaderBuilder { /// final projection, excluding any nested columns. fn compute_cache_projection(&self, row_group_idx: usize, filter: &RowFilter) -> ProjectionMask { let meta = self.metadata.row_group(row_group_idx); - match self.compute_cache_projection_inner(filter) { + let cache_projection = Self::compute_cache_projection_inner( + filter, + &self.projection, + &self.metadata, + self.max_predicate_cache_size, + ); + match cache_projection { Some(projection) => projection, None => ProjectionMask::none(meta.columns().len()), } } - fn compute_cache_projection_inner(&self, filter: &RowFilter) -> Option<ProjectionMask> { + /// An associated function, so that [`Self::new`] can call it before the + /// builder exists. + fn compute_cache_projection_inner( + filter: &RowFilter, + projection: &ProjectionMask, + metadata: &ParquetMetaData, + max_predicate_cache_size: usize, + ) -> Option<ProjectionMask> { // Do not compute the projection mask if the predicate cache is disabled - if self.max_predicate_cache_size == 0 { + if max_predicate_cache_size == 0 { return None; } let mut cache_projection = filter.predicates.first()?.projection().clone(); for predicate in &filter.predicates { cache_projection.union(predicate.projection()); } - cache_projection.intersect(&self.projection); - self.exclude_nested_columns_from_cache(&cache_projection) - } - - /// Exclude leaves belonging to roots that span multiple parquet leaves (i.e. nested columns) - fn exclude_nested_columns_from_cache(&self, mask: &ProjectionMask) -> Option<ProjectionMask> { - mask.without_nested_types(self.metadata.file_metadata().schema_descr()) + cache_projection.intersect(projection); + // Exclude leaves belonging to roots that span multiple parquet leaves (i.e. nested columns) + cache_projection.without_nested_types(metadata.file_metadata().schema_descr()) } /// Get the offset index for the specified row group, if any
