From 5fa9feb8ee6a217c5a4bb2af705f7f80a0169da4 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 19:06:01 +0300 Subject: [PATCH] Support equality with one in DF55 WindowTopN safely --- .../physical-optimizer/src/window_topn.rs | 69 ++++++-- .../sqllogictest/test_files/window_topn.slt | 152 ++++++++++++++++-- 2 files changed, 197 insertions(+), 24 deletions(-) diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index 20bd8b0d38a1c..4f0287913e3c3 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -99,6 +99,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 /// @@ -108,11 +109,10 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; /// - The window function is `ROW_NUMBER` or `RANK` (not `DENSE_RANK`) /// - The window function has a `PARTITION BY` clause (global top-K is /// already handled by `SortExec` with `fetch`) -/// - For `RANK`: a non-empty `ORDER BY` clause (otherwise all rows tie -/// at rank 1 — the optimization is useless and the boundary-tie storage -/// would be unbounded) +/// - At least one effective `ORDER BY` key remains after removing keys +/// already covered by the partition prefix /// - 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)] @@ -158,6 +158,17 @@ impl WindowTopN { } let fn_kind = supported_window_fn(&window_exprs[window_expr_idx])?; + // Pruning rows before the window can change sibling expressions such as + // LEAD. Ranking functions with the same keys only depend on retained rows. + let matched_expr = &window_exprs[window_expr_idx]; + if !window_exprs.iter().all(|expr| { + supported_window_fn(expr).is_some() + && expr.partition_by() == matched_expr.partition_by() + && expr.order_by() == matched_expr.order_by() + }) { + return None; + } + // Step 5: Validate PARTITION BY / ORDER BY and collect sort keys from the window expr let partition_by = window_exprs[window_expr_idx].partition_by(); let partition_prefix_len = partition_by.len(); @@ -168,14 +179,7 @@ impl WindowTopN { return None; } - // For RANK: an empty ORDER BY makes every row tie at rank 1 — - // the optimization is degenerate (we'd retain the entire input) - // and tie storage would be unbounded. let order_by = window_exprs[window_expr_idx].order_by(); - if matches!(fn_kind, WindowFnKind::Rank) && order_by.is_empty() { - return None; - } - // Step 6: Build PartitionedTopKExec from the window's partition/order keys let expr_iterator = partition_by .iter() @@ -183,6 +187,12 @@ impl WindowTopN { .chain(order_by.iter().cloned()); let expr = LexOrdering::new(expr_iterator)?; + // TopK requires an order key beyond the partition prefix. LexOrdering + // can remove redundant ORDER BY keys such as ORDER BY pk. + if expr.len() <= partition_prefix_len { + return None; + } + let partitioned_topk = PartitionedTopKExec::try_new( Arc::clone(window_exec_typed.input()), expr, @@ -250,13 +260,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, @@ -275,6 +288,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, }; } @@ -288,6 +302,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, }; } @@ -359,3 +374,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 44cb31153b004..1d0f2f1b09c6b 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -865,30 +865,52 @@ 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 *, 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 @@ -901,9 +923,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 @@ -1085,6 +1107,112 @@ set datafusion.execution.batch_size = 1; statement ok set datafusion.optimizer.enable_window_topn = true; +# A sibling LEAD needs the row after rn=1; pruning before the window would +# turn next_val into NULL. Keep the regular window plan for this case. +statement ok +CREATE TABLE window_topn_sibling_t (id INT, pk INT, val INT) AS VALUES + (1, 1, 10), (2, 1, 20), + (5, 2, 5), (6, 2, 15), + (9, 3, 50), (10, 3, 75); + +query IIIII rowsort +SELECT id, pk, val, rn, next_val FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) AS rn, + LEAD(val) OVER (PARTITION BY pk ORDER BY val) AS next_val + FROM window_topn_sibling_t +) WHERE rn = 1; +---- +1 1 10 1 20 +5 2 5 1 15 +9 3 50 1 75 + +query TT +EXPLAIN SELECT * FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) AS rn, + LEAD(val) OVER (PARTITION BY pk ORDER BY val) AS next_val + FROM window_topn_sibling_t +) WHERE rn = 1; +---- +logical_plan +01)Projection: window_topn_sibling_t.id, window_topn_sibling_t.pk, window_topn_sibling_t.val, row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn, lead(window_topn_sibling_t.val) PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS next_val +02)--Filter: row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW = UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, lead(window_topn_sibling_t.val) PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: window_topn_sibling_t projection=[id, pk, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn, lead(window_topn_sibling_t.val) PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as next_val] +02)--FilterExec: row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 = 1 +03)----RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=1, maintains_sort_order=true +04)------BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, lead(window_topn_sibling_t.val) PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "lead(window_topn_sibling_t.val) PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int32 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +05)--------SortPreservingMergeExec: [pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST] +06)----------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[true] +07)------------DataSourceExec: partitions=5, partition_sizes=[2, 1, 1, 1, 1] + +# Both plans have no effective order key and must avoid PartitionedTopKExec. +query I +SELECT COUNT(*) FROM ( + SELECT * FROM ( + SELECT *, ROW_NUMBER() OVER (PARTITION BY pk) AS rn + FROM window_topn_sibling_t + ) WHERE rn = 1 +); +---- +3 + +query I +SELECT COUNT(*) FROM ( + SELECT * FROM ( + SELECT *, ROW_NUMBER() OVER (PARTITION BY pk ORDER BY pk) AS rn + FROM window_topn_sibling_t + ) WHERE rn = 1 +); +---- +3 + +query TT +EXPLAIN SELECT * FROM ( + SELECT *, ROW_NUMBER() OVER (PARTITION BY pk) AS rn + FROM window_topn_sibling_t +) WHERE rn = 1; +---- +logical_plan +01)Projection: window_topn_sibling_t.id, window_topn_sibling_t.pk, window_topn_sibling_t.val, row_number() PARTITION BY [window_topn_sibling_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS rn +02)--Filter: row_number() PARTITION BY [window_topn_sibling_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING = UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [window_topn_sibling_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] +04)------TableScan: window_topn_sibling_t projection=[id, pk, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_sibling_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@3 as rn] +02)--FilterExec: row_number() PARTITION BY [window_topn_sibling_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@3 = 1 +03)----RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=1, maintains_sort_order=true +04)------BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_sibling_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Field { "row_number() PARTITION BY [window_topn_sibling_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING], mode=[Sorted] +05)--------SortPreservingMergeExec: [pk@1 ASC NULLS LAST] +06)----------SortExec: expr=[pk@1 ASC NULLS LAST], preserve_partitioning=[true] +07)------------DataSourceExec: partitions=5, partition_sizes=[2, 1, 1, 1, 1] + +query TT +EXPLAIN SELECT * FROM ( + SELECT *, ROW_NUMBER() OVER (PARTITION BY pk ORDER BY pk) AS rn + FROM window_topn_sibling_t +) WHERE rn = 1; +---- +logical_plan +01)Projection: window_topn_sibling_t.id, window_topn_sibling_t.pk, window_topn_sibling_t.val, row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.pk ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn +02)--Filter: row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.pk ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW = UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.pk ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: window_topn_sibling_t projection=[id, pk, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.pk ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] +02)--FilterExec: row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.pk ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 = 1 +03)----RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=1, maintains_sort_order=true +04)------BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.pk ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_sibling_t.pk] ORDER BY [window_topn_sibling_t.pk ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +05)--------SortPreservingMergeExec: [pk@1 ASC NULLS LAST] +06)----------SortExec: expr=[pk@1 ASC NULLS LAST], preserve_partitioning=[true] +07)------------DataSourceExec: partitions=5, partition_sizes=[2, 1, 1, 1, 1] + +statement ok +DROP TABLE window_topn_sibling_t; + query TT EXPLAIN SELECT * FROM ( SELECT c1, c2, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY c2 DESC) as rn