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 8 commits 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.
|
Thanks for you reviews @zhuqi-lucas, i update in c0369a6 |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thank you @haohuaijin and @zhuqi-lucas -- I spent quite a while reading this one in detail; I think it is looking very good -- I had some comment and encapsulation nits which I would like to fix before merging but I don't think they are strictly required
|
I also ran |
|
Thanks for your details reviews and suggestion @alamb , i apply all suggestion and fix coverage issue. |
# Which issue does this PR close? - Part of #10624. # Rationale for this change `RowSelection` exposes `row_count()` (selected rows) and `skipped_row_count()` (skipped rows), but no way to get the total number of rows a selection spans. Callers that need the total (e.g. to validate a selection against a row group's row count, as #10702 does) must call both methods, which iterates a selector-backed selection twice and performs two popcounts on a mask-backed selection just to have them cancel out. # What changes are included in this PR? Adds `RowSelection::total_row_count()`, which computes the total in a single pass: - selector-backed: one sum over the selectors - mask-backed: `mask.len()`, O(1) with no popcount # Are these changes tested? Yes, a unit test covers both backings plus the empty selection. # Are there any user-facing changes? New public method `RowSelection::total_row_count()`. No changes to existing APIs. --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
|
update code to use #10704 |
# Which issue does this PR close? - Part of apache#10624. # Rationale for this change `RowSelection` exposes `row_count()` (selected rows) and `skipped_row_count()` (skipped rows), but no way to get the total number of rows a selection spans. Callers that need the total (e.g. to validate a selection against a row group's row count, as apache#10702 does) must call both methods, which iterates a selector-backed selection twice and performs two popcounts on a mask-backed selection just to have them cancel out. # What changes are included in this PR? Adds `RowSelection::total_row_count()`, which computes the total in a single pass: - selector-backed: one sum over the selectors - mask-backed: `mask.len()`, O(1) with no popcount # Are these changes tested? Yes, a unit test covers both backings plus the empty selection. # Are there any user-facing changes? New public method `RowSelection::total_row_count()`. No changes to existing APIs. --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
3 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.