perf: Extend WindowTopN to dense_rank

benchmarks/queries/h2o/window.sql

@@ -196,3 +196,53 @@ SELECT pk, largest_v2 FROM (
            RANK() OVER (PARTITION BY (id3 % 100000) ORDER BY v2 DESC) AS rk_v2
     FROM large WHERE v2 IS NOT NULL
 ) sub_query WHERE rk_v2 <= 2;
+
+-- Window Top-N (DENSE_RANK top-2 per partition, ~100 partitions)
+-- The DENSE_RANK queries below mirror the RANK cardinality sweep above.
+-- DENSE_RANK semantics keep every row whose ORDER BY value is among the
+-- K distinct-smallest values in the partition, so total kept per partition
+-- is unbounded in rows-per-distinct-value — exercises PartitionedTopKDenseRank's
+-- HashMap-of-groups path.
+SELECT pk, largest_v2 FROM (
+    SELECT (id3 % 100) AS pk, v2 AS largest_v2,
+           DENSE_RANK() OVER (PARTITION BY (id3 % 100) ORDER BY v2 DESC) AS dr_v2
+    FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE dr_v2 <= 2;
+
+-- Window Top-N (DENSE_RANK top-2 per partition, ~1K partitions)
+SELECT pkey, largest_v2 FROM (
+    SELECT (id3 % 1000) AS pkey, v2 AS largest_v2,
+           DENSE_RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY v2 DESC) AS dr_v2
+    FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE dr_v2 <= 2;
+
+-- Window Top-N (DENSE_RANK top-2 per partition, ~1K partitions, heavy ties)
+-- v2 % 10 forces 10 distinct OBY values; most rows share the top-2 distinct
+-- values so appends dominate — exercises the "Case A" append-to-existing-Vec
+-- fast path in PartitionedTopKDenseRank.
+SELECT pkey, largest_v2 FROM (
+    SELECT (id3 % 1000) AS pkey, v2 AS largest_v2,
+           DENSE_RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY (v2 % 10) DESC) AS dr_v2
+    FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE dr_v2 <= 2;
+
+-- Window Top-N (DENSE_RANK top-2 per partition, ~10K partitions, low ties)
+SELECT id2, largest_v2 FROM (
+    SELECT id2, v2 AS largest_v2,
+           DENSE_RANK() OVER (PARTITION BY id2 ORDER BY v2 DESC) AS dr_v2
+    FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE dr_v2 <= 2;
+
+-- Window Top-N (DENSE_RANK top-2 per partition, ~10K partitions, heavy ties)
+SELECT id2, largest_v2 FROM (
+    SELECT id2, v2 AS largest_v2,
+           DENSE_RANK() OVER (PARTITION BY id2 ORDER BY (v2 % 10) DESC) AS dr_v2
+    FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE dr_v2 <= 2;
+
+-- Window Top-N (DENSE_RANK top-2 per partition, ~100K partitions)
+SELECT pk, largest_v2 FROM (
+    SELECT (id3 % 100000) AS pk, v2 AS largest_v2,
+           DENSE_RANK() OVER (PARTITION BY (id3 % 100000) ORDER BY v2 DESC) AS dr_v2
+    FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE dr_v2 <= 2;

datafusion/core/tests/physical_optimizer/window_topn.rs

@@ -488,8 +488,12 @@ fn build_ranking_topn_plan(
     Ok(filter)
 }
 
-/// Build a RANK plan with NO ORDER BY: every row ties at rank 1 — degenerate.
-fn build_rank_no_order_by_plan(limit_value: i64) -> Result<Arc<dyn ExecutionPlan>> {
+/// Build a RANK / DENSE_RANK plan with NO ORDER BY: every row ties at rank 1 — degenerate.
+fn build_no_order_by_plan(
+    udwf_factory: fn() -> Arc<datafusion_expr::WindowUDF>,
+    udwf_name: &str,
+    limit_value: i64,
+) -> Result<Arc<dyn ExecutionPlan>> {
     let s = schema();
     let input: Arc<dyn ExecutionPlan> = Arc::new(PlaceholderRowExec::new(Arc::clone(&s)));
 
@@ -503,7 +507,7 @@ fn build_rank_no_order_by_plan(limit_value: i64) -> Result<Arc<dyn ExecutionPlan
     let partition_by = vec![col("pk", &s)?];
 
     let window_expr = Arc::new(StandardWindowExpr::new(
-        create_udwf_window_expr(&rank_udwf(), &[], &s, "rank".to_string(), false)?,
+        create_udwf_window_expr(&udwf_factory(), &[], &s, udwf_name.to_string(), false)?,
         &partition_by,
         &[], // empty ORDER BY
         Arc::new(WindowFrame::new_bounds(
@@ -520,7 +524,7 @@ fn build_rank_no_order_by_plan(limit_value: i64) -> Result<Arc<dyn ExecutionPlan
         true,
     )?);
 
-    let rk_col = Arc::new(Column::new("rank", 2));
+    let rk_col = Arc::new(Column::new(udwf_name, 2));
     let limit_lit = lit(ScalarValue::UInt64(Some(limit_value as u64)));
     let predicate = Arc::new(BinaryExpr::new(rk_col, Operator::LtEq, limit_lit));
     let filter: Arc<dyn ExecutionPlan> =
@@ -582,7 +586,7 @@ fn rank_no_order_by_no_change() -> Result<()> {
     // Without ORDER BY, every row ties at rank 1 — the optimization is
     // degenerate (entire input would be retained, ties storage unbounded).
     // The rule must skip.
-    let plan = build_rank_no_order_by_plan(3)?;
+    let plan = build_no_order_by_plan(rank_udwf, "rank", 3)?;
     let before = plan_str(plan.as_ref());
     let optimized = optimize(plan)?;
     let after = plan_str(optimized.as_ref());
@@ -593,17 +597,98 @@ fn rank_no_order_by_no_change() -> Result<()> {
     Ok(())
 }
 
+// ----------------------------------------------------------------------
+// DENSE_RANK rule tests
+// ----------------------------------------------------------------------
+
 #[test]
-fn dense_rank_no_change() -> Result<()> {
-    // DENSE_RANK is not yet supported by the rule. The plan must pass
-    // through unchanged.
+fn basic_dense_rank_dr_lteq_3() -> Result<()> {
     let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 3, Operator::LtEq)?;
+    let optimized = optimize(plan)?;
+    assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+    BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+      PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
+        PlaceholderRowExec
+    "#);
+    Ok(())
+}
+
+#[test]
+fn dense_rank_dr_lt_4_becomes_fetch_3() -> Result<()> {
+    let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 4, Operator::Lt)?;
+    let optimized = optimize(plan)?;
+    assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+    BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+      PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
+        PlaceholderRowExec
+    "#);
+    Ok(())
+}
+
+#[test]
+fn dense_rank_flipped_3_gteq_dr() -> Result<()> {
+    let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 3, Operator::GtEq)?;
+    let optimized = optimize(plan)?;
+    assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+    BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+      PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
+        PlaceholderRowExec
+    "#);
+    Ok(())
+}
+
+#[test]
+fn dense_rank_flipped_4_gt_dr_becomes_fetch_3() -> Result<()> {
+    let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 4, Operator::Gt)?;
+    let optimized = optimize(plan)?;
+    assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+    BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+      PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
+        PlaceholderRowExec
+    "#);
+    Ok(())
+}
+
+#[test]
+fn dense_rank_no_order_by_no_change() -> Result<()> {
+    // Without ORDER BY, every row ties at dense_rank 1 — the optimization
+    // is degenerate (entire input would be retained). The rule must skip.
+    let plan = build_no_order_by_plan(dense_rank_udwf, "dense_rank", 3)?;
     let before = plan_str(plan.as_ref());
     let optimized = optimize(plan)?;
     let after = plan_str(optimized.as_ref());
     assert_eq!(
         before, after,
-        "DENSE_RANK is unsupported and must not be rewritten"
+        "DENSE_RANK with empty ORDER BY must not be rewritten"
     );
     Ok(())
 }
+
+// ----------------------------------------------------------------------
+// Shared guard: `fn < 1` keeps nothing
+// ----------------------------------------------------------------------
+
+#[test]
+fn predicate_lt_1_no_change() -> Result<()> {
+    // `fn < 1` (and the flipped `1 > fn`) yields limit_n = 0. Since
+    // ROW_NUMBER / RANK / DENSE_RANK are always >= 1, the predicate keeps
+    // nothing and the rule must skip — a fetch=0 PartitionedTopK* would
+    // otherwise panic on its `k > 0` assertion at execution time.
+    type UdwfFactory = fn() -> Arc<datafusion_expr::WindowUDF>;
+    let cases: [(UdwfFactory, &str); 3] = [
+        (row_number_udwf, "row_number"),
+        (rank_udwf, "rank"),
+        (dense_rank_udwf, "dense_rank"),
+    ];
+    for (factory, name) in cases {
+        let plan = build_ranking_topn_plan(factory, name, 1, Operator::Lt)?;
+        let before = plan_str(plan.as_ref());
+        let optimized = optimize(plan)?;
+        let after = plan_str(optimized.as_ref());
+        assert_eq!(
+            before, after,
+            "`{name} < 1` (limit 0) must not be rewritten"
+        );
+    }
+    Ok(())
+}

datafusion/physical-optimizer/src/window_topn.rs

@@ -26,7 +26,7 @@
 //! ) WHERE rn <= K;
 //! ```
 //!
-//! or with `RANK()` in place of `ROW_NUMBER()`:
+//! or with `RANK()` / `DENSE_RANK()` in place of `ROW_NUMBER()`:
 //!
 //! ```sql
 //! SELECT * FROM (
@@ -40,8 +40,8 @@
 //! the `FilterExec` and `SortExec`.
 //!
 //! The appropriate [`WindowFnKind`] is forwarded to `PartitionedTopKExec`.
-//! RANK requires a non-empty `ORDER BY` clause (otherwise all rows tie at
-//! rank 1 and the optimization is degenerate).
+//! `RANK` and `DENSE_RANK` require a non-empty `ORDER BY` clause (otherwise
+//! all rows tie at rank 1 and the optimization is degenerate).
 //!
 //! See [`PartitionedTopKExec`] for details on the replacement operator.
 //!
@@ -68,9 +68,9 @@ use datafusion_physical_plan::sorts::partitioned_topk::{
 use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
 
-/// Physical optimizer rule that converts per-partition `ROW_NUMBER` and
-/// `RANK` top-K queries into a more efficient plan using
-/// [`PartitionedTopKExec`].
+/// Physical optimizer rule that converts per-partition `ROW_NUMBER`,
+/// `RANK`, and `DENSE_RANK` top-K queries into a more efficient plan
+/// using [`PartitionedTopKExec`].
 ///
 /// # Pattern Detected
 ///
@@ -86,12 +86,14 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
 /// ```text
 /// [optional ProjectionExec]
 ///   BoundedWindowAggExec(<ranking fn> PARTITION BY ... ORDER BY ...)
-///     PartitionedTopKExec(fn=<row_number|rank>, partition_keys, order_keys, fetch=K)
+///     PartitionedTopKExec(fn=<row_number|rank|dense_rank>, partition_keys, order_keys, fetch=K)
 /// ```
 ///
 /// The `FilterExec` is removed entirely. The `SortExec` is replaced by
-/// `PartitionedTopKExec`, which maintains a per-partition top-K heap (and,
-/// for `RANK`, a sibling ties `Vec`) instead of sorting the whole dataset.
+/// `PartitionedTopKExec`, which maintains per-partition top-K state (a
+/// heap for `ROW_NUMBER`, a heap plus boundary ties for `RANK`, a
+/// K-bounded distinct-ob map for `DENSE_RANK`) instead of sorting the
+/// whole dataset.
 ///
 /// # Supported Predicates
 ///
@@ -105,12 +107,19 @@ 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 → SortExec`
-/// - The window function is `ROW_NUMBER` or `RANK` (not `DENSE_RANK`)
+/// - 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`: 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 `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 `>`
 ///
@@ -141,6 +150,14 @@ impl WindowTopN {
         // Step 2: Extract limit from predicate (rn <= K, rn < K, etc.)
         let (col_idx, limit_n) = extract_window_limit(filter.predicate())?;
 
+        // `fn < 1` (or the flipped `1 > fn`) yields limit_n = 0. Since
+        // ROW_NUMBER / RANK / DENSE_RANK are always >= 1, such a predicate
+        // keeps nothing — bail out and let the ordinary FilterExec return
+        // the empty result. (The PartitionedTopK* operators require k > 0.)
+        if limit_n == 0 {
+            return None;
+        }
+
         // Step 3: Walk through optional ProjectionExec and RepartitionExec to find BoundedWindowAggExec
         let child = filter.input();
         let (window_exec, intermediates) = find_window_below(child)?;
@@ -159,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();
 
@@ -172,12 +211,21 @@ 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.
-        if matches!(fn_kind, WindowFnKind::Rank)
-            && 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;
         }
 
@@ -315,7 +363,8 @@ fn scalar_to_usize(value: &ScalarValue) -> Option<usize> {
 /// the UDF name. Returns:
 /// - `Some(WindowFnKind::RowNumber)` for `"row_number"`
 /// - `Some(WindowFnKind::Rank)` for `"rank"`
-/// - `None` for everything else (e.g. `dense_rank`)
+/// - `Some(WindowFnKind::DenseRank)` for `"dense_rank"`
+/// - `None` for everything else
 fn supported_window_fn(
     expr: &Arc<dyn datafusion_physical_expr::window::WindowExpr>,
 ) -> Option<WindowFnKind> {
@@ -325,6 +374,7 @@ fn supported_window_fn(
     match udf.fun().name() {
         "row_number" => Some(WindowFnKind::RowNumber),
         "rank" => Some(WindowFnKind::Rank),
+        "dense_rank" => Some(WindowFnKind::DenseRank),
         _ => None,
     }
 }