address review comments

Author: SubhamSinghal

datafusion/physical-optimizer/src/window_topn.rs

@@ -108,11 +108,18 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
 /// - Config flag `enable_window_topn` is `true`
 /// - The plan matches `FilterExec → [ProjectionExec] → BoundedWindowAggExec → SortExec`
 /// - 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
+///   sibling that reads pruned rows (e.g. `LEAD`, or an aggregate whose
+///   frame is not strictly backward-looking) would be computed over the
+///   pruned input and give wrong results.
 /// - The window function has a `PARTITION BY` clause (global top-K is
 ///   already handled by `SortExec` with `fetch`)
-/// - For `RANK` / `DENSE_RANK`: a non-empty `ORDER BY` clause (otherwise
-///   all rows tie at rank 1 — the optimization is useless and the
-///   retained set would be unbounded)
+/// - At least one `ORDER BY` key survives past the `PARTITION BY` prefix
+///   (so the operator has a non-empty ORDER BY). This rejects both a
+///   missing `ORDER BY` and one fully covered by the partition prefix
+///   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 `>`
 ///
@@ -169,6 +176,28 @@ impl WindowTopN {
         }
         let fn_kind = supported_window_fn(&window_exprs[window_expr_idx])?;
 
+        // Tail-pruning drops the rows that rank after the retained top-K,
+        // and every window expression in this `BoundedWindowAggExec` is
+        // then evaluated over the *pruned* input. The rewrite is only valid
+        // if each expression's value for a *retained* row is unaffected by
+        // the dropped rows. ROW_NUMBER / RANK / DENSE_RANK over the same
+        // PARTITION BY / ORDER BY satisfy this — each depends only on rows
+        // at or before the current row in that order, all of which are
+        // retained. A sibling like `LEAD(x)` reads following (pruned) rows,
+        // so at the retained boundary it would resolve to a dropped row and
+        // give a wrong result. Bail out unless every window expression is a
+        // supported ranking function sharing the matched expression's
+        // partition/order keys.
+        let matched_expr = &window_exprs[window_expr_idx];
+        let all_prune_safe = window_exprs.iter().all(|e| {
+            supported_window_fn(e).is_some()
+                && e.partition_by() == matched_expr.partition_by()
+                && e.order_by() == matched_expr.order_by()
+        });
+        if !all_prune_safe {
+            return None;
+        }
+
         // Step 5: child of window is SortExec (verified above)
         let sort_child = sort_exec.input();
 
@@ -182,12 +211,21 @@ impl WindowTopN {
             return None;
         }
 
-        // For RANK / DENSE_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.
-        if matches!(fn_kind, WindowFnKind::Rank | WindowFnKind::DenseRank)
-            && window_exprs[window_expr_idx].order_by().is_empty()
-        {
+        // `PartitionedTopKExec` derives its ORDER BY keys from the
+        // SortExec ordering *beyond* the partition prefix
+        // (`expr[partition_prefix_len..]`). That slice is empty in two
+        // cases:
+        //   * no ORDER BY at all (e.g. `ROW_NUMBER() OVER (PARTITION BY pk)`);
+        //   * ORDER BY keys fully covered by the partition prefix (e.g.
+        //     `DENSE_RANK() OVER (PARTITION BY pk ORDER BY pk)`, whose
+        //     deduplicated sort ordering is just `[pk]`).
+        // With zero order keys the operator panics on execution (it
+        // requires at least one), and for RANK / DENSE_RANK every row
+        // would tie at rank 1 (a degenerate, unbounded retained set).
+        // `order_by()` alone does not catch the second case — it reports
+        // `[pk]` even though no order key survives past the partition
+        // prefix — so guard on the effective order-key count instead.
+        if sort_exec.expr().len() <= partition_prefix_len {
             return None;
         }
 

datafusion/physical-plan/src/topk/mod.rs

@@ -1819,6 +1819,14 @@ impl PartitionedTopKRank {
 /// `RecordBatch`. `TieEntry` is reused verbatim from [`RankPartitionState`].
 struct DenseRankPartitionState {
     groups: HashMap<Vec<u8>, Vec<TieEntry>>,
+    /// Cached admission boundary: the largest tracked ob value once
+    /// `groups` is full. A `HashMap` is unordered, so finding it otherwise
+    /// costs an O(K) scan of the keys on *every* new distinct ob value —
+    /// O(N·K) overall on mostly-distinct input. Caching makes the common
+    /// "new ob is worse than the boundary → drop" path O(1); the O(K) scan
+    /// is paid only when the cache is stale (`None`) — after a fill or an
+    /// eviction changes the key set.
+    max_key: Option<Vec<u8>>,
 }
 
 impl DenseRankPartitionState {
@@ -1839,7 +1847,7 @@ impl DenseRankPartitionState {
                         .sum::<usize>()
             })
             .sum();
-        table_overhead + contents
+        table_overhead + contents + self.max_key.as_ref().map_or(0, |k| k.capacity())
     }
 }
 
