From bf9336c1e918f5229d75e664c28ef6d1f205c75d Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Thu, 20 Aug 2026 15:23:47 +0800 Subject: [PATCH 1/2] fix(parquet): strip empty row groups in PreparedAccessPlan::prepare A row group whose `RowSelection` selects zero rows after page-index pruning is silently advanced past by arrow-rs's push decoder inside `try_next_reader`, without handing back a reader. The rest of DataFusion (per-RG metadata maps, the runtime dynamic pruner) assumes a 1:1 correspondence between the prepared plan and the readers the decoder hands back, so a silently-skipped empty row group leaves that bookkeeping off by one. Strip such empty row groups in `PreparedAccessPlan::prepare` so the prepared plan never contains a row group the decoder will skip. The flat `RowSelection` is split per row group (mirroring arrow-rs) and only non-empty segments are kept. Closes #24287. --- .../datasource-parquet/src/access_plan.rs | 86 +++++++++++++++++++ datafusion/datasource-parquet/src/sort.rs | 24 +++--- 2 files changed, 97 insertions(+), 13 deletions(-) diff --git a/datafusion/datasource-parquet/src/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs index 1e9bae0ff6ba3..1aaa3a846ed64 100644 --- a/datafusion/datasource-parquet/src/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -573,10 +573,61 @@ 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) } } +/// Strip row groups whose post-pruning `RowSelection` selects zero rows. +/// +/// arrow-rs's push decoder silently advances past such row groups inside +/// `try_next_reader`, but the rest of DataFusion (per-RG metadata maps and the +/// runtime dynamic-pruner) assumes a 1:1 correspondence between the prepared +/// plan and the readers the decoder hands back. Removing these empty entries +/// here keeps that invariant so downstream per-RG bookkeeping stays in sync +/// with the decoder. +/// +/// The flat `RowSelection` is split per row group with +/// [`RowSelection::split_off`] (mirroring arrow-rs's own logic) and the +/// surviving segments are concatenated back into the result selection. When +/// `row_selection` is `None` (no page-index pruning, no user-supplied +/// selection) no row group can be empty and the inputs are returned unchanged. +fn strip_empty_row_groups( + row_group_indexes: Vec, + row_selection: Option, + row_group_meta_data: &[RowGroupMetaData], +) -> (Vec, Option) { + let Some(mut remaining) = row_selection else { + return (row_group_indexes, None); + }; + + let mut kept_indexes = Vec::with_capacity(row_group_indexes.len()); + let mut kept_selectors: Vec = Vec::new(); + + for &rg_idx in row_group_indexes.iter() { + let rg_row_count = row_group_meta_data[rg_idx].num_rows() as usize; + // `split_off` cuts off the first `rg_row_count` rows worth of + // selection — this row group's segment. `remaining` keeps the rest. + let rg_segment = remaining.split_off(rg_row_count); + if rg_segment.row_count() > 0 { + kept_indexes.push(rg_idx); + kept_selectors.extend(rg_segment.iter().copied()); + } + // Empty segment ⇒ arrow-rs would have silently skipped this row group + // anyway; drop it from our plan so per-RG bookkeeping stays in sync. + } + + 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 @@ -1049,6 +1100,41 @@ mod test { /// [`RowGroupMetaData`] that returns 4 row groups with 10, 20, 30, 40 rows /// respectively + #[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_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()); + } + static ROW_GROUP_METADATA: LazyLock> = LazyLock::new(|| { let schema_descr = get_test_schema_descr(); let row_counts = [10, 20, 30, 40]; diff --git a/datafusion/datasource-parquet/src/sort.rs b/datafusion/datasource-parquet/src/sort.rs index ea33fb0e2ecb2..0f73723a1de91 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] From dd773704e0e22acaa37749e73da76c827f0ada15 Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Thu, 20 Aug 2026 22:47:14 +0800 Subject: [PATCH 2/2] Address @adriangb review: single-pass cursor, reframe as hardening, add reachable-case + non-contiguous tests - Rewrite strip_empty_row_groups to walk the flat selection once with OverallRowSelectionCursor instead of per-row-group RowSelection::split_off (which reallocated the selector tail on each call). - Reframe the doc comment around the well-formed-plan invariant and the one reachable producer (scan_selection intersecting to empty), rather than a runtime-pruner desync that is gated off today. - Move the unit tests below ROW_GROUP_METADATA so its doc comment no longer attaches to a test. - Add a non-contiguous (vec![1, 3]) test that pins absolute indexing of row_group_meta_data, and a regression test through prepare() exercising the scan_selection empty-intersection path. --- .../datasource-parquet/src/access_plan.rs | 184 ++++++++++++------ 1 file changed, 129 insertions(+), 55 deletions(-) diff --git a/datafusion/datasource-parquet/src/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs index 1aaa3a846ed64..6d96c68e18cb0 100644 --- a/datafusion/datasource-parquet/src/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -580,43 +580,65 @@ impl ParquetAccessPlan { } } -/// Strip row groups whose post-pruning `RowSelection` selects zero rows. +/// 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 such row groups inside -/// `try_next_reader`, but the rest of DataFusion (per-RG metadata maps and the -/// runtime dynamic-pruner) assumes a 1:1 correspondence between the prepared -/// plan and the readers the decoder hands back. Removing these empty entries -/// here keeps that invariant so downstream per-RG bookkeeping stays in sync -/// with the decoder. +/// 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. /// -/// The flat `RowSelection` is split per row group with -/// [`RowSelection::split_off`] (mirroring arrow-rs's own logic) and the -/// surviving segments are concatenated back into the result selection. When -/// `row_selection` is `None` (no page-index pruning, no user-supplied -/// selection) no row group can be empty and the inputs are returned unchanged. +/// 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, row_selection: Option, row_group_meta_data: &[RowGroupMetaData], ) -> (Vec, Option) { - let Some(mut remaining) = row_selection else { + 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 = Vec::new(); for &rg_idx in row_group_indexes.iter() { let rg_row_count = row_group_meta_data[rg_idx].num_rows() as usize; - // `split_off` cuts off the first `rg_row_count` rows worth of - // selection — this row group's segment. `remaining` keeps the rest. - let rg_segment = remaining.split_off(rg_row_count); - if rg_segment.row_count() > 0 { + // 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); - kept_selectors.extend(rg_segment.iter().copied()); + } 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); } - // Empty segment ⇒ arrow-rs would have silently skipped this row group - // anyway; drop it from our plan so per-RG bookkeeping stays in sync. } let result_selection = if kept_selectors.is_empty() { @@ -1100,41 +1122,6 @@ mod test { /// [`RowGroupMetaData`] that returns 4 row groups with 10, 20, 30, 40 rows /// respectively - #[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_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()); - } - static ROW_GROUP_METADATA: LazyLock> = LazyLock::new(|| { let schema_descr = get_test_schema_descr(); let row_counts = [10, 20, 30, 40]; @@ -1516,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]); + } }