adriangb commented on code in PR #23696:
URL: https://github.com/apache/datafusion/pull/23696#discussion_r3766050626
##########
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:
This seems like a good candidate to factor out into a function returning a
named struct w/ fields
##########
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 seems like a good general improvement that already applies to `main`,
should we factor it out as an independent PR if so?
##########
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:
If target isn't present in rg_plan (an invariant violation, the decoder's
frontier names an RG our plan doesn't know), this loop silently drains. Can we
add a check that `target` is found and if it's not throw an `internal_err!()`
instead of possibly returning bogus results?
##########
datafusion/datasource-parquet/src/opener/mod.rs:
##########
@@ -3455,6 +3510,44 @@ mod test {
))
}
+ #[tokio::test]
+ async fn test_input_file_name_projection() {
Review Comment:
This test seems to have moved with no changes - location diff noise?
##########
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:
should this branch set `filter_installed = true` or
`prepared.file_metrics.row_filter_skipped_fully_matched.add(1)`?
--
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]