@@ -1869,8 +1877,9 @@ impl DenseRankPartitionState {
 ///   new `TieEntry` (one entry per contributing batch).
 /// - `ob_key` new, `state.groups.len() < k` → insert the run as a new
 ///   group.
-/// - `ob_key` new, `state.groups.len() == k` → find the current max via
-///   an O(K) scan of `state.groups.keys()`:
+/// - `ob_key` new, `state.groups.len() == k` → the largest tracked ob
+///   value is the admission boundary, cached in `max_key` (recomputed via
+///   an O(K) scan of `state.groups.keys()` only when the cache is stale):
 ///   - `ob_key < max` → remove the max key (evict the entire max-key
 ///     group — up to many rows) and insert the run. The evicted group's
 ///     row count is added to the `row_replacements` metric.
@@ -2024,6 +2033,7 @@ impl PartitionedTopKDenseRank {
                     .entry(pk)
                     .or_insert_with(|| DenseRankPartitionState {
                         groups: HashMap::new(),
+                        max_key: None,
                     });
 
             // Bucket by ob key. `ob_runs` is a reused scratch map (taken
@@ -2061,22 +2071,30 @@ impl PartitionedTopKDenseRank {
                             batch_bytes: input_batch_bytes,
                         }],
                     );
+                    // A new distinct key may raise the boundary; drop the
+                    // cached max so it is recomputed on the next Case C.
+                    state.max_key = None;
                     continue;
                 }
 
