refactor(parquet): strip empty row groups so PreparedAccessPlan stays well-formed

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]);
+    }
 }

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]