feat(physical-optimizer): support FilterExec with embedded projection in WindowTopN

The WindowTopN rule bailed out unconditionally when the FilterExec at
the top of the pattern carried an embedded projection (from an earlier
filter/projection pushdown pass), even when the underlying pattern
otherwise matched. That skipped rewrite path caused ROW_NUMBER
top-K-per-group queries to fall back to a full sort in production.

Capture the FilterExec's projection indices at the start, run the
existing PartitionedTopKExec rewrite as usual, and re-apply the
captured projection as an outer ProjectionExec so the transformed plan
preserves the original output schema.

Adds a regression test in datafusion/core/tests/physical_optimizer/window_topn.rs
covering the FilterExec-with-projection shape.

Author: zhuqi-lucas

datafusion/core/tests/physical_optimizer/window_topn.rs

@@ -33,7 +33,7 @@ use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::window_topn::WindowTopN;
 use datafusion_physical_plan::displayable;
-use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::filter::{FilterExec, FilterExecBuilder};
 use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::sorts::sort::SortExec;
@@ -607,3 +607,69 @@ fn dense_rank_no_change() -> Result<()> {
     );
     Ok(())
 }
+
+/// Regression: FilterExec that carries an embedded projection (e.g. from
+/// an earlier filter/projection pushdown pass) used to make the rule bail
+/// out entirely. The rule now captures the projection, applies the
+/// PartitionedTopKExec rewrite, and wraps the result in a ProjectionExec
+/// that reproduces the original output schema.
+#[test]
+fn filter_with_projection_still_rewrites() -> Result<()> {
+    let s = schema();
+    let input: Arc<dyn ExecutionPlan> = Arc::new(PlaceholderRowExec::new(Arc::clone(&s)));
+
+    let ordering = LexOrdering::new(vec![
+        PhysicalSortExpr::new_default(col("pk", &s)?).asc(),
+        PhysicalSortExpr::new_default(col("val", &s)?).asc(),
+    ])
+    .unwrap();
+    let sort: Arc<dyn ExecutionPlan> =
+        Arc::new(SortExec::new(ordering, input).with_preserve_partitioning(true));
+
+    let partition_by = vec![col("pk", &s)?];
+    let order_by = vec![PhysicalSortExpr::new_default(col("val", &s)?).asc()];
+    let window_expr = Arc::new(StandardWindowExpr::new(
+        create_udwf_window_expr(
+            &row_number_udwf(),
+            &[],
+            &s,
+            "row_number".to_string(),
+            false,
+        )?,
+        &partition_by,
+        &order_by,
+        Arc::new(WindowFrame::new_bounds(
+            WindowFrameUnits::Rows,
+            WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
+            WindowFrameBound::CurrentRow,
+        )),
+    ));
+    let window: Arc<dyn ExecutionPlan> = Arc::new(BoundedWindowAggExec::try_new(
+        vec![window_expr],
+        sort,
+        InputOrderMode::Sorted,
+        true,
+    )?);
+
+    // Filter: row_number@2 <= 3, with an embedded projection that keeps
+    // only [pk, val] (drops the row_number column) — the shape produced
+    // by filter/projection pushdown when downstream doesn't need the
+    // window column.
+    let rn_col = Arc::new(Column::new("row_number", 2));
+    let limit_lit = lit(ScalarValue::UInt64(Some(3)));
+    let predicate = Arc::new(BinaryExpr::new(rn_col, Operator::LtEq, limit_lit));
+    let filter: Arc<dyn ExecutionPlan> = Arc::new(
+        FilterExecBuilder::new(predicate, window)
+            .apply_projection(Some(vec![0, 1]))?
+            .build()?,
+    );
+
+    let optimized = optimize(filter)?;
+    assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+    ProjectionExec: expr=[pk@0 as pk, val@1 as val]
+      BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+        PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC]
+          PlaceholderRowExec
+    "#);
+    Ok(())
+}

datafusion/physical-optimizer/src/window_topn.rs

@@ -56,6 +56,7 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::Operator;
+use datafusion_physical_expr::PhysicalExpr;
 use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal};
 use datafusion_physical_expr::window::StandardWindowExpr;
 use datafusion_physical_plan::ExecutionPlan;
@@ -132,10 +133,14 @@ impl WindowTopN {
         // Step 1: Match FilterExec at the top
         let filter = plan.downcast_ref::<FilterExec>()?;
 
-        // Don't handle filters with projections
-        if filter.projection().is_some() {
-            return None;
-        }
+        // A projection embedded in the FilterExec (from an earlier
+        // filter/projection pushdown pass) is captured here and re-applied
+        // via a wrapping ProjectionExec at the end so the rewrite preserves
+        // the original output schema.
+        let filter_projection: Option<Vec<usize>> = filter
+            .projection()
+            .as_ref()
+            .map(|p| p.iter().copied().collect());
 
         // Step 2: Extract limit from predicate (rn <= K, rn < K, etc.)
         let (col_idx, limit_n) = extract_window_limit(filter.predicate())?;
@@ -195,13 +200,41 @@ impl WindowTopN {
             .ok()?;
 
         // Step 9: If ProjectionExec was between Filter and Window, rebuild it
-        let result = match proj_between {
+        let mut result = match proj_between {
             Some(proj) => Arc::clone(&child_as_arc(proj))
                 .with_new_children(vec![new_window])
                 .ok()?,
             None => new_window,
         };
 
+        // Step 10: Re-apply the FilterExec's embedded projection (if any)
+        // as an outer ProjectionExec. The projection indices refer to
+        // columns in `filter.input().schema()`, which equals `result`'s
+        // schema at this point (Steps 8-9 preserve schema), so the
+        // indices remain valid.
+        if let Some(indices) = filter_projection {
+            let input_schema = result.schema();
+            let field_count = input_schema.fields().len();
+            // Validate before indexing: an out-of-range index would panic
+            // in `input_schema.field(idx)`. Bail out of the rewrite instead
+            // so a malformed FilterExec projection can never crash the
+            // optimizer.
+            if indices.iter().any(|&idx| idx >= field_count) {
+                return None;
+            }
+            let projection_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = indices
+                .iter()
+                .map(|&idx| {
+                    let field = input_schema.field(idx);
+                    (
+                        Arc::new(Column::new(field.name(), idx)) as Arc<dyn PhysicalExpr>,
+                        field.name().clone(),
+                    )
+                })
+                .collect();
+            result = Arc::new(ProjectionExec::try_new(projection_exprs, result).ok()?);
+        }
+
         Some(result)
     }
 }
@@ -264,9 +297,7 @@ impl PhysicalOptimizerRule for WindowTopN {
 /// - `10 >= rn` → `Some((2, 10))`
 /// - `rn = 1` → `None` (equality not supported)
 /// - `val <= 5` → `Some((1, 5))` (caller must verify it's a window column)
-fn extract_window_limit(
-    predicate: &Arc<dyn datafusion_physical_expr::PhysicalExpr>,
-) -> Option<(usize, usize)> {
+fn extract_window_limit(predicate: &Arc<dyn PhysicalExpr>) -> Option<(usize, usize)> {
     let binary = predicate.downcast_ref::<BinaryExpr>()?;
     let op = binary.op();
     let left = binary.left();