Support equality with one in opt-in WindowTopN rule

datafusion/physical-optimizer/src/window_topn.rs

@@ -103,6 +103,7 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
 /// - `rn < K` → fetch = K - 1
 /// - `K >= rn` (flipped) → fetch = K
 /// - `K > rn` (flipped) → fetch = K - 1
+/// - `rn = 1` or `1 = rn` → fetch = 1
 ///
 /// # When the Rule Fires
 ///
@@ -126,7 +127,7 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
 ///   such as `PARTITION BY pk ORDER BY pk`; for `RANK` / `DENSE_RANK`
 ///   such orderings also make every row tie at rank 1 (degenerate).
 /// - The filter predicate compares the window output column to an integer
-///   literal using `<=`, `<`, `>=`, or `>`
+///   literal using `<=`, `<`, `>=`, `>`, or equality with 1
 ///
 /// [`PartitionedTopKExec`]: datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec
 #[derive(Default, Clone, Debug)]
@@ -351,13 +352,16 @@ impl PhysicalOptimizerRule for WindowTopN {
 /// | `Column(idx) < Literal(N)` | `(idx, N-1)` |
 /// | `Literal(N) >= Column(idx)` | `(idx, N)` |
 /// | `Literal(N) > Column(idx)` | `(idx, N-1)` |
+/// | `Column(idx) = Literal(1)` | `(idx, 1)` |
+/// | `Literal(1) = Column(idx)` | `(idx, 1)` |
 ///
 /// # Examples
 ///
 /// - `rn <= 5` → `Some((2, 5))` (assuming rn is column index 2)
 /// - `rn < 3` → `Some((2, 2))`
 /// - `10 >= rn` → `Some((2, 10))`
-/// - `rn = 1` → `None` (equality not supported)
+/// - `rn = 1` → `Some((2, 1))`
+/// - `rn = 2` → `None` (removing the filter would include rank 1)
 /// - `val <= 5` → `Some((1, 5))` (caller must verify it's a window column)
 fn extract_window_limit(
     predicate: &Arc<dyn datafusion_physical_expr::PhysicalExpr>,
@@ -376,6 +380,7 @@ fn extract_window_limit(
         return match *op {
             Operator::LtEq => Some((col.index(), n)),
             Operator::Lt => Some((col.index(), n - 1)),
+            Operator::Eq if n == 1 => Some((col.index(), 1)),
             _ => None,
         };
     }
@@ -389,6 +394,7 @@ fn extract_window_limit(
         return match *op {
             Operator::GtEq => Some((col.index(), n)),
             Operator::Gt => Some((col.index(), n - 1)),
+            Operator::Eq if n == 1 => Some((col.index(), 1)),
             _ => None,
         };
     }
@@ -498,3 +504,33 @@ fn column_index_below(
         }
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion_physical_expr::PhysicalExpr;
+
+    #[test]
+    fn equality_with_one_is_the_only_rewritable_equality() {
+        for (value, flipped, expected) in [
+            (Some(1), false, Some((2, 1))),
+            (Some(1), true, Some((2, 1))),
+            (Some(0), false, None),
+            (Some(2), false, None),
+            (Some(2), true, None),
+            (None, false, None),
+        ] {
+            let column: Arc<dyn PhysicalExpr> = Arc::new(Column::new("rn", 2));
+            let literal: Arc<dyn PhysicalExpr> =
+                Arc::new(Literal::new(ScalarValue::UInt64(value)));
+            let (left, right) = if flipped {
+                (literal, column)
+            } else {
+                (column, literal)
+            };
+            let predicate: Arc<dyn PhysicalExpr> =
+                Arc::new(BinaryExpr::new(left, Operator::Eq, right));
+            assert_eq!(extract_window_limit(&predicate), expected);
+        }
+    }
+}

datafusion/sqllogictest/test_files/window_topn.slt

@@ -941,30 +941,61 @@ statement ok
 DROP TABLE window_topn_rank_t;
 
 ###############################################################################
-# RANK() — equality predicate (negative: rule supports only </<=/>=/>)
+# Ranking functions — equality with one
 ###############################################################################
 #
-# `extract_window_limit` matches only `<, <=, >, >=`. Equality predicates
-# `rk = N` are NOT optimized by this rule (regardless of N). DuckDB
-# special-cases `rk = 1` as equivalent to `rk <= 1`; we don't. The two
-# tests below pin current behavior so that an accidental rule extension
-# (or regression) shows up.
+# Ranking functions return positive integers, so `rank = 1` is equivalent to
+# `rank <= 1`. Equality with larger values must retain the filter.
 
 statement ok
 CREATE TABLE window_topn_rank_eq_t (id INT, pk INT, val INT) AS VALUES
-  (1, 1, 10), (2, 1, 20), (3, 1, 30),
+  (1, 1, 10), (2, 1, 10), (3, 1, 20),
   (4, 2, 5),  (5, 2, 15), (6, 2, 25);
 
-# Test R12: `rk = 1` — correct results, but plan should still contain
-# FilterExec + SortExec (rule did NOT fire).
+# Test R12: `rk = 1` keeps all tied winners.
 query III rowsort
 SELECT id, pk, val FROM (
   SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_eq_t
 ) WHERE rk = 1;
 ----
 1 1 10
+2 1 10
 4 2 5
 
+query III rowsort
+SELECT id, pk, val FROM (
+  SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_eq_t
+) WHERE 1 = rk;
+----
+1 1 10
+2 1 10
+4 2 5
+
+query III rowsort
+SELECT id, pk, val FROM (
+  SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_eq_t
+) WHERE rk = 1;
+----
+1 1 10
+2 1 10
+4 2 5
+
+query III rowsort
+SELECT id, pk, val FROM (
+  SELECT *, ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val, id) as rn FROM window_topn_rank_eq_t
+) WHERE rn = 1;
+----
+1 1 10
+4 2 5
+
+# `rank = 2` is not equivalent to `rank <= 2`.
+query III rowsort
+SELECT id, pk, val FROM (
+  SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_eq_t
+) WHERE rk = 2;
+----
+5 2 15
+
 query TT
 EXPLAIN SELECT * FROM (
   SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_eq_t
@@ -977,10 +1008,34 @@ logical_plan
 04)------TableScan: window_topn_rank_eq_t projection=[id, pk, val]
 physical_plan
 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
-02)--FilterExec: rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 = 1
-03)----BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[false]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=rank, fetch=1, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+SET datafusion.explain.physical_plan_only = true;
+
+query TT
+EXPLAIN SELECT * FROM (
+  SELECT *, ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val, id) as rn FROM window_topn_rank_eq_t
+) WHERE rn = 1;
+----
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST, window_topn_rank_eq_t.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
+02)--PartitionedTopKExec: fn=row_number, fetch=1, partition=[pk@1], order=[val@2 ASC NULLS LAST, id@0 ASC NULLS LAST], emit=[row_number() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST, window_topn_rank_eq_t.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+query TT
+EXPLAIN SELECT * FROM (
+  SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_eq_t
+) WHERE rk = 1;
+----
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, dense_rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
+02)--PartitionedTopKExec: fn=dense_rank, fetch=1, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[dense_rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+SET datafusion.explain.physical_plan_only = false;
 
 statement ok
 DROP TABLE window_topn_rank_eq_t;
@@ -1689,4 +1744,4 @@ statement ok
 SET datafusion.optimizer.enable_window_topn = false;
 
 statement ok
-DROP TABLE window_topn_shift;
+DROP TABLE window_topn_shift;