fix: WindowTopN resolves the filter column through intermediate proje…

datafusion/core/tests/physical_optimizer/window_topn.rs

@@ -46,10 +46,13 @@ use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::metrics::MetricValue;
 use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
 use datafusion_physical_plan::projection::ProjectionExec;
-use datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec;
+use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::sorts::partitioned_topk::{
+    PartitionedTopKExec, WindowFnKind,
+};
 use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::windows::{BoundedWindowAggExec, create_udwf_window_expr};
-use datafusion_physical_plan::{ExecutionPlan, InputOrderMode};
+use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, Partitioning};
 use insta::assert_snapshot;
 
 fn schema() -> Arc<Schema> {
@@ -458,6 +461,149 @@ fn two_window_exprs_keep_the_window_node() -> Result<()> {
     Ok(())
 }
 
+/// `ROW_NUMBER() AS rn, RANK() AS rk` over `PARTITION BY pk ORDER BY val`.
+/// Output: `[pk@0, val@1, rn@2, rk@3]`.
+fn build_rn_rk_window() -> Result<Arc<dyn ExecutionPlan>> {
+    let s = schema();
+    let input: Arc<dyn ExecutionPlan> = Arc::new(PlaceholderRowExec::new(Arc::clone(&s)));
+    let partition_by = vec![col("pk", &s)?];
+    let order_by = vec![PhysicalSortExpr::new_default(col("val", &s)?).asc()];
+    let window_expr = |udwf: Arc<datafusion_expr::WindowUDF>, name: &str| {
+        Ok::<_, datafusion_common::DataFusionError>(Arc::new(StandardWindowExpr::new(
+            create_udwf_window_expr(&udwf, &[], &s, name.to_string(), false)?,
+            &partition_by,
+            &order_by,
+            Arc::new(WindowFrame::new_bounds(
+                WindowFrameUnits::Rows,
+                WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
+                WindowFrameBound::CurrentRow,
+            )),
+        )))
+    };
+    Ok(Arc::new(BoundedWindowAggExec::try_new(
+        vec![
+            window_expr(row_number_udwf(), "rn")?,
+            window_expr(rank_udwf(), "rk")?,
+        ],
+        input,
+        InputOrderMode::Sorted,
+        true,
+    )?))
+}
+
+/// A `ProjectionExec` passing through `input`'s columns at `indices`.
+fn project(
+    input: Arc<dyn ExecutionPlan>,
+    indices: &[usize],
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let schema = input.schema();
+    let exprs: Vec<(Arc<dyn datafusion_physical_expr::PhysicalExpr>, String)> = indices
+        .iter()
+        .map(|&i| {
+            let name = schema.field(i).name();
+            (
+                Arc::new(Column::new(name, i))
+                    as Arc<dyn datafusion_physical_expr::PhysicalExpr>,
+                name.to_string(),
+            )
+        })
+        .collect();
+    Ok(Arc::new(ProjectionExec::try_new(exprs, input)?))
+}
+
+/// `FilterExec(<column idx of input> <= 3)`, with the literal in the column's type.
+fn filter_lteq_3(
+    input: Arc<dyn ExecutionPlan>,
+    idx: usize,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let schema = input.schema();
+    let field = schema.field(idx);
+    let predicate = Arc::new(BinaryExpr::new(
+        Arc::new(Column::new(field.name(), idx)),
+        Operator::LtEq,
+        lit(ScalarValue::UInt64(Some(3)).cast_to(field.data_type())?),
+    ));
+    Ok(Arc::new(FilterExec::try_new(predicate, input)?))
+}
+
+/// The ranking function the rewrite chose, or `None` if it did not fire.
+fn rewritten_fn(plan: Arc<dyn ExecutionPlan>) -> Result<Option<WindowFnKind>> {
+    let optimized = optimize(plan)?;
+    Ok(find_partitioned_topk(&optimized).map(|p| {
+        p.downcast_ref::<PartitionedTopKExec>()
+            .expect("find_partitioned_topk returns a PartitionedTopKExec")
+            .fn_kind()
+    }))
+}
+
+/// Stacked projections are resolved top-down: each one maps an index into its
+/// output to one into its input. Bottom-up, `rk@0` would follow `pk@0` and then
+/// `rk@2` to `rn` in the window's output.
+#[test]
+fn stacked_projections_resolve_top_down() -> Result<()> {
+    let lower = project(build_rn_rk_window()?, &[0, 2, 3])?; // [pk, rn, rk]
+    let upper = project(lower, &[2, 0])?; // [rk, pk]
+    assert_eq!(
+        rewritten_fn(filter_lteq_3(upper, 0)?)?,
+        Some(WindowFnKind::Rank)
+    );
+    Ok(())
+}
+
+/// A `RepartitionExec` keeps its input's schema, so the index passes through it.
+#[test]
+fn repartition_between_passes_the_index_through() -> Result<()> {
+    let repartition: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new(
+        build_rn_rk_window()?,
+        Partitioning::RoundRobinBatch(4),
+    )?);
+    let projection = project(repartition, &[0, 3])?; // [pk, rk]
+    assert_eq!(
+        rewritten_fn(filter_lteq_3(projection, 1)?)?,
+        Some(WindowFnKind::Rank)
+    );
+    Ok(())
+}
+
+/// A projection that reorders the window's columns: the filter's `pk@2` maps
+/// back to `pk@0`, an input column, so the rule leaves the plan alone. Read as
+/// an index into the window's output, it would be `rn`.
+#[test]
+fn reordered_projection_input_column_no_change() -> Result<()> {
+    let projection = project(build_rn_rk_window()?, &[2, 3, 0])?; // [rn, rk, pk]
+    let plan = filter_lteq_3(projection, 2)?;
+    let before = plan_str(plan.as_ref());
+    assert_eq!(plan_str(optimize(plan)?.as_ref()), before);
+    Ok(())
+}
+
+/// A column the projection computes has no single window column behind it,
+/// so the rule leaves the plan alone. `x` sits at index 2 so that treating it
+/// as a pass-through would land on `rn` rather than be rejected as an input
+/// column.
+#[test]
+fn computed_projection_column_no_change() -> Result<()> {
+    let window = build_rn_rk_window()?;
+    let rk_plus_one: Arc<dyn datafusion_physical_expr::PhysicalExpr> =
+        Arc::new(BinaryExpr::new(
+            Arc::new(Column::new("rk", 3)),
+            Operator::Plus,
+            lit(ScalarValue::UInt64(Some(1))),
+        ));
+    let projection: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new(
+        vec![
+            (Arc::new(Column::new("pk", 0)) as _, "pk".to_string()),
+            (Arc::new(Column::new("val", 1)) as _, "val".to_string()),
+            (rk_plus_one, "x".to_string()),
+        ],
+        window,
+    )?);
+    let plan = filter_lteq_3(projection, 2)?;
+    let before = plan_str(plan.as_ref());
+    assert_eq!(plan_str(optimize(plan)?.as_ref()), before);
+    Ok(())
+}
+
 // ----------------------------------------------------------------------
 // RANK rule tests
 // ----------------------------------------------------------------------

