diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index 758ad66de301..f970d058651c 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -101,6 +101,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 /// @@ -121,7 +122,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)] @@ -300,13 +301,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, @@ -325,6 +329,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, }; } @@ -338,6 +343,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, }; } @@ -411,3 +417,33 @@ fn find_window_below(plan: &Arc) -> Option = Arc::new(Column::new("rn", 2)); + let literal: Arc = + Arc::new(Literal::new(ScalarValue::UInt64(value))); + let (left, right) = if flipped { + (literal, column) + } else { + (column, literal) + }; + let predicate: Arc = + Arc::new(BinaryExpr::new(left, Operator::Eq, right)); + assert_eq!(extract_window_limit(&predicate), expected); + } + } +} diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt index 5c3757aa5f77..337bd291f6c1 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -994,30 +994,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 @@ -1030,9 +1061,9 @@ 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] +02)--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] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=rank, fetch=1, partition=[pk@1], order=[val@2 ASC NULLS LAST] 05)--------DataSourceExec: partitions=1, partition_sizes=[1] statement ok