This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24509-f239959350080b5bfe4403d4de5e83bcfc03dc0b in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 1979d7e7943d8a640c32263b03006a5c762e0bfc Author: Qi Zhu <[email protected]> AuthorDate: Thu Aug 20 16:42:34 2026 +0000 refactor(parquet): strip empty row groups so PreparedAccessPlan stays well-formed (#24509) ## Which issue does this PR close? - Closes #24287. ## Rationale for this change Split out from #23696 (per @adriangb's [decomposition proposal](https://github.com/apache/datafusion/pull/23696#issuecomment-5346992228)). This keeps `PreparedAccessPlan` well-formed: a row group whose `RowSelection` selects zero rows is dropped so `with_row_groups(...)` never names a row group the push decoder would immediately skip. arrow-rs advances past an all-skipped row group inside `try_next_reader` without handing back a reader, so keeping such a row group in the plan would leave the plan's row-group list one entry longer than the readers produced. This is mostly **hardening** — the two in-tree flat-selection producers already avoid landing an empty `Selection` (`page_filter.rs` maps an all-skip selection to `Skip`; `try_new_from_overall_row_selection` normalizes via `RowGroupAccessBuilder::into_access`), and `RowGroupPruner` is gated off whenever a selection is live. There is, however, **one reachable producer**: `ParquetAccessPlan::scan_selection` intersects an existing `Selection` with a new one and can leave a row group selecting nothing (its `rows_selected > 0` guard checks the incoming selection, not the intersection). A well-formed plan here is also a precondition for re-enabling runtime pruning under a live selection (#24358) and for #23696. ## What changes are included in this PR? - `access_plan.rs`: `PreparedAccessPlan::prepare` calls a new generic `strip_empty_row_groups(row_group_indexes, row_selection, row_group_meta_data)`. It walks the flat `RowSelection` once with `OverallRowSelectionCursor` (the same single-pass approach `try_new_from_overall_row_selection` uses, rather than a per-row-group `split_off`) and keeps only row groups that still select a row. - `sort.rs`: `test_prepared_access_plan_reverse_empty_selection` — an all-skipped plan now strips to an empty plan. No `fully_matched` coupling, no public API change, no new metric, no slt churn. ## Are these changes tested? Yes — unit tests in `access_plan.rs`: - `test_strip_empty_row_groups_drops_only_empties` — mixed empty/non-empty row groups. - `test_strip_empty_row_groups_non_contiguous_indexes` — `vec![1, 3]` (RGs 0/2 skipped), pinning that `row_group_meta_data` is indexed absolutely. - `test_strip_empty_row_groups_none_selection_unchanged` — `None` passes through. - `test_prepare_strips_row_group_emptied_by_intersecting_selections` — the reachable `scan_selection` empty-intersection path, exercised through `prepare()`. ## Are there any user-facing changes? No. ## Follow-up @adriangb also suggested a complementary one-line normalization in `scan_selection` itself (empty intersection → `Skip`), which would keep `ParquetAccessPlan` well-formed for its other consumers (`row_group_indexes()`, metrics, `is_fully_matched`) too. Left out here to keep this PR scoped to #24287; happy to do it as a small follow-up. --- datafusion/datasource-parquet/src/access_plan.rs | 160 +++++++++++++++++++++++ datafusion/datasource-parquet/src/sort.rs | 24 ++-- 2 files changed, 171 insertions(+), 13 deletions(-) diff --git a/datafusion/datasource-parquet/src/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs index 1e9bae0ff6..6d96c68e18 100644 --- a/datafusion/datasource-parquet/src/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -573,10 +573,83 @@ impl ParquetAccessPlan { let row_group_indexes = self.row_group_indexes(); let row_selection = self.into_overall_row_selection(row_group_meta_data)?; + let (row_group_indexes, row_selection) = + strip_empty_row_groups(row_group_indexes, row_selection, row_group_meta_data); + PreparedAccessPlan::new(row_group_indexes, row_selection) } } +/// Drop row groups whose post-pruning `RowSelection` selects zero rows, so the +/// prepared plan stays well-formed: `with_row_groups(...)` never names a row +/// group the decoder would immediately skip. +/// +/// arrow-rs's push decoder silently advances past an all-skipped row group +/// inside `try_next_reader` without handing back a reader, so keeping such a +/// row group in the plan would leave its row-group list one entry longer than +/// the readers the decoder produces. An empty selection is reachable today via +/// [`ParquetAccessPlan::scan_selection`], which intersects an existing +/// `Selection` with a new one and can leave a row group selecting nothing. +/// A well-formed plan here is also a precondition for re-enabling runtime +/// pruning under a live selection (#24358) and for #23696. +/// +/// Walks the flat `RowSelection` once with [`OverallRowSelectionCursor`], +/// rather than a per-row-group [`RowSelection::split_off`] (which reallocates +/// the selector tail on every call, i.e. O(row groups × selectors)), keeping +/// only the row groups that still select at least one row. `row_group_meta_data` +/// is the **full file** metadata, indexed absolutely by row-group index, while +/// `row_selection` covers only the scanned row groups — matching +/// `into_overall_row_selection`, which emits nothing for a skipped row group. +/// When `row_selection` is `None` no row group can be empty and the inputs are +/// returned unchanged. +fn strip_empty_row_groups( + row_group_indexes: Vec<usize>, + row_selection: Option<RowSelection>, + row_group_meta_data: &[RowGroupMetaData], +) -> (Vec<usize>, Option<RowSelection>) { + let Some(selection) = row_selection else { + return (row_group_indexes, None); + }; + + let mut cursor = OverallRowSelectionCursor::new(selection); + let mut kept_indexes = Vec::with_capacity(row_group_indexes.len()); + let mut kept_selectors: Vec<RowSelector> = Vec::new(); + + for &rg_idx in row_group_indexes.iter() { + let rg_row_count = row_group_meta_data[rg_idx].num_rows() as usize; + // Pull this row group's fragments off the shared cursor in a single + // pass (no `split_off` tail reallocation). + let start = kept_selectors.len(); + let mut selected = 0usize; + let mut taken = 0usize; + while taken < rg_row_count { + let Some(fragment) = cursor.take(rg_row_count - taken) else { + break; + }; + taken += fragment.row_count; + if !fragment.skip { + selected += fragment.row_count; + } + kept_selectors.push(fragment); + } + if selected > 0 { + kept_indexes.push(rg_idx); + } else { + // Empty row group: arrow-rs would silently skip it. Drop it and its + // fragments so the plan stays 1:1 with the readers. + kept_selectors.truncate(start); + } + } + + let result_selection = if kept_selectors.is_empty() { + None + } else { + Some(RowSelection::from(kept_selectors)) + }; + + (kept_indexes, result_selection) +} + /// Represents a prepared, fully resolved [`ParquetAccessPlan`] /// /// The [`RowSelection`] represents the result of applying all pruning such as @@ -1430,4 +1503,91 @@ mod test { // Ordered by min(a) ASC only: 3, 4, 5. assert_eq!(result.row_group_indexes, vec![1, 2, 0]); } + + #[test] + fn test_strip_empty_row_groups_drops_only_empties() { + // 4 row groups of [10, 20, 30, 40] rows. RG 1 and RG 3 select nothing + // after pruning, so they must be dropped; RG 0 and RG 2 survive with + // their re-concatenated selections intact. + let selection = RowSelection::from(vec![ + RowSelector::select(10), // RG 0: keep all 10 + RowSelector::skip(30), // RG 1 (20) fully skipped + RG 2's leading 10 + RowSelector::select(20), // RG 2: keep 20 + RowSelector::skip(40), // RG 3: skip all 40 + ]); + + let (indexes, result) = strip_empty_row_groups( + vec![0, 1, 2, 3], + Some(selection), + &ROW_GROUP_METADATA, + ); + + // RG 1 and RG 3 are dropped; the surviving indexes stay in order. + assert_eq!(indexes, vec![0, 2]); + let result = result.expect("survivors keep a selection"); + assert_eq!(result.row_count(), 30); // 10 from RG 0 + 20 from RG 2 + assert_eq!(result.skipped_row_count(), 10); // RG 2's leading skip + } + + #[test] + fn test_strip_empty_row_groups_non_contiguous_indexes() { + // Only RG 1 (20 rows) and RG 3 (40 rows) are scanned — RG 0 and RG 2 + // are whole-group skips, so the selection covers 20 + 40 = 60 rows. + // This pins down that `strip` indexes `row_group_meta_data` + // *absolutely* (meta[1]=20, meta[3]=40): a relative reading would take + // 10 then 20 rows and misalign the split. + let selection = RowSelection::from(vec![ + RowSelector::skip(20), // RG 1: skip all 20 -> dropped + RowSelector::select(40), // RG 3: keep all 40 + ]); + + let (indexes, result) = + strip_empty_row_groups(vec![1, 3], Some(selection), &ROW_GROUP_METADATA); + + assert_eq!(indexes, vec![3]); + let result = result.expect("RG 3 survives"); + assert_eq!(result.row_count(), 40); + assert_eq!(result.skipped_row_count(), 0); + } + + #[test] + fn test_strip_empty_row_groups_none_selection_unchanged() { + // With no row selection no row group can be empty, so the inputs pass + // through unchanged. + let (indexes, result) = + strip_empty_row_groups(vec![0, 1, 2, 3], None, &ROW_GROUP_METADATA); + assert_eq!(indexes, vec![0, 1, 2, 3]); + assert!(result.is_none()); + } + + #[test] + fn test_prepare_strips_row_group_emptied_by_intersecting_selections() { + // The reachable producer: `scan_selection` intersects two disjoint + // selections within RG 1 to nothing, leaving an empty `Selection` in + // the plan (its own `rows_selected > 0` guard checks the incoming + // selection, not the intersection). `prepare` must strip that row + // group so the prepared plan never names a row group the decoder would + // silently skip. + let mut plan = ParquetAccessPlan::new_all(4); // RGs [10, 20, 30, 40] + + // RG 1 (20 rows): select rows 0..10, then intersect with rows 10..20. + plan.scan_selection( + 1, + RowSelection::from(vec![RowSelector::select(10), RowSelector::skip(10)]), + ); + plan.scan_selection( + 1, + RowSelection::from(vec![RowSelector::skip(10), RowSelector::select(10)]), + ); + // RG 2 keeps a genuine partial selection so a flat selection is emitted. + plan.scan_selection( + 2, + RowSelection::from(vec![RowSelector::skip(10), RowSelector::select(20)]), + ); + + let prepared = plan.prepare(&ROW_GROUP_METADATA).expect("prepare"); + + // RG 1 (emptied by the intersection) is stripped; RG 0/2/3 remain. + assert_eq!(prepared.row_group_indexes, vec![0, 2, 3]); + } } diff --git a/datafusion/datasource-parquet/src/sort.rs b/datafusion/datasource-parquet/src/sort.rs index ea33fb0e2e..0f73723a1d 100644 --- a/datafusion/datasource-parquet/src/sort.rs +++ b/datafusion/datasource-parquet/src/sort.rs @@ -407,7 +407,8 @@ mod tests { #[test] fn test_prepared_access_plan_reverse_empty_selection() { - // Test: all rows are skipped + // Test: all rows are skipped. After `strip_empty_row_groups`, every row + // group's selection is empty, so the whole plan strips to nothing. let metadata = create_test_metadata(vec![100, 100, 100]); let mut access_plan = ParquetAccessPlan::new_all(3); @@ -423,21 +424,18 @@ mod tests { .prepare(rg_metadata) .expect("Failed to create PreparedAccessPlan"); + // All row groups are empty after pruning, so they are stripped and the + // prepared plan is empty (rather than carrying a selection that skips + // every row). + assert!(prepared_plan.row_group_indexes.is_empty()); + assert!(prepared_plan.row_selection.is_none()); + + // Reversing an empty plan stays empty. let reversed_plan = prepared_plan .reverse(&metadata) .expect("Failed to reverse PreparedAccessPlan"); - - // Should still skip all rows - let total_selected: usize = reversed_plan - .row_selection - .as_ref() - .unwrap() - .iter() - .filter(|s| !s.skip) - .map(|s| s.row_count) - .sum(); - - assert_eq!(total_selected, 0); + assert!(reversed_plan.row_group_indexes.is_empty()); + assert!(reversed_plan.row_selection.is_none()); } #[test] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
