feat: Add row-group-local RowSelection support to the push decoder - #10702
feat: Add row-group-local RowSelection support to the push decoder#10702haohuaijin wants to merge 1 commit into
Conversation
haohuaijin
changed the title
This PR adds |
02cae11 to
a9157e1
Compare
There was a problem hiding this comment.
Thanks @haohuaijin for this work!
From the DataFusion side: root fix for apache/datafusion#24352 / apache/datafusion#24355 — into_builder preserving local selections is exactly it. Core looks correct and well-tested; small notes inline. Happy to take the DataFusion migration (apache/datafusion#24358) once this lands.
| fields, | ||
| batch_size, | ||
| row_groups, | ||
| row_group_plan, |
There was a problem hiding this comment.
async build threads row_group_plan through but never calls into_global() — what happens if with_row_group_selections reaches here? A one-line test pinning either "works" or "clean error" would prevent a silent misbehave.
| /// ]) | ||
| /// # } | ||
| /// ``` | ||
| pub fn with_row_group_selections( |
There was a problem hiding this comment.
Worth adding tests while fresh: PerRowGroup + offset/limit, a duplicate RG index (doc says "decoded once per entry"), a selection shorter than its RG (trailing-skip), and empty vec![] (currently reads nothing).
| let Some(row_group_idx) = self.queued.front() else { | ||
| return Ok(None); | ||
| }; | ||
| if self.budget.is_exhausted() |
There was a problem hiding this comment.
nit: this early-exit only fires for Global; PerRowGroup drains via the selected_rows == 0 → continue path below — a short comment on that asymmetry would help.
2 participants
Which issue does this PR close?
Rationale for this change
DataFusion makes row-group-local selection decisions (
ParquetAccessPlan), but the reader APIs only accept selected row groups plus a single globalRowSelection. Callers must concatenate per-row-group selections into one global selection, which arrow-rs then re-partitions back into per-row-group selections during decoding. This round trip is wasted work and loses each selection's representation (bitmap vs. selector).What changes are included in this PR?
RowGroupSelection(a row group index plus an optional row-group-localRowSelection) andParquetPushDecoderBuilder::with_row_group_selections. Entries decode in the supplied order, omitted row groups are skipped,Nonereads the whole row group, and each selection keeps its bitmap or selector representation.with_row_groups/with_row_selection: the setters share an internal state machine (RowGroupPlan) that reports conflicting combinations as an error frombuild()regardless of call order. The legacy API combination is unchanged.build()validates per-row-group plans eagerly: out-of-bounds indices and selections longer than their row group are errors; shorter selections skip the trailing rows.ParquetPushDecoder::into_builderpreserves remaining local selections (still in local coordinates), so adaptive scans compose with the new API.with_row_groupson the push decoder now returns aParquetErrorduring decoding instead of panicking.The sync and async builders are unchanged; the async builder already delegates to the push decoder, so extending the API to it is a small follow-up if needed.
Are these changes tested?
Yes, new tests cover bitmap- and selector-backed local selections (including out-of-order row groups and short selections), skip/replace semantics, mutual exclusion in all four call orders, build-time validation,
into_builderround-trips, and the unchanged legacy combination. All existing tests pass.Are there any user-facing changes?
New public API:
RowGroupSelectionandParquetPushDecoderBuilder::with_row_group_selections, with doc examples. No breaking changes; one behavior change: out-of-boundswith_row_groupsindices on the push decoder now error during decoding instead of panicking.