datafusion/physical-optimizer/src/window_topn.rs

@@ -122,6 +122,9 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
 /// All of the following must be true:
 /// - Config flag `enable_window_topn` is `true`
 /// - The plan matches `FilterExec → [ProjectionExec] → BoundedWindowAggExec`
+/// - Every `ProjectionExec` between the filter and the window passes the
+///   filtered column through as a bare column, possibly at another position
+///   (as when a column the query no longer needs is dropped)
 /// - The window function is `ROW_NUMBER`, `RANK`, or `DENSE_RANK`
 /// - Every window expression in the `BoundedWindowAggExec` is `ROW_NUMBER`,
 ///   `RANK`, or `DENSE_RANK` over the same `PARTITION BY` / `ORDER BY`. A
@@ -178,7 +181,10 @@ impl WindowTopN {
         let child = filter.input();
         let (window_exec, intermediates) = find_window_below(child)?;
 
-        // Step 4: Verify col_idx references a supported window function output column
+        // Step 4: Verify col_idx references a supported window function output column.
+        // `col_idx` indexes the filter's input; resolve it to the window's output
+        // first, since a projection in between can drop or reorder columns.
+        let col_idx = column_index_below(col_idx, &intermediates)?;
         let window_exec_typed = window_exec.downcast_ref::<BoundedWindowAggExec>()?;
         let input_field_count = window_exec_typed.input().schema().fields().len();
         if col_idx < input_field_count {
@@ -440,8 +446,10 @@ type PlanAndIntermediates = (Arc<dyn ExecutionPlan>, Vec<Arc<dyn ExecutionPlan>>
 ///
 /// Handles sequences of `ProjectionExec` and `RepartitionExec`.
 /// This is safe because `PartitionedTopKExec` can be pushed below them:
-/// projections only provide aliases, and pushing the limit below repartitions
-/// is safe because the limit is computed per-partition.
+/// projections only rename, drop, reorder or compute over the window's
+/// columns (the filter's column is mapped back through them by
+/// [`column_index_below`]), and pushing the limit below repartitions is safe
+/// because the limit is computed per-partition.
 ///
 /// Returns the window exec and a list of intermediate nodes to rebuild,
 /// or `None` if no `BoundedWindowAggExec` is found.
@@ -463,3 +471,28 @@ fn find_window_below(plan: &Arc<dyn ExecutionPlan>) -> Option<PlanAndIntermediat
         }
     }
 }
+
+/// Maps `col_idx`, an index into the output of `intermediates[0]` (the
+/// filter's input), to the index of the same column in the output of the
+/// node below the last intermediate (the window).
+///
+/// `intermediates` is top-down, as [`find_window_below`] returns it, so each
+/// step reads an index into the node's output and yields one into its input.
+/// A `RepartitionExec` keeps its input's schema. A `ProjectionExec` must pass
+/// the column through as a bare [`Column`]: one it computes (e.g. `rn + 1`)
+/// has no single input column to follow, so this returns `None`.
+fn column_index_below(
+    col_idx: usize,
+    intermediates: &[Arc<dyn ExecutionPlan>],
+) -> Option<usize> {
+    intermediates.iter().try_fold(col_idx, |idx, node| {
+        if let Some(projection) = node.downcast_ref::<ProjectionExec>() {
+            let expr = &projection.expr().get(idx)?.expr;
+            Some(expr.downcast_ref::<Column>()?.index())
+        } else if node.downcast_ref::<RepartitionExec>().is_some() {
+            Some(idx)
+        } else {
+            None
+        }
+    })
+}