zhuqi-lucas commented on code in PR #24509:
URL: https://github.com/apache/datafusion/pull/24509#discussion_r3822663442
##########
datafusion/datasource-parquet/src/access_plan.rs:
##########
@@ -573,10 +573,61 @@ impl ParquetAccessPlan {
let row_group_indexes = self.row_group_indexes();
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);
+
PreparedAccessPlan::new(row_group_indexes, 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 and the
+/// runtime dynamic-pruner) assumes a 1:1 correspondence between the prepared
+/// plan and the readers the decoder hands back. Removing these empty entries
+/// here keeps that invariant so downstream per-RG bookkeeping stays in sync
+/// with the decoder.
Review Comment:
Thanks @adriangb, good call. Reframed the doc around the well-formed-plan
invariant and the one reachable producer (`scan_selection` intersecting to
empty), and dropped the "runtime dynamic-pruner desync" description since that
path is gated off. Retitled `fix:` → `refactor:` and reworded the PR body to
match.
##########
datafusion/datasource-parquet/src/access_plan.rs:
##########
@@ -573,10 +573,61 @@ impl ParquetAccessPlan {
let row_group_indexes = self.row_group_indexes();
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);
+
PreparedAccessPlan::new(row_group_indexes, 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 and the
+/// runtime dynamic-pruner) assumes a 1:1 correspondence between the prepared
+/// plan and the readers the decoder hands back. Removing these empty entries
+/// here keeps that invariant so downstream per-RG bookkeeping stays in sync
+/// with the decoder.
+///
+/// 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.
+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(mut remaining) = row_selection else {
+ return (row_group_indexes, None);
+ };
+
+ 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;
Review Comment:
Thanks — added `test_strip_empty_row_groups_non_contiguous_indexes`
(`vec![1, 3]`, RGs 0/2 skipped) which pins the absolute indexing, and the doc
now states `row_group_meta_data` is the full-file metadata indexed absolutely
while `row_selection` covers only the scanned groups.
##########
datafusion/datasource-parquet/src/access_plan.rs:
##########
@@ -573,10 +573,61 @@ impl ParquetAccessPlan {
let row_group_indexes = self.row_group_indexes();
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);
+
PreparedAccessPlan::new(row_group_indexes, 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 and the
+/// runtime dynamic-pruner) assumes a 1:1 correspondence between the prepared
+/// plan and the readers the decoder hands back. Removing these empty entries
+/// here keeps that invariant so downstream per-RG bookkeeping stays in sync
+/// with the decoder.
+///
+/// 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.
+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(mut remaining) = row_selection else {
+ return (row_group_indexes, None);
+ };
+
+ 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;
+ // `split_off` cuts off the first `rg_row_count` rows worth of
+ // selection — this row group's segment. `remaining` keeps the rest.
+ let rg_segment = remaining.split_off(rg_row_count);
Review Comment:
Thanks, agreed — this was the one I most wanted to fix too. Rewrote it to
walk the selection once with `OverallRowSelectionCursor` (same single-pass
approach as `try_new_from_overall_row_selection`), so no more per-row-group
`split_off` tail reallocation.
##########
datafusion/datasource-parquet/src/access_plan.rs:
##########
@@ -1049,6 +1100,41 @@ mod test {
/// [`RowGroupMetaData`] that returns 4 row groups with 10, 20, 30, 40 rows
/// respectively
Review Comment:
Thanks — moved both tests below `ROW_GROUP_METADATA` so its doc comment
attaches to the static again.
--
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]
