feat: TopK cumulative RG pruning after reorder (works with WHERE)

Instead of threshold-based pruning (which fails with WHERE clauses
due to unknown qualifying row counts), use cumulative row counting:
after reorder + reverse, accumulate rows from the front until we
have enough for the TopK fetch limit (K), then prune the rest.

This works for sort pushdown with or without WHERE because it only
depends on row counts + RG ordering, not threshold values or types.

Adds fetch field to DynamicFilterPhysicalExpr (set by SortExec) so
the parquet reader knows the TopK K value.

Keeps stats init for the no-WHERE sort pushdown case (20-58x speedup)
as a complementary optimization that also helps cross-file pruning
via the shared DynamicFilter.

Author: zhuqi-lucas

datafusion/datasource-parquet/src/access_plan.rs

@@ -381,6 +381,22 @@ impl PreparedAccessPlan {
         })
     }
 
+    /// Return a reference to the row group indexes.
+    pub(crate) fn row_group_indexes(&self) -> &[usize] {
+        &self.row_group_indexes
+    }
+
+    /// Keep only the first `count` row groups, dropping the rest.
+    /// Used for TopK cumulative pruning after reorder + reverse.
+    pub(crate) fn truncate_row_groups(mut self, count: usize) -> Self {
+        self.row_group_indexes.truncate(count);
+        // Clear row_selection since it's tied to the original RG set
+        if self.row_selection.is_some() {
+            self.row_selection = None;
+        }
+        self
+    }
+
     /// Reorder row groups by their min statistics for the given sort order.
     ///
     /// This helps TopK queries find optimal values first. For ASC sort,

datafusion/datasource-parquet/src/opener.rs

@@ -1222,6 +1222,36 @@ impl RowGroupsPrunedParquetOpen {
             )?;
         }
 
+        // TopK cumulative pruning: after reorder + reverse, the RGs are in
+        // optimal order for the sort. Keep accumulating rows from the front
+        // until we have enough for the TopK fetch limit (K), then prune the
+        // rest. Works for sort pushdown with or without WHERE because it only
+        // depends on row counts + ordering, not threshold values.
+        if prepared.sort_order_for_reorder.is_some()
+            && let Some(predicate) = &prepared.predicate
+            && let Some(df) = find_dynamic_filter(predicate)
+            && let Some(fetch) = df.fetch()
+        {
+            let rg_indexes = prepared_plan.row_group_indexes();
+            let mut cumulative = 0usize;
+            let mut keep_count = 0;
+            for &idx in rg_indexes {
+                cumulative += file_metadata.row_group(idx).num_rows() as usize;
+                keep_count += 1;
+                if cumulative >= fetch {
+                    break;
+                }
+            }
+            if keep_count < rg_indexes.len() {
+                let pruned = rg_indexes.len() - keep_count;
+                debug!(
+                    "TopK cumulative prune: keeping {keep_count} of {} RGs ({cumulative} rows >= fetch={fetch}), pruning {pruned}",
+                    rg_indexes.len()
+                );
+                prepared_plan = prepared_plan.truncate_row_groups(keep_count);
+            }
+        }
+
         let arrow_reader_metrics = ArrowReaderMetrics::enabled();
         let read_plan = build_projection_read_plan(
             prepared.projection.expr_iter(),
@@ -3343,6 +3373,7 @@ mod test {
             vec![Arc::clone(&col_expr)],
             super::lit(true),
             vec![desc_opts],
+            Some(10),
         ));
         let predicate: Arc<dyn PhysicalExpr> = dynamic_filter.clone();
 
@@ -3389,6 +3420,7 @@ mod test {
             vec![Arc::clone(&col_expr)],
             super::lit(true),
             vec![asc_opts],
+            Some(10),
         ));
         let predicate: Arc<dyn PhysicalExpr> = dynamic_filter.clone();
 
@@ -3480,6 +3512,7 @@ mod test {
             vec![col1, col2],
             super::lit(true),
             vec![opts, opts],
+            Some(10),
         ));
         let predicate: Arc<dyn PhysicalExpr> = dynamic_filter.clone();
 

datafusion/physical-expr/src/expressions/dynamic_filters.rs

@@ -80,6 +80,9 @@ pub struct DynamicFilterPhysicalExpr {
     /// parquet reader can use these to initialize the filter threshold from
     /// column statistics before reading any data.
     sort_options: Option<Vec<SortOptions>>,
+    /// Optional TopK fetch limit (K in LIMIT K).
+    /// Used by the parquet reader for cumulative RG pruning after reorder.
+    fetch: Option<usize>,
 }
 
 #[derive(Debug)]
@@ -184,6 +187,7 @@ impl DynamicFilterPhysicalExpr {
             data_type: Arc::new(RwLock::new(None)),
             nullable: Arc::new(RwLock::new(None)),
             sort_options: None,
+            fetch: None,
         }
     }
 
@@ -196,9 +200,11 @@ impl DynamicFilterPhysicalExpr {
         children: Vec<Arc<dyn PhysicalExpr>>,
         inner: Arc<dyn PhysicalExpr>,
         sort_options: Vec<SortOptions>,
+        fetch: Option<usize>,
     ) -> Self {
         let mut this = Self::new(children, inner);
         this.sort_options = Some(sort_options);
+        this.fetch = fetch;
         this
     }
 
@@ -207,6 +213,11 @@ impl DynamicFilterPhysicalExpr {
         self.sort_options.as_deref()
     }
 
+    /// Returns the TopK fetch limit (K), if available.
+    pub fn fetch(&self) -> Option<usize> {
+        self.fetch
+    }
+
     fn remap_children(
         children: &[Arc<dyn PhysicalExpr>],
         remapped_children: Option<&Vec<Arc<dyn PhysicalExpr>>>,
@@ -396,6 +407,7 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr {
             data_type: Arc::clone(&self.data_type),
             nullable: Arc::clone(&self.nullable),
             sort_options: self.sort_options.clone(),
+            fetch: self.fetch,
         }))
     }
 

datafusion/physical-plan/src/sorts/sort.rs

@@ -925,6 +925,7 @@ impl SortExec {
                 children,
                 lit(true),
                 sort_options,
+                self.fetch,
             ),
         ))))
     }