-                // Case C: new ob, at K distinct keys. Find the current
-                // max (K-th distinct-best) via an O(K) scan — cold path.
-                // Scoped so the immutable borrow ends before mutation.
-                let evict_key: Option<Vec<u8>> = {
-                    let max_key = state
-                        .groups
-                        .keys()
-                        .map(|key| key.as_slice())
-                        .max()
-                        .expect("state.groups has k >= 1 keys");
-                    (ob_key.as_slice() < max_key).then(|| max_key.to_vec())
-                };
-                if let Some(evicted_key) = evict_key {
+                // Case C: new ob, at K distinct keys. The largest tracked
+                // ob value is the admission boundary; it is cached in
+                // `max_key` and recomputed via an O(K) scan only when the
+                // cache is stale, so the common "ob >= max → drop" path is
+                // O(1). Scoped so the immutable borrow ends before mutation.
+                if state.max_key.is_none() {
+                    state.max_key = state.groups.keys().max().cloned();
+                }
+                let is_smaller = state
+                    .max_key
+                    .as_deref()
+                    .map(|max_key| ob_key.as_slice() < max_key)
+                    .expect("state.groups has k >= 1 keys");
+
+                if is_smaller {
+                    // Evict the entire max-key group and drop the cached
+                    // max (recomputed on the next Case C).
+                    let evicted_key = state.max_key.take().expect("max key present");
                     let evicted =
                         state.groups.remove(&evicted_key).expect("max key present");
                     replacements +=
@@ -2139,7 +2157,7 @@ impl PartitionedTopKDenseRank {
         let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size);
 
         for pk in sorted_pks {
-            let DenseRankPartitionState { groups } =
+            let DenseRankPartitionState { groups, max_key: _ } =
                 states.remove(&pk).expect("key from states.keys()");
             // HashMap is unordered — sort the <= K distinct ob keys so
             // rows emit ascending (byte-comparable encoding == sort order).

datafusion/sqllogictest/test_files/window_topn.slt

@@ -329,9 +329,115 @@ physical_plan
 04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[false]
 05)--------DataSourceExec: partitions=1, partition_sizes=[1]
 
+# Sibling safety: an aggregate sibling (SUM) is computed over the pruned
+# input, so the rule must NOT fire even though the filter is on ROW_NUMBER.
+# SUM over the default RANGE frame includes tie-peers that pruning would
+# drop, giving a wrong sum. The plan keeps SortExec + FilterExec — no
+# PartitionedTopKExec.
+query TT
+EXPLAIN SELECT * FROM (
+  SELECT *,
+    ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+    SUM(val) OVER (PARTITION BY pk ORDER BY val) as running_sum
+  FROM window_topn_t
+) WHERE rn <= 3;
+----
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn, sum(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as running_sum]
+02)--FilterExec: row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 <= 3
+03)----BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, sum(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, 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]
+
+# Sibling safety: a LEAD sibling reads following (pruned) rows, so the
+# rule must NOT fire — the boundary row's LEAD would resolve to a pruned
+# row. The plan keeps SortExec + FilterExec.
+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_t
+) WHERE rn <= 3;
+----
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn, lead(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as next_val]
+02)--FilterExec: row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 <= 3
+03)----BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, lead(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "lead(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int32 }, 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]
+
 statement ok
 SET datafusion.explain.physical_plan_only = false;
 
+# Sibling safety (correctness): LEAD reads following rows, so if the rule
+# fired, the pruned input would give a wrong LEAD at the retained boundary
+# — each rn=2 row's next_val would become NULL instead of the (pruned)
+# rn=3 row's value. The guard keeps the normal plan, so LEAD sees the full
+# partition and next_val is correct (30 / 25 / 100 for the rn=2 rows).
+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_t
+) WHERE rn <= 2;
+----
+1 1 10 1 20
+10 3 75 2 100
+2 1 20 2 30
+5 2 5 1 15
+6 2 15 2 25
+9 3 50 1 75
+
+# A tie table for sibling-safety correctness: pk=1 has three peers at
+# val=5, one 10, one 20 (row_number 1..5; rank 1,1,1,4,5; dense_rank
+# 1,1,1,2,3).
+statement ok
+CREATE TABLE window_topn_sib_t (id INT, pk INT, val INT) AS VALUES
+  (1, 1, 5),
+  (2, 1, 5),
+  (3, 1, 5),
+  (4, 1, 10),
+  (5, 1, 20);
+
+# SUM sibling (correctness): fetch=2 but there are three peers at val=5,
+# so one peer would be pruned. The guard keeps the normal plan, so the
+# RANGE frame sees all three 5s and running_sum = 15. If the rule fired,
+# the pruned 2-row input would sum to 10. (id excluded — which two of the
+# three tied rows get row_number 1/2 is not deterministic.)
+query III rowsort
+SELECT pk, val, running_sum FROM (
+  SELECT *,
+    ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+    SUM(val)     OVER (PARTITION BY pk ORDER BY val) as running_sum
+  FROM window_topn_sib_t
+) WHERE rn <= 2;
+----
+1 5 15
+1 5 15
+
+# All three ranking siblings (ROW_NUMBER + RANK + DENSE_RANK) over the same
+# window are prune-safe, so the rule DOES fire (filter on rn). Verify the
+# rank/dense_rank values stay correct under the optimized plan — including
+# the val=10 row where row_number=4, rank=4, dense_rank=2 all differ.
+query IIII rowsort
+SELECT pk, val, rnk, dr FROM (
+  SELECT *,
+    ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+    RANK()       OVER (PARTITION BY pk ORDER BY val) as rnk,
+    DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr
+  FROM window_topn_sib_t
+) WHERE rn <= 4;
+----
+1 10 4 2
+1 5 1 1
+1 5 1 1
+1 5 1 1
+
+statement ok
+DROP TABLE window_topn_sib_t;
+
 # Test 15: ROW_NUMBER with DESC ordering — correctness
 query III rowsort
 SELECT id, pk, val FROM (
@@ -1292,6 +1398,34 @@ SELECT id, pk, val FROM (
 ) WHERE dr < 1;
 ----
 
+# Test DR11b: ORDER BY keys fully covered by the PARTITION BY prefix
+# (`PARTITION BY pk ORDER BY pk`). The deduplicated sort ordering is just
+# `[pk]`, so no order key survives past the partition prefix. The rule
+# must NOT fire — a PartitionedTopKExec built with zero order expressions
+# panics on execution. Every row ties at dense_rank 1, so `dr <= 2` keeps
+# the whole table (returned via the ordinary window + filter plan).
+query III rowsort
+SELECT id, pk, val FROM (
+  SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY pk) as dr FROM window_topn_dense_rank_t
+) WHERE dr <= 2;
+----
+1 1 10
+10 2 25
+11 3 100
+12 3 200
+13 3 200
+14 3 200
+15 3 200
+16 3 300
+2 1 20
+3 1 20
+4 1 20
+5 1 30
+6 1 40
+7 2 5
+8 2 5
+9 2 15
+
 statement ok
 DROP TABLE window_topn_dense_rank_t;