diff --git a/parquet/src/arrow/push_decoder/mod.rs b/parquet/src/arrow/push_decoder/mod.rs
index 6dc5520bb975..ba075b4c08b6 100644
--- a/parquet/src/arrow/push_decoder/mod.rs
+++ b/parquet/src/arrow/push_decoder/mod.rs
@@ -538,6 +538,30 @@ impl ParquetPushDecoder {
self.state.row_groups_remaining()
}
+ /// Returns the row-group index that the next call to
+ /// [`Self::try_next_reader`] will yield a reader for, after applying
+ /// any internal skipping (row selection emptiness, exhausted budget,
+ /// finished state).
+ ///
+ /// Safe to call at any time. When called mid-row-group (i.e. while a
+ /// previously-emitted [`ParquetRecordBatchReader`] is still being
+ /// drained), the returned index refers to the row group that
+ /// [`Self::try_next_reader`] will produce *after* the current one.
+ ///
+ /// Returns `Ok(None)` when:
+ /// - the decoder has no more row groups to read, or
+ /// - every remaining row group would be skipped.
+ ///
+ /// This method does not mutate decoder state. It is useful for
+ /// callers that maintain per-row-group state in lock-step with the
+ /// decoder (e.g. dynamic row-group pruners) to determine which row
+ /// group the next reader corresponds to, since
+ /// [`Self::try_next_reader`] may silently advance past row groups
+ /// based on filtering and other criteria.
+ pub fn peek_next_row_group(&self) -> Result, ParquetError> {
+ self.state.peek_next_row_group()
+ }
+
/// Decompose this decoder back into a [`ParquetPushDecoderBuilder`] for the
/// row groups that have *not* yet been decoded.
///
@@ -840,6 +864,26 @@ impl ParquetDecoderState {
}
}
+ /// See [`ParquetPushDecoder::peek_next_row_group`] for the public
+ /// API contract. This inner method delegates to the underlying
+ /// [`RemainingRowGroups`] for both `ReadingRowGroup` and
+ /// `DecodingRowGroup`: mid-row-group the answer is the row group
+ /// `try_next_reader` will produce *after* the active one finishes,
+ /// which is exactly what `RemainingRowGroups::peek_next_row_group`
+ /// computes from the queued frontier.
+ fn peek_next_row_group(&self) -> Result , ParquetError> {
+ match self {
+ ParquetDecoderState::ReadingRowGroup {
+ remaining_row_groups,
+ } => remaining_row_groups.peek_next_row_group(),
+ ParquetDecoderState::DecodingRowGroup {
+ remaining_row_groups,
+ ..
+ } => remaining_row_groups.peek_next_row_group(),
+ ParquetDecoderState::Finished => Ok(None),
+ }
+ }
+
fn into_builder(self) -> Result {
let remaining_row_groups = match self {
ParquetDecoderState::ReadingRowGroup {
@@ -1743,6 +1787,206 @@ mod test {
expect_finished(decoder.try_decode());
}
+ /// `peek_next_row_group` reports the index of the row group the
+ /// next `try_next_reader` call will hand back, matching the
+ /// frontier's internal skip logic.
+ #[test]
+ fn test_peek_next_row_group_basic() {
+ let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .build()
+ .unwrap();
+
+ // Two row groups (0, 1). At boundary before any read, peek should
+ // see RG 0.
+ assert_eq!(decoder.peek_next_row_group().unwrap(), Some(0));
+ assert!(decoder.is_at_row_group_boundary());
+
+ let ranges = expect_needs_data(decoder.try_next_reader());
+ push_ranges_to_decoder(&mut decoder, ranges);
+ let reader = expect_data(decoder.try_next_reader());
+ // Once the reader for RG 0 has been handed off, the decoder is
+ // back at a boundary waiting for RG 1 — peek must reflect that
+ // (the active reader lives outside the decoder).
+ assert!(decoder.is_at_row_group_boundary());
+ assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
+
+ // Drain RG 0's reader and consume RG 1.
+ for batch in reader {
+ let _ = batch.unwrap();
+ }
+ let ranges = expect_needs_data(decoder.try_next_reader());
+ push_ranges_to_decoder(&mut decoder, ranges);
+ let reader = expect_data(decoder.try_next_reader());
+ for batch in reader {
+ let _ = batch.unwrap();
+ }
+
+ // No row groups left.
+ assert_eq!(decoder.peek_next_row_group().unwrap(), None);
+ }
+
+ /// `peek_next_row_group` honors `with_row_groups` — restricting the
+ /// scan to a single row group means peek reports only that one and
+ /// then `None`.
+ #[test]
+ fn test_peek_next_row_group_respects_with_row_groups() {
+ let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .with_row_groups(vec![1])
+ .build()
+ .unwrap();
+
+ assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
+ }
+
+ /// When a row-selection segment leaves the next row group with zero
+ /// selected rows, `peek_next_row_group` mirrors
+ /// `next_readable_row_group`'s skip: it returns the *following*
+ /// row group instead of the empty one.
+ #[test]
+ fn test_peek_next_row_group_skips_empty_selection() {
+ // Each row group has 200 rows. Skip all 200 of RG 0 plus 50 of
+ // RG 1; the next reader will be for RG 1, not RG 0.
+ let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .with_row_selection(RowSelection::from(vec![
+ RowSelector::skip(250),
+ RowSelector::select(100),
+ ]))
+ .build()
+ .unwrap();
+
+ assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
+ }
+
+ /// `peek_next_row_group` returns `None` on a finished decoder.
+ #[test]
+ fn test_peek_next_row_group_finished() {
+ let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .with_row_groups(vec![])
+ .build()
+ .unwrap();
+
+ // No row groups requested ⇒ already finished, no peek.
+ expect_finished(decoder.try_next_reader());
+ assert_eq!(decoder.peek_next_row_group().unwrap(), None);
+ }
+
+ /// `peek_next_row_group` mirrors `next_readable_row_group`'s
+ /// budget-skipping logic: with an `OFFSET` that consumes the
+ /// entire first row group, peek must return the *following*
+ /// row group, not RG 0 (no predicates, so budget skips apply).
+ #[test]
+ fn test_peek_next_row_group_skips_for_budget() {
+ // Each row group has 200 rows. OFFSET 200 drains RG 0 entirely
+ // and lands the decoder's first emitted reader at RG 1.
+ let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .with_offset(200)
+ .build()
+ .unwrap();
+
+ assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
+ }
+
+ /// OFFSET larger than every remaining row group's row count
+ /// exhausts the budget, so peek should report `None` — matching
+ /// `next_readable_row_group`'s behavior of producing `Finished`.
+ #[test]
+ fn test_peek_next_row_group_budget_drains_all() {
+ // 2 row groups × 200 rows = 400 total. OFFSET 500 cannot land
+ // anywhere — every RG is skipped under the budget.
+ let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .with_offset(500)
+ .build()
+ .unwrap();
+
+ assert_eq!(decoder.peek_next_row_group().unwrap(), None);
+ }
+
+ /// `peek_next_row_group` is safe to call while a row group's reader
+ /// is still being drained via `try_decode`. In `DecodingRowGroup`
+ /// state it must report the row group `try_next_reader` will yield
+ /// a reader for *after* the active one — never `None` just because
+ /// the decoder is mid-row-group.
+ #[test]
+ fn test_peek_next_row_group_during_decoding_row_group() {
+ // Two row groups (200 rows each), small batch size so the
+ // decoder stays in `DecodingRowGroup` across multiple
+ // `try_decode` calls within RG 0.
+ let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .with_batch_size(100)
+ .build()
+ .unwrap();
+ decoder
+ .push_range(test_file_range(), TEST_FILE_DATA.clone())
+ .unwrap();
+
+ // First batch of RG 0 — after this, the decoder is in
+ // `DecodingRowGroup` state with the reader retained internally
+ // and one more 100-row batch still owed for RG 0. Peek must
+ // therefore look past the active RG to RG 1.
+ let _batch0 = expect_data(decoder.try_decode());
+ assert!(!decoder.is_at_row_group_boundary());
+ assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
+
+ // Drain the rest of RG 0; peek still reports RG 1.
+ let _batch1 = expect_data(decoder.try_decode());
+ assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
+
+ // Move into RG 1 — peek now sees no further row groups.
+ let _batch2 = expect_data(decoder.try_decode());
+ assert!(!decoder.is_at_row_group_boundary());
+ assert_eq!(decoder.peek_next_row_group().unwrap(), None);
+
+ let _batch3 = expect_data(decoder.try_decode());
+ expect_finished(decoder.try_decode());
+ }
+
+ /// Peeking is a read-only operation: calling it repeatedly between
+ /// `try_next_reader` calls must never change which row group the
+ /// reader path actually produces. Drives the decoder all the way
+ /// through and asserts each peek/read pair agrees.
+ #[test]
+ fn test_peek_next_row_group_does_not_mutate_state() {
+ let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
+ .unwrap()
+ .build()
+ .unwrap();
+
+ // Two row groups expected — drive both, asserting peek/read agree.
+ for expected_rg in [0usize, 1usize] {
+ assert!(decoder.is_at_row_group_boundary());
+
+ // Multiple peeks before reading must all agree, and must not
+ // disturb the upcoming read.
+ let first_peek = decoder.peek_next_row_group().unwrap();
+ let second_peek = decoder.peek_next_row_group().unwrap();
+ let third_peek = decoder.peek_next_row_group().unwrap();
+ assert_eq!(first_peek, Some(expected_rg));
+ assert_eq!(first_peek, second_peek);
+ assert_eq!(second_peek, third_peek);
+
+ // Now read for real and confirm the decoder hands back exactly
+ // what peek promised.
+ let ranges = expect_needs_data(decoder.try_next_reader());
+ push_ranges_to_decoder(&mut decoder, ranges);
+ let reader = expect_data(decoder.try_next_reader());
+ for batch in reader {
+ let _ = batch.unwrap();
+ }
+ }
+
+ // Decoder is drained. Peek must agree with `try_next_reader`'s
+ // terminal state.
+ assert_eq!(decoder.peek_next_row_group().unwrap(), None);
+ expect_finished(decoder.try_next_reader());
+ }
+
/// `into_builder` between row groups recovers a builder for the
/// not-yet-decoded row groups; rebuilding it with a new row filter
/// applies that filter to the subsequent row groups while leaving the
diff --git a/parquet/src/arrow/push_decoder/remaining.rs b/parquet/src/arrow/push_decoder/remaining.rs
index d1070d2aa69f..d2658189b7c1 100644
--- a/parquet/src/arrow/push_decoder/remaining.rs
+++ b/parquet/src/arrow/push_decoder/remaining.rs
@@ -49,7 +49,7 @@ struct NextRowGroup {
budget: RowBudget,
}
-#[derive(Debug)]
+#[derive(Debug, Clone)]
struct RowGroupFrontier {
/// Metadata used to resolve row counts for queued row groups.
parquet_metadata: Arc,
@@ -93,6 +93,23 @@ impl RowGroupFrontier {
self.budget = budget;
}
+ /// Peek at the next row-group index [`Self::next_readable_row_group`]
+ /// would hand out, without mutating any state. Returns `None` if every
+ /// remaining row group would be skipped under the current
+ /// selection/budget, or if the queue is empty.
+ ///
+ /// Runs the real [`Self::next_readable_row_group`] advance logic on a
+ /// throwaway clone of the frontier, so peek can never drift from the
+ /// read path. The clone copies the queued row-group indices and optional
+ /// row-selection (a `Vec`); see
+ /// [`RemainingRowGroups::peek_next_row_group`].
+ fn peek_next_row_group(&self) -> Result, ParquetError> {
+ Ok(self
+ .clone()
+ .next_readable_row_group()?
+ .map(|next_row_group| next_row_group.row_group_idx))
+ }
+
fn clear_remaining(&mut self) {
self.selection = None;
self.row_groups.clear();
@@ -299,6 +316,28 @@ impl RemainingRowGroups {
self.frontier.row_groups.len()
}
+ /// Peek at the file-level row-group index that the next call to
+ /// [`Self::try_next_reader`] will produce a reader for, after
+ /// simulating the same skip logic [`Self::try_next_reader`] applies
+ /// internally (row-selection emptiness + offset/limit budget). Does
+ /// not mutate state.
+ ///
+ /// Returns `None` when the active row group is still being decoded,
+ /// when no row groups remain, or when every remaining row group
+ /// would be skipped under the current selection/budget.
+ ///
+ /// Cost: one clone of the queued row-group indices and optional
+ /// row-selection per call (the frontier is cloned so the real advance
+ /// logic can run non-destructively). For callers that peek once per
+ /// row-group boundary this is O(remaining row groups + selectors) per
+ /// boundary.
+ pub fn peek_next_row_group(&self) -> Result , ParquetError> {
+ if self.row_group_reader_builder.has_active_row_group() {
+ return Ok(None);
+ }
+ self.frontier.peek_next_row_group()
+ }
+
/// returns [`ParquetRecordBatchReader`] suitable for reading the next
/// group of rows from the Parquet data, or the list of data ranges still
/// needed to proceed