adriangb commented on code in PR #24509:
URL: https://github.com/apache/datafusion/pull/24509#discussion_r3821711474
##########
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:
This paragraph describes a desync that can't currently occur, which makes
the function look like it's fixing a live bug rather than establishing an
invariant.
Neither in-tree producer can put an empty `Selection` here: `page_filter.rs`
maps an all-skip selection to `access_plan.skip(row_group_index)` instead of
`scan_selection`, and `try_new_from_overall_row_selection` normalizes the same
case through `RowGroupAccessBuilder::into_access` (its own test asserts
`RowGroupAccess::Skip` for the fully-skipped group). And the "runtime
dynamic-pruner" named here is switched off entirely whenever
`has_row_selection` is true, so it never observes the misalignment.
Suggest re-scoping this to what it actually buys — the plan stays
well-formed so `with_row_groups(...)` never names a row group the decoder will
skip, which is a precondition for re-enabling pruning under a live selection
(#24358) and for #23696 — rather than describing a bookkeeping failure that's
already guarded.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
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:
Worth a note here: this indexes `row_group_meta_data` *absolutely*, while
`remaining` only covers the row groups that are actually scanned. That pairing
is correct — `into_overall_row_selection` emits nothing for
`RowGroupAccess::Skip` — but it's the subtlest thing in the function, and both
new tests pass `vec![0, 1, 2, 3]`, so nothing pins it down.
A case with a non-contiguous list (e.g. `vec![1, 3]`, RGs 0 and 2 skipped)
would lock in exactly this. A one-line `debug_assert!` on the precondition, or
a sentence saying the slice is the full file metadata, would also help — note
that `into_overall_row_selection` uses `zip` in the same situation, so a short
slice truncates silently there but panics here.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
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:
These tests were inserted between this doc comment and the
`ROW_GROUP_METADATA` static it describes, so `/// [\`RowGroupMetaData\`] that
returns 4 row groups with 10, 20, 30, 40 rows respectively` now documents
`test_strip_empty_row_groups_drops_only_empties`. Moving both tests below the
static (line 1138) fixes it.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
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:
This reintroduces the pattern that `try_new_from_overall_row_selection` was
deliberately written to avoid. Its comment says it directly:
> Keep this as a single pass over the selector stream rather than repeatedly
calling `RowSelection::split_off` per row group. The `split_off` version is
simpler, but it clones/retains substantially more selector buffer capacity for
highly fragmented selections.
On a selectors-backed selection each `split_off` bottoms out in
`Vec::split_off`, allocating and copying the tail, so this loop is O(row_groups
× selectors). Since `into_overall_row_selection` always builds via `collect()`
from a `RowSelector` iterator, the input is always selectors-backed and always
takes that path.
`OverallRowSelectionCursor` is right there and would drive this in one pass
— consistent with the sibling function and effectively free. This is on the
per-file open path so the absolute cost is probably small, but it's the change
I'd most want before merge.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
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
+ #[test]
+ fn test_strip_empty_row_groups_drops_only_empties() {
+ // 4 row groups of [10, 20, 30, 40] rows. RG 1 and RG 3 select nothing
+ // after pruning, so they must be dropped; RG 0 and RG 2 survive with
+ // their re-concatenated selections intact.
+ let selection = RowSelection::from(vec![
+ RowSelector::select(10), // RG 0: keep all 10
+ RowSelector::skip(30), // RG 1 (20) fully skipped + RG 2's
leading 10
+ RowSelector::select(20), // RG 2: keep 20
+ RowSelector::skip(40), // RG 3: skip all 40
+ ]);
+
+ let (indexes, result) = strip_empty_row_groups(
+ vec![0, 1, 2, 3],
+ Some(selection),
+ &ROW_GROUP_METADATA,
+ );
+
+ // RG 1 and RG 3 are dropped; the surviving indexes stay in order.
+ assert_eq!(indexes, vec![0, 2]);
+ let result = result.expect("survivors keep a selection");
+ assert_eq!(result.row_count(), 30); // 10 from RG 0 + 20 from RG 2
+ assert_eq!(result.skipped_row_count(), 10); // RG 2's leading skip
+ }
+
+ #[test]
+ fn test_strip_empty_row_groups_none_selection_unchanged() {
+ // With no row selection no row group can be empty, so the inputs pass
+ // through unchanged.
+ let (indexes, result) =
+ strip_empty_row_groups(vec![0, 1, 2, 3], None,
&ROW_GROUP_METADATA);
+ assert_eq!(indexes, vec![0, 1, 2, 3]);
+ assert!(result.is_none());
+ }
Review Comment:
Both tests call `strip_empty_row_groups` directly — nothing exercises it
through `ParquetAccessPlan::prepare`, and nothing lands in
`datafusion/core/tests/parquet/external_access_plan.rs`.
That gap matters because there *is* a reachable producer of an empty
`Selection`, and it isn't the one the doc comment names. `scan_selection`
intersects an existing `Selection` with the new one and never re-normalizes an
empty result back to `Skip`:
```rust
RowGroupAccess::Selection(existing_selection) => {
RowGroupAccess::Selection(existing_selection.intersection(&selection))
}
```
So an externally supplied `ParquetRowSelection` plus page-index pruning
whose surviving pages are disjoint from the user's selection *within a row
group* lands an empty `Selection` in `prepare()` — the page-index path's own
`rows_selected > 0` guard doesn't catch it, because it checks the incoming
selection, not the intersection.
That scenario is the regression test this PR wants, and it also suggests a
complementary one-line fix in `scan_selection` itself: normalizing an empty
intersection to `Skip` keeps `ParquetAccessPlan` well-formed for
`row_group_indexes()`, the metrics, and `is_fully_matched` too — not just for
`prepare()`.
---
_Generated by [Claude Code](https://claude.ai/code)_
--
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]
