zhuqi-lucas commented on code in PR #23696:
URL: https://github.com/apache/datafusion/pull/23696#discussion_r3766670405
##########
datafusion/datasource-parquet/src/opener/mod.rs:
##########
@@ -1464,25 +1471,66 @@ impl RowGroupsPrunedParquetOpen {
};
let prepared_access_plan = prepare_access_plan(access_plan)?;
+ // Build `rg_plan` parallel to the decoder's view: the
+ // `prepared_access_plan` has already had its empty-selection
+ // row groups stripped, so 1:1 correspondence with the readers
+ // arrow-rs will hand back is restored. We zip with the
+ // `fully_matched` flag so the stream can toggle the per-row
+ // `RowFilter` per RG.
let rg_plan: VecDeque<RgPlanEntry> = prepared_access_plan
.row_group_indexes
.iter()
.copied()
- .map(|rg_index| RgPlanEntry { rg_index })
+ .zip(prepared_access_plan.fully_matched.iter().copied())
+ .map(|(rg_index, fully_matched)| RgPlanEntry {
+ rg_index,
+ fully_matched,
+ })
.collect();
+ // Decide the initial row filter state based on the first RG to
+ // read. If that RG is `fully_matched` the per-row predicate is
+ // a no-op for every row, so we install an empty `RowFilter`
+ // (arrow-rs's `has_predicates` check then short-circuits the
+ // per-row eval) and the stream toggles back to the real filter
+ // at the first non-fully-matched RG boundary.
+ //
+ // `RowFilterContext` carries everything `build_row_filter`
+ // needs so the stream can regenerate the filter later — the
+ // installed filter is owned by the decoder and is not
+ // recoverable once replaced.
+ let first_rg_fully_matched = rg_plan.front().is_some_and(|e|
e.fully_matched);
+ let initial_filter = precomputed_context
+ .as_ref()
+ .and_then(|ctx| ctx.build_row_filter());
+ let row_filter_context = precomputed_context;
+
let mut builder =
decoder_config.build(prepared_access_plan,
reader_metadata.clone());
- if let Some(row_filter) = row_filter_generator.next_filter() {
- builder = builder.with_row_filter(row_filter);
- if let Some(max_predicate_cache_size) =
prepared.max_predicate_cache_size
- {
- builder =
-
builder.with_max_predicate_cache_size(max_predicate_cache_size);
+ let mut filter_installed = false;
+ if let Some(row_filter) = initial_filter {
+ if first_rg_fully_matched {
+ builder = builder.with_row_filter(
+ parquet::arrow::arrow_reader::RowFilter::new(vec![]),
Review Comment:
Good catch — fixed in b31b0b4: the first-RG-fully-matched path now calls
`row_filter_skipped_fully_matched.add(1)`, matching the mid-scan toggle.
`filter_installed` stays false (no real filter is installed).
##########
datafusion/datasource-parquet/src/push_decoder.rs:
##########
@@ -414,6 +498,123 @@ impl PushDecoderStreamState {
}
}
+ /// Keep `rg_plan.front()` aligned with the row group the decoder will emit
+ /// next. `try_next_reader` silently skips row groups whose row selection
is
+ /// empty (e.g. page-index pruning removed every page), which would
otherwise
+ /// leave `rg_plan` off-by-one from the decoder's frontier.
+ fn sync_rg_plan_to_decoder_frontier(&mut self) -> Result<(),
DataFusionError> {
+ match self
+ .decoder
+ .as_ref()
+ .expect("decoder present")
+ .peek_next_row_group()
+ .map_err(DataFusionError::from)?
+ {
+ Some(actual) => self.advance_rg_plan_to(actual),
+ // Decoder has nothing left to emit — drain our plan so the stream
+ // finishes cleanly.
+ None => self.rg_plan.clear(),
+ }
+ Ok(())
+ }
+
+ /// Pop `rg_plan` entries until its front is `target` (or it empties).
+ fn advance_rg_plan_to(&mut self, target: usize) {
Review Comment:
Agreed — fixed in b31b0b4. `advance_rg_plan_to` now returns an internal
error if `target` is not in `rg_plan` (decoder/plan divergence) instead of
silently draining the plan.
##########
datafusion/datasource-parquet/src/opener/mod.rs:
##########
@@ -1435,18 +1434,26 @@ impl RowGroupsPrunedParquetOpen {
prepared.virtual_state.as_deref(),
)?;
- let (decoder, rg_plan) = {
+ let (decoder, rg_plan, filter_installed, row_filter_context) = {
Review Comment:
Good idea — will factor the initial decoder / rg_plan / filter setup into a
helper returning a named struct (e.g. `InitialDecoderState`) as a follow-up,
keeping this commit focused on the correctness fixes.
##########
datafusion/datasource-parquet/src/access_plan.rs:
##########
@@ -571,12 +571,81 @@ impl ParquetAccessPlan {
row_group_meta_data: &[RowGroupMetaData],
) -> Result<PreparedAccessPlan> {
let row_group_indexes = self.row_group_indexes();
+ // Carry `fully_matched` flags in the same order as
+ // `row_group_indexes` so downstream code (per-RG `RowFilter` skip)
+ // can look them up positionally.
+ let fully_matched: Vec<bool> = row_group_indexes
+ .iter()
+ .map(|&idx| self.fully_matched[idx])
+ .collect();
let row_selection =
self.into_overall_row_selection(row_group_meta_data)?;
- PreparedAccessPlan::new(row_group_indexes, row_selection)
+ let (row_group_indexes, fully_matched, row_selection) =
strip_empty_row_groups(
+ row_group_indexes,
+ fully_matched,
+ row_selection,
+ row_group_meta_data,
+ );
+
+ PreparedAccessPlan::new(row_group_indexes, fully_matched,
row_selection)
}
}
+/// Strip row groups whose post-pruning `RowSelection` selects zero rows.
+///
+/// arrow-rs's push decoder silently advances past such row groups inside
+/// `try_next_reader`, but the rest of DataFusion (per-RG metadata maps,
+/// the runtime dynamic-pruner, the per-RG `RowFilter` toggle) assumes a
+/// 1:1 correspondence between the prepared plan and the readers the
+/// decoder hands back. Removing these empty entries here keeps that
+/// invariant and lets downstream code consult per-RG state — like
+/// [`PreparedAccessPlan::fully_matched`] — without going out of sync.
+///
+/// The flat `RowSelection` is split per row group with
+/// [`RowSelection::split_off`] (mirroring arrow-rs's own logic) and the
+/// surviving segments are concatenated back into the result selection.
+/// When `row_selection` is `None` (no page-index pruning, no
+/// user-supplied selection) no row group can be empty and the inputs are
+/// returned unchanged.
Review Comment:
This came in with #22450 (not on `main` yet) and is currently coupled to
`fully_matched` — it keeps `row_group_indexes`, `fully_matched`, and the
`RowSelection` aligned while stripping empty RGs. The empty-RG stripping itself
is generic, so happy to extract that part into its own PR as a follow-up once
this lands.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
