Structurally eliminate rg_plan / RowSelection drift in the parquet runtime prune path (retain_row_groups)
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
Follow-up to #24352 / #24354 (and the second instance #24355), from @adriangb's review on #24354.
## Problem
The parquet runtime row-group prune path in `PushDecoderStreamState` maintains state **parallel** to the arrow-rs push decoder and updates it independently of the decoder's own frontier:
- `rg_plan: VecDeque` — a DataFusion-side queue of pending row groups, and
- a carried flat `RowSelection` over the concatenation of the *remaining* row groups (arrow-rs side).
Both must stay aligned with the decoder's row-group frontier, but nothing enforces that structurally, and each drift is a silent wrong-results bug:
- **#24352** — a row group whose post-predicate selection is empty is silently finished without handing back a reader, so `rg_plan` trails the decoder by one and a later rebuild re-reads an already-delivered row group. Fixed in #24354 by syncing `rg_plan` to the frontier via `peek_next_row_group()`.
- **#24355** — `into_builder().with_row_groups(new_indices)` drops row groups without slicing the carried `RowSelection` to match, so selectors intended for a dropped RG are applied to the next surviving one.
The #24354 fix also couples `rg_plan`'s correctness to `row_group_pruner.is_some()` (it syncs only when a pruner exists) — sound today because the pruner is the only consumer that rebuilds from `rg_plan`, but fragile: the invariant is enforced only where a pruner happens to be present.
## Proposed direction
Give the decoder ownership of the drop so the row-group queue and the `RowSelection` cannot be updated independently. As @adriangb suggested:
- **arrow-rs**: add `ParquetPushDecoder::retain_row_groups(impl FnMut(usize) -> bool)` (or `remaining_row_groups() -> impl ExactSizeIterator`) that filters the frontier **in place**, slicing the `RowSelection` alongside the row-group queue.
- **DataFusion**: replace `rg_plan` / `RgPlanEntry` / `sync_rg_plan_to_decoder_frontier` / `advance_rg_plan_to` / the `Data`-arm pop / and the `into_builder()` + `is_at_row_group_boundary()` dance with a single
```rust
decoder.retain_row_groups(|rg| !pruner.should_prune(&[rg]));
```
This removes the parallel state entirely, so there is no alignment left to drift — closing the whole family (#24352, #24355) at the root and dropping the `pruner.is_some()` coupling.
Either repo can host the change; the in-place `retain` really wants to live in arrow-rs, which owns both the queue and the selection.
cc @alamb @adriangb @hhhizzz
Contributor guide
Assessment
This issue has not been assessed yet.