feat: Add row-group-local RowSelection support to the push decoder by haohuaijin · Pull Request #10702 · apache/arrow-rs

haohuaijin · 2026-08-15T15:26:16Z

Which issue does this PR close?

Closes Support row-group-local row selections in the Parquet push decoder #10624.

Rationale for this change

DataFusion makes row-group-local selection decisions (ParquetAccessPlan), but the reader APIs only accept selected row groups plus a single global RowSelection. 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?

New public API on the push decoder: RowGroupSelection (a row group index plus an optional row-group-local RowSelection) and ParquetPushDecoderBuilder::with_row_group_selections. Entries decode in the supplied order, omitted row groups are skipped, None reads the whole row group, and each selection keeps its bitmap or selector representation.
Mutually exclusive with with_row_groups / with_row_selection: the setters share an internal state machine (RowGroupPlan) that reports conflicting combinations as an error from build() 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_builder preserves remaining local selections (still in local coordinates), so adaptive scans compose with the new API.
Minor behavior improvement: an out-of-bounds index from with_row_groups on the push decoder now returns a ParquetError during 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_builder round-trips, and the unchanged legacy combination. All existing tests pass.

Are there any user-facing changes?

New public API: RowGroupSelection and ParquetPushDecoderBuilder::with_row_group_selections, with doc examples. No breaking changes; one behavior change: out-of-bounds with_row_groups indices on the push decoder now error during decoding instead of panicking.

haohuaijin · 2026-08-15T15:42:08Z

In my mind the trick will be how to mesh the existing apis (with_row_groups and with_row_selection) and make sure we have a coherent documentation of what will happen if the user calls each one

This PR adds ParquetPushDecoderBuilder::with_row_group_selections so callers can supply row-group-local RowSelections directly. The new configuration is mutually exclusive with with_row_groups / with_row_selection: an internal state machine catches any conflicting combination in any call order and reports it as an error from build(), the legacy APIs keep their exact semantics, and each setter's docs spell out how it interacts with the others.

cc @alamb @zhuqi-lucas

zhuqi-lucas

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.

haohuaijin · 2026-08-17T09:36:02Z

Thanks for you reviews @zhuqi-lucas, i update in c0369a6

alamb

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

alamb · 2026-08-18T18:25:31Z

+
+/// Row-selection configuration shared by the Arrow reader builders.
+///
+/// `Global` is the legacy configuration formed by `with_row_groups` and


nit: it would be nice to put these docs directly on the enum variants, along with a description of what they mean, not just where they came from / what they are used for

Something like this, perhaps:

/// Row-selection configuration shared by the Arrow reader builders. pub(crate) enum RowGroupPlan { /// (IS THIS TRUE?) Selection: first select any row_groups, adn then apply the RowSelection to /// any remaining rows. /// formed by `with_row_groups` and ....) Global { ... }

(I am happy to make these changes myself, but I wanted to double check first)

alamb · 2026-08-18T18:28:10Z

+}
+
+impl RowGroupSelection {
+    /// Creates a row-group-local selection.


it might be good to define what row-group-local means more precisely -- specifically something like "The selection is relative to the rows in that selection (e.g. row 100 means the 100th row in the row group)"

alamb · 2026-08-18T18:55:35Z

+        // `row_groups` and `selection` straddle `projection`/`filter` so that the
+        // legacy configuration keeps its historical field order; only the
+        // row-group-local fields are new.
+        let (row_groups, selection, row_group_selections) = match &self.row_group_plan {


this feels overly complicated. Why not derive Debug for RowGroupPlan and display it as normal here

alamb · 2026-08-18T18:57:06Z

    pub(crate) batch_size: usize,

-    pub(crate) row_groups: Option<Vec<usize>>,
+    pub(crate) row_group_plan: RowGroupPlan,


alamb · 2026-08-18T18:59:03Z

-            row_groups: Some(row_groups),
-            ..self
-        }
+    ///


If we are going to keep this wording I suggest making async stream builders a link too -- something like

- /// On [`ParquetPushDecoderBuilder`] and async stream builders, which - /// additionally offer `with_row_group_selections`, this cannot be combined - /// with that method; attempting to do so returns an error from `build`. + /// On [`ParquetPushDecoderBuilder`] and [`ParquetRecordBatchStreamBuilder`], + /// which additionally offer `with_row_group_selections`, this cannot be + /// combined with that method; attempting to do so returns an error from + /// `build`. /// /// [`ParquetPushDecoderBuilder`]: crate::arrow::push_decoder::ParquetPushDecoderBuilder + /// [`ParquetRecordBatchStreamBuilder`]: crate::arrow::async_reader::ParquetRecordBatchStreamBuilder

alamb · 2026-08-18T19:13:30Z

            has_predicates: _,
        } = frontier;
+        let row_group_plan = match queued {
+            QueuedRowGroups::Global {


Can we also make this a method on QueuedRowGroups (like queued.into_plan() or something like that?

alamb · 2026-08-18T19:15:15Z

-    ) -> Self {
-        Self {
+    ) -> Result<Self, ParquetError> {
+        let queued = match row_group_plan {


This looks like it would more naturally be a constructor / from impl on RowGroupPlan -- to construct QueuedRowGroups from a RowGroupPlan

That would also encapsulate the compleixty more

alamb · 2026-08-18T19:17:38Z

+    /// error from [`Self::build`]. Calling this method more than once replaces
+    /// the previous row-group-local configuration.
+    ///
+    /// ```no_run


It took me some time to figure out (with some help) that we duplicated with_row_group_selections twice because it isn't supported via the serialized reader

alamb · 2026-08-18T19:19:25Z

+    /// than its row group skips the trailing rows, while a selection longer
+    /// than its row group returns an error from [`Self::build`]. Each
+    /// selection is passed through without being re-partitioned, so no
+    /// conversion between the bitmap and selector representations is forced


I think the mention of different representations of the selection is confusing here (it seems like an irrelevant implementation detail)

alamb · 2026-08-18T19:36:58Z

I also ran cargo llvm-cov and it reports several of the cases in the decoder aren't covered. This might be resolved by refactoring some of the access plans into their own methods

parquet: Add row-group-local RowSelection support to the push decoder

a9157e1

github-actions Bot added the parquet Changes to the parquet crate label Aug 15, 2026

haohuaijin changed the title ~~parquet: Add row-group-local RowSelection support to the push decoder~~ feat: Add row-group-local RowSelection support to the push decoder Aug 15, 2026

haohuaijin force-pushed the push-decoder-row-group-selections branch from 02cae11 to a9157e1 Compare August 16, 2026 03:24

haohuaijin mentioned this pull request Aug 16, 2026

feat: Add RowSelection::total_row_count #10704

Open

zhuqi-lucas reviewed Aug 16, 2026

View reviewed changes

Comment thread parquet/src/arrow/async_reader/mod.rs

zhuqi-lucas reviewed Aug 16, 2026

View reviewed changes

Comment thread parquet/src/arrow/push_decoder/mod.rs

zhuqi-lucas reviewed Aug 16, 2026

View reviewed changes

Comment thread parquet/src/arrow/push_decoder/remaining.rs

update

c0369a6

alamb approved these changes Aug 18, 2026

View reviewed changes

Uh oh!

Conversation

haohuaijin commented Aug 15, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

haohuaijin commented Aug 15, 2026

Uh oh!

zhuqi-lucas left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

haohuaijin commented Aug 17, 2026

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Aug 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zhuqi-lucas left a comment •

edited

Loading