apache / apache/datafusion

Structurally eliminate rg_plan / RowSelection drift in the parquet runtime prune path (retain_row_groups)

Open
#24,358 2 comments 0 reactions 1 assignee Claimed by @zhuqi-lucas View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.