zhuqi-lucas commented on code in PR #25608:
URL: https://github.com/apache/datafusion/pull/25608#discussion_r4078504623
##########
datafusion/datasource-parquet/src/access_plan.rs:
##########
@@ -565,137 +576,83 @@ impl ParquetAccessPlan {
self.row_groups
}
- /// Prepare this plan and resolve to the final `PreparedAccessPlan`
+ /// Validate selections and prepare row groups for the decoder without
+ /// converting local selections into a single file-level selection.
pub(crate) fn prepare(
- mut self,
+ self,
row_group_meta_data: &[RowGroupMetaData],
) -> Result<PreparedAccessPlan> {
- let row_group_indexes = self.row_group_indexes();
- // `fully_matched` is indexed by absolute row-group index; take it out
- // before `into_overall_row_selection` consumes `self`.
- let fully_matched_by_index = std::mem::take(&mut self.fully_matched);
- let row_selection =
self.into_overall_row_selection(row_group_meta_data)?;
-
- let (row_group_indexes, row_selection) =
- strip_empty_row_groups(row_group_indexes, row_selection,
row_group_meta_data);
-
- // Carry `fully_matched` flags in the same order as the *surviving*
- // `row_group_indexes` so downstream code (per-RG `RowFilter` skip) can
- // look them up positionally. Mapping after the strip keeps
- // `strip_empty_row_groups` generic (no `fully_matched` parameter).
- let fully_matched: Vec<bool> = row_group_indexes
+ assert_eq!(row_group_meta_data.len(), self.row_groups.len());
+ // Keep the scan policies that previously depended on the presence of
+ // an overall selection, even when its partial groups are all empty.
+ let has_row_selection = self
Review Comment:
This restriction gets more interesting after this PR. Statistics reordering
and runtime pruning were disabled in the presence of a selection because the
coordinates were global — dropping a row group at runtime invalidated every
offset after it in the flat selection. With row-group-local selections that
reason is gone: a dropped group takes its own selection with it and nothing
else shifts.
Lifting it is probably what actually closes #24358. Worth a follow-up issue
if you're not planning to do it here.
##########
datafusion/datasource-parquet/src/access_plan.rs:
##########
@@ -565,137 +576,83 @@ impl ParquetAccessPlan {
self.row_groups
}
- /// Prepare this plan and resolve to the final `PreparedAccessPlan`
+ /// Validate selections and prepare row groups for the decoder without
+ /// converting local selections into a single file-level selection.
pub(crate) fn prepare(
- mut self,
+ self,
row_group_meta_data: &[RowGroupMetaData],
) -> Result<PreparedAccessPlan> {
- let row_group_indexes = self.row_group_indexes();
- // `fully_matched` is indexed by absolute row-group index; take it out
- // before `into_overall_row_selection` consumes `self`.
- let fully_matched_by_index = std::mem::take(&mut self.fully_matched);
- let row_selection =
self.into_overall_row_selection(row_group_meta_data)?;
-
- let (row_group_indexes, row_selection) =
- strip_empty_row_groups(row_group_indexes, row_selection,
row_group_meta_data);
-
- // Carry `fully_matched` flags in the same order as the *surviving*
- // `row_group_indexes` so downstream code (per-RG `RowFilter` skip) can
- // look them up positionally. Mapping after the strip keeps
- // `strip_empty_row_groups` generic (no `fully_matched` parameter).
- let fully_matched: Vec<bool> = row_group_indexes
+ assert_eq!(row_group_meta_data.len(), self.row_groups.len());
+ // Keep the scan policies that previously depended on the presence of
+ // an overall selection, even when its partial groups are all empty.
+ let has_row_selection = self
+ .row_groups
.iter()
- .map(|&idx| fully_matched_by_index[idx])
- .collect();
-
- PreparedAccessPlan::new(row_group_indexes, fully_matched,
row_selection)
- }
-}
-
-/// Drop row groups whose post-pruning `RowSelection` selects zero rows, so the
-/// prepared plan stays well-formed: `with_row_groups(...)` never names a row
-/// group the decoder would immediately skip.
-///
-/// arrow-rs's push decoder silently advances past an all-skipped row group
-/// inside `try_next_reader` without handing back a reader, so keeping such a
-/// row group in the plan would leave its row-group list one entry longer than
-/// the readers the decoder produces. An empty selection is reachable today via
-/// [`ParquetAccessPlan::scan_selection`], which intersects an existing
-/// `Selection` with a new one and can leave a row group selecting nothing.
-/// A well-formed plan here is also a precondition for re-enabling runtime
-/// pruning under a live selection (#24358) and for #23696.
-///
-/// Walks the flat `RowSelection` once with [`OverallRowSelectionCursor`],
-/// rather than a per-row-group [`RowSelection::split_off`] (which reallocates
-/// the selector tail on every call, i.e. O(row groups × selectors)), keeping
-/// only the row groups that still select at least one row.
`row_group_meta_data`
-/// is the **full file** metadata, indexed absolutely by row-group index, while
-/// `row_selection` covers only the scanned row groups — matching
-/// `into_overall_row_selection`, which emits nothing for a skipped row group.
-/// When `row_selection` is `None` no row group can be empty and the inputs are
-/// returned unchanged.
-fn strip_empty_row_groups(
- row_group_indexes: Vec<usize>,
- row_selection: Option<RowSelection>,
- row_group_meta_data: &[RowGroupMetaData],
-) -> (Vec<usize>, Option<RowSelection>) {
- let Some(selection) = row_selection else {
- return (row_group_indexes, None);
- };
-
- let mut cursor = OverallRowSelectionCursor::new(selection);
- let mut kept_indexes = Vec::with_capacity(row_group_indexes.len());
- let mut kept_selectors: Vec<RowSelector> = Vec::new();
-
- for &rg_idx in row_group_indexes.iter() {
- let rg_row_count = row_group_meta_data[rg_idx].num_rows() as usize;
- // Pull this row group's fragments off the shared cursor in a single
- // pass (no `split_off` tail reallocation).
- let start = kept_selectors.len();
- let mut selected = 0usize;
- let mut taken = 0usize;
- while taken < rg_row_count {
- let Some(fragment) = cursor.take(rg_row_count - taken) else {
- break;
+ .any(|access| matches!(access, RowGroupAccess::Selection(_)));
+ let mut row_groups = Vec::with_capacity(self.row_groups.len());
+ for (index, (access, fully_matched)) in self
+ .row_groups
+ .into_iter()
+ .zip(self.fully_matched)
+ .enumerate()
+ {
+ let selection = match access {
+ RowGroupAccess::Skip => continue,
+ RowGroupAccess::Scan => {
+ if has_row_selection &&
row_group_meta_data[index].num_rows() == 0 {
Review Comment:
Why is the zero-row skip gated on `has_row_selection`? An empty row group
contributes nothing either way, so the asymmetry looks like it's preserving
some exact prior behaviour — a short comment naming which one would help the
next reader.
##########
datafusion/datasource-parquet/src/push_decoder.rs:
##########
@@ -100,10 +101,7 @@ impl DecoderBuilderConfig<'_> {
if self.force_filter_selections {
builder =
builder.with_row_selection_policy(RowSelectionPolicy::Selectors);
}
- if let Some(row_selection) = prepared_access_plan.row_selection {
- builder = builder.with_row_selection(row_selection);
- }
- builder =
builder.with_row_groups(prepared_access_plan.row_group_indexes);
+ builder = builder.with_row_group_selections(row_group_selections);
Review Comment:
Might be worth spelling out in the description which drift this closes and
which it doesn't. `sync_rg_plan_to_decoder_frontier` further down is for a
different one — arrow-rs finishes a row group whose post-predicate selection is
empty without handing back a reader, so `rg_plan` trails the decoder by one
(#24352). That's untouched here and that seems right, but #24358 reads like it
covers both, so it would be easy to over-read this as closing it.
--
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]
