You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/14 07:28:05 UTC
(datafusion) branch gh-readonly-queue/main/pr-22885-ff774c6c3ce3349aaa9f778b3843fc94908504b8 created (now 79692c18f1)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-22885-ff774c6c3ce3349aaa9f778b3843fc94908504b8
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 79692c18f1 perf: Extend WindowTopN to support RANK (#22885)
This branch includes the following new commits:
new 79692c18f1 perf: Extend WindowTopN to support RANK (#22885)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: perf: Extend WindowTopN to support RANK (#22885)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22885-ff774c6c3ce3349aaa9f778b3843fc94908504b8
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 79692c18f1bb49bea09555a1b7a5236e4461eefb
Author: Subham Singhal <Su...@users.noreply.github.com>
AuthorDate: Tue Jul 14 12:57:42 2026 +0530
perf: Extend WindowTopN to support RANK (#22885)
## Which issue does this PR close?
- Related to https://github.com/apache/datafusion/issues/6899
(DENSE_RANK to follow in a separate PR will close it).
## Rationale for this change
PR #21479 introduced `WindowTopN` for `ROW_NUMBER` only; `RANK` and
`DENSE_RANK` were explicitly out of scope. This PR extends the rule to
`RANK`, replacing the full sort under `Filter(rk≤K) → Window(RANK) →
Sort` with a per-partition heap-of-K plus a boundary-tie buffer.
## What changes are included in this PR?
- **`datafusion/physical-plan/src/topk/mod.rs`** — new `pub(crate)
struct PartitionedTopKRank` (sibling of
`PartitionedTopK` from #23096) with per-partition `RankPartitionState {
TopKHeap, Vec<TieEntry> }`.
- **`datafusion/physical-plan/src/sorts/partitioned_topk.rs`** —
`WindowFnKind` enum (`RowNumber` / `Rank`).
`do_partitioned_topk` dispatches on `fn_kind` to
`PartitionedTopK::try_new` or `PartitionedTopKRank::try_new`;
- **`datafusion/physical-optimizer/src/window_topn.rs`** —
`is_row_number` → `supported_window_fn(expr) ->
Option<WindowFnKind>`; empty-`order_by` guard for RANK; `WindowFnKind`
plumbed through `PartitionedTopKExec::try_new`.
- **`datafusion/sqllogictest/test_files/window_topn.slt`** — RANK SLT
cases: basic, strict (`<`), flipped (`>=` / `>`),
boundary ties, ties spanning ob values, empty-`ORDER BY` (rule must NOT
fire), mixed window functions, ASC/DESC × NULLS FIRST/LAST, QUALIFY.
- **`datafusion/core/tests/physical_optimizer/window_topn.rs`** — 6 new
RANK rule unit tests covering predicate
matching, partition-by/order-by guards, dense_rank skip.
- **`benchmarks/queries/h2o/window.sql`** — six new RANK queries
(Q14–Q17, Q22, Q23) covering partition counts from ~100 to ~100K, low
and heavy tie densities.
h2o `window` benchmark, 10M-row `large` table, RANK top-2, 3-iteration
average. Toggle via `DATAFUSION_OPTIMIZER_ENABLE_WINDOW_TOPN`.
| Variant | Partitions | OFF (rule disabled) | ON (rule enabled) | Δ |
|---|---:|---:|---:|:---:|
| RANK low ties (`id3 % 100`) | ~100 | 305 ms | **107 ms** | **2.84×
faster** ✓ |
| RANK low ties (`id3 % 1000`) | ~1K | 263 ms | **120 ms** | **2.19×
faster** ✓ |
| RANK heavy ties (`id3 % 1000`, `v2 % 10` OB) | ~1K | 282 ms | **125
ms** | **2.25× faster** ✓ |
| RANK low ties (`id2`) | ~10K | 363 ms | **140 ms** | **2.59× faster**
✓ |
| RANK heavy ties (`id2`, `v2 % 10` OB) | ~10K | 291 ms | **143 ms** |
**2.04× faster** ✓ |
| RANK low ties (`id3 % 100K`) | ~100K | 241 ms | 422 ms | 1.75× slower
|
## Are these changes tested?
Yes:
- `cargo test -p datafusion-physical-plan --lib` — 1455 passed
- `cargo test -p datafusion-physical-optimizer --lib` — 27 passed
- `cargo test -p datafusion --test core_integration
physical_optimizer::window_topn::` — 13 passed (7 ROW_NUMBER + 6 RANK)
- `cargo test --test sqllogictests -- window_topn` — passed
## Are there any user-facing changes?
The existing `optimizer.enable_window_topn` config flag (default
`false`) now also covers `RANK` queries. No public API additions
---
benchmarks/queries/h2o/window.sql | 48 ++
.../core/tests/physical_optimizer/window_topn.rs | 192 ++++-
datafusion/physical-optimizer/src/window_topn.rs | 93 ++-
.../physical-plan/src/sorts/partitioned_topk.rs | 150 +++-
datafusion/physical-plan/src/topk/mod.rs | 794 ++++++++++++++++++++-
datafusion/sqllogictest/test_files/window_topn.slt | 482 ++++++++++++-
6 files changed, 1653 insertions(+), 106 deletions(-)
diff --git a/benchmarks/queries/h2o/window.sql b/benchmarks/queries/h2o/window.sql
index 346a8e4713..37df0a28ae 100644
--- a/benchmarks/queries/h2o/window.sql
+++ b/benchmarks/queries/h2o/window.sql
@@ -148,3 +148,51 @@ SELECT pk, largest2_v2 FROM (
ROW_NUMBER() OVER (PARTITION BY id3 % 100000 ORDER BY v2 DESC) AS order_v2
FROM large WHERE v2 IS NOT NULL
) sub_query WHERE order_v2 <= 2;
+
+-- Window Top-N (RANK top-2 per partition, ~100 partitions)
+-- The RANK queries below mirror the ROW_NUMBER cardinality sweep
+-- above and add heavy-ties variants. RANK semantics retain boundary
+-- ties (`WHERE rk <= K` may keep more than K rows per partition), so
+-- this exercises PartitionedTopKRank's ties-Vec path.
+SELECT pk, largest_v2 FROM (
+ SELECT (id3 % 100) AS pk, v2 AS largest_v2,
+ RANK() OVER (PARTITION BY (id3 % 100) ORDER BY v2 DESC) AS rk_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE rk_v2 <= 2;
+
+-- Window Top-N (RANK top-2 per partition, ~1K partitions)
+SELECT pkey, largest_v2 FROM (
+ SELECT (id3 % 1000) AS pkey, v2 AS largest_v2,
+ RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY v2 DESC) AS rk_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE rk_v2 <= 2;
+
+-- Window Top-N (RANK top-2 per partition, ~1K partitions, heavy ties)
+-- v2 % 10 forces 10 distinct OBY values, so most rows tie at the boundary
+-- and exercise PartitionedTopKRank's ties-Vec path.
+SELECT pkey, largest_v2 FROM (
+ SELECT (id3 % 1000) AS pkey, v2 AS largest_v2,
+ RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY (v2 % 10) DESC) AS rk_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE rk_v2 <= 2;
+
+-- Window Top-N (RANK top-2 per partition, ~10K partitions, low ties)
+SELECT id2, largest_v2 FROM (
+ SELECT id2, v2 AS largest_v2,
+ RANK() OVER (PARTITION BY id2 ORDER BY v2 DESC) AS rk_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE rk_v2 <= 2;
+
+-- Window Top-N (RANK top-2 per partition, ~10K partitions, heavy ties)
+SELECT id2, largest_v2 FROM (
+ SELECT id2, v2 AS largest_v2,
+ RANK() OVER (PARTITION BY id2 ORDER BY (v2 % 10) DESC) AS rk_v2
+ FROM large WHERE v2 IS NOT NULL
+) sub_query WHERE rk_v2 <= 2;
+
+-- Window Top-N (RANK top-2 per partition, ~100K partitions)
+SELECT pk, largest_v2 FROM (
+ SELECT (id3 % 100000) AS pk, v2 AS largest_v2,
+ 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;
diff --git a/datafusion/core/tests/physical_optimizer/window_topn.rs b/datafusion/core/tests/physical_optimizer/window_topn.rs
index e3f73a8535..07a1db127e 100644
--- a/datafusion/core/tests/physical_optimizer/window_topn.rs
+++ b/datafusion/core/tests/physical_optimizer/window_topn.rs
@@ -25,6 +25,7 @@ use datafusion_common::ScalarValue;
use datafusion_common::config::ConfigOptions;
use datafusion_expr::Operator;
use datafusion_expr::{WindowFrame, WindowFrameBound, WindowFrameUnits};
+use datafusion_functions_window::rank::{dense_rank_udwf, rank_udwf};
use datafusion_functions_window::row_number::row_number_udwf;
use datafusion_physical_expr::expressions::{BinaryExpr, Column, col, lit};
use datafusion_physical_expr::window::StandardWindowExpr;
@@ -226,7 +227,7 @@ fn basic_row_number_rn_lteq_3() -> Result<()> {
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fetch=3, partition=[pk@0], order=[val@1 ASC]
+ PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC]
PlaceholderRowExec
"#);
Ok(())
@@ -238,7 +239,7 @@ fn rn_lt_3_becomes_fetch_2() -> Result<()> {
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fetch=2, partition=[pk@0], order=[val@1 ASC]
+ PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@0], order=[val@1 ASC]
PlaceholderRowExec
"#);
Ok(())
@@ -300,7 +301,7 @@ fn flipped_3_gteq_rn() -> Result<()> {
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fetch=3, partition=[pk@0], order=[val@1 ASC]
+ PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC]
PlaceholderRowExec
"#);
Ok(())
@@ -418,8 +419,191 @@ fn with_projection_between() -> Result<()> {
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
ProjectionExec: expr=[pk@0 as pk, val@1 as val, row_number@2 as row_number]
BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fetch=3, partition=[pk@0], order=[val@1 ASC]
+ PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC]
PlaceholderRowExec
"#);
Ok(())
}
+
+// ----------------------------------------------------------------------
+// RANK rule tests
+// ----------------------------------------------------------------------
+
+/// Build: FilterExec(rk op limit) → BoundedWindowAggExec(<udwf> PBY pk OBY val) → SortExec(pk, val)
+///
+/// `udwf_factory` selects the window UDWF (rank, dense_rank, ...) and
+/// `udwf_name` is the column name produced by that UDWF (matters because
+/// the rule resolves the filter column by index, but the snapshot prints
+/// the name).
+fn build_ranking_topn_plan(
+ udwf_factory: fn() -> Arc<datafusion_expr::WindowUDF>,
+ udwf_name: &str,
+ limit_value: i64,
+ op: Operator,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let s = schema();
+ let input: Arc<dyn ExecutionPlan> = Arc::new(PlaceholderRowExec::new(Arc::clone(&s)));
+
+ let ordering = LexOrdering::new(vec![
+ PhysicalSortExpr::new_default(col("pk", &s)?).asc(),
+ PhysicalSortExpr::new_default(col("val", &s)?).asc(),
+ ])
+ .unwrap();
+
+ let sort: Arc<dyn ExecutionPlan> =
+ Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true));
+
+ let partition_by = vec![col("pk", &s)?];
+ let order_by = vec![PhysicalSortExpr::new_default(col("val", &s)?).asc()];
+
+ let window_expr = Arc::new(StandardWindowExpr::new(
+ create_udwf_window_expr(&udwf_factory(), &[], &s, udwf_name.to_string(), false)?,
+ &partition_by,
+ &order_by,
+ Arc::new(WindowFrame::new_bounds(
+ WindowFrameUnits::Rows,
+ WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
+ WindowFrameBound::CurrentRow,
+ )),
+ ));
+
+ let window: Arc<dyn ExecutionPlan> = Arc::new(BoundedWindowAggExec::try_new(
+ vec![window_expr],
+ sort,
+ InputOrderMode::Sorted,
+ true,
+ )?);
+
+ let rk_col = Arc::new(Column::new(udwf_name, 2));
+ let limit_lit = lit(ScalarValue::UInt64(Some(limit_value as u64)));
+ // Place column on whichever side matches the operator's expectation.
+ let predicate: Arc<dyn datafusion_physical_expr::PhysicalExpr> = match op {
+ Operator::LtEq | Operator::Lt => Arc::new(BinaryExpr::new(rk_col, op, limit_lit)),
+ Operator::GtEq | Operator::Gt => Arc::new(BinaryExpr::new(limit_lit, op, rk_col)),
+ _ => unreachable!("only </<=/>=/> are supported by the rule"),
+ };
+ let filter: Arc<dyn ExecutionPlan> =
+ Arc::new(FilterExec::try_new(predicate, window)?);
+
+ 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>> {
+ let s = schema();
+ let input: Arc<dyn ExecutionPlan> = Arc::new(PlaceholderRowExec::new(Arc::clone(&s)));
+
+ let ordering =
+ LexOrdering::new(vec![PhysicalSortExpr::new_default(col("pk", &s)?).asc()])
+ .unwrap();
+
+ let sort: Arc<dyn ExecutionPlan> =
+ Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true));
+
+ 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)?,
+ &partition_by,
+ &[], // empty ORDER BY
+ Arc::new(WindowFrame::new_bounds(
+ WindowFrameUnits::Rows,
+ WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
+ WindowFrameBound::CurrentRow,
+ )),
+ ));
+
+ let window: Arc<dyn ExecutionPlan> = Arc::new(BoundedWindowAggExec::try_new(
+ vec![window_expr],
+ sort,
+ InputOrderMode::Sorted,
+ true,
+ )?);
+
+ let rk_col = Arc::new(Column::new("rank", 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> =
+ Arc::new(FilterExec::try_new(predicate, window)?);
+
+ Ok(filter)
+}
+
+#[test]
+fn basic_rank_rk_lteq_3() -> Result<()> {
+ let plan = build_ranking_topn_plan(rank_udwf, "rank", 3, Operator::LtEq)?;
+ let optimized = optimize(plan)?;
+ assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+ BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
+ PlaceholderRowExec
+ "#);
+ Ok(())
+}
+
+#[test]
+fn rank_rk_lt_4_becomes_fetch_3() -> Result<()> {
+ let plan = build_ranking_topn_plan(rank_udwf, "rank", 4, Operator::Lt)?;
+ let optimized = optimize(plan)?;
+ assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+ BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
+ PlaceholderRowExec
+ "#);
+ Ok(())
+}
+
+#[test]
+fn rank_flipped_3_gteq_rk() -> Result<()> {
+ let plan = build_ranking_topn_plan(rank_udwf, "rank", 3, Operator::GtEq)?;
+ let optimized = optimize(plan)?;
+ assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+ BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
+ PlaceholderRowExec
+ "#);
+ Ok(())
+}
+
+#[test]
+fn rank_flipped_4_gt_rk_becomes_fetch_3() -> Result<()> {
+ let plan = build_ranking_topn_plan(rank_udwf, "rank", 4, Operator::Gt)?;
+ let optimized = optimize(plan)?;
+ assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+ BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
+ PlaceholderRowExec
+ "#);
+ Ok(())
+}
+
+#[test]
+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 before = plan_str(plan.as_ref());
+ let optimized = optimize(plan)?;
+ let after = plan_str(optimized.as_ref());
+ assert_eq!(
+ before, after,
+ "RANK with empty ORDER BY must not be rewritten"
+ );
+ Ok(())
+}
+
+#[test]
+fn dense_rank_no_change() -> Result<()> {
+ // DENSE_RANK is not yet supported by the rule. The plan must pass
+ // through unchanged.
+ let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 3, Operator::LtEq)?;
+ 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"
+ );
+ Ok(())
+}
diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs
index 40dbddfbdf..3f88e86c67 100644
--- a/datafusion/physical-optimizer/src/window_topn.rs
+++ b/datafusion/physical-optimizer/src/window_topn.rs
@@ -26,12 +26,27 @@
//! ) WHERE rn <= K;
//! ```
//!
+//! or with `RANK()` in place of `ROW_NUMBER()`:
+//!
+//! ```sql
+//! SELECT * FROM (
+//! SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk
+//! FROM t
+//! ) WHERE rk <= K;
+//! ```
+//!
//! And replaces the `FilterExec → BoundedWindowAggExec → SortExec` pipeline
//! with `BoundedWindowAggExec → PartitionedTopKExec(fetch=K)`, removing both
//! the `FilterExec` and `SortExec`.
//!
-//! See [`PartitionedTopKExec`]
-//! for details on the replacement operator.
+//! 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).
+//!
+//! See [`PartitionedTopKExec`] for details on the replacement operator.
+//!
+//! [`PartitionedTopKExec`]: datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec
+//! [`WindowFnKind`]: datafusion_physical_plan::sorts::partitioned_topk::WindowFnKind
use std::sync::Arc;
@@ -46,19 +61,22 @@ use datafusion_physical_expr::window::StandardWindowExpr;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::filter::FilterExec;
use datafusion_physical_plan::projection::ProjectionExec;
-use datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec;
+use datafusion_physical_plan::sorts::partitioned_topk::{
+ PartitionedTopKExec, WindowFnKind,
+};
use datafusion_physical_plan::sorts::sort::SortExec;
use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
-/// Physical optimizer rule that converts per-partition `ROW_NUMBER` top-K
-/// queries into a more efficient plan using [`PartitionedTopKExec`].
+/// Physical optimizer rule that converts per-partition `ROW_NUMBER` and
+/// `RANK` top-K queries into a more efficient plan using
+/// [`PartitionedTopKExec`].
///
/// # Pattern Detected
///
/// ```text
-/// FilterExec(rn <= K)
+/// FilterExec(<ranking fn output> <= K)
/// [optional ProjectionExec]
-/// BoundedWindowAggExec(ROW_NUMBER PARTITION BY ... ORDER BY ...)
+/// BoundedWindowAggExec(<ranking fn> PARTITION BY ... ORDER BY ...)
/// SortExec(partition_keys, order_keys)
/// ```
///
@@ -66,13 +84,13 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
///
/// ```text
/// [optional ProjectionExec]
-/// BoundedWindowAggExec(ROW_NUMBER PARTITION BY ... ORDER BY ...)
-/// PartitionedTopKExec(partition_keys, order_keys, fetch=K)
+/// BoundedWindowAggExec(<ranking fn> PARTITION BY ... ORDER BY ...)
+/// PartitionedTopKExec(fn=<row_number|rank>, partition_keys, order_keys, fetch=K)
/// ```
///
-/// The `FilterExec` is removed entirely (all output rows have `rn ∈ {1..K}`).
-/// The `SortExec` is replaced by `PartitionedTopKExec` which maintains a
-/// per-partition top-K heap instead of sorting the entire dataset.
+/// 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.
///
/// # Supported Predicates
///
@@ -86,9 +104,12 @@ 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` (not `RANK`, `DENSE_RANK`, etc.)
-/// - `ROW_NUMBER` has a `PARTITION BY` clause (global top-K is already
-/// handled by `SortExec` with `fetch`)
+/// - 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)
/// - The filter predicate compares the window output column to an integer
/// literal using `<=`, `<`, `>=`, or `>`
///
@@ -123,7 +144,7 @@ impl WindowTopN {
let child = filter.input();
let (window_exec, proj_between) = find_window_below(child)?;
- // Step 4: Verify col_idx references a ROW_NUMBER window output column
+ // Step 4: Verify col_idx references a supported window function output column
let input_field_count = window_exec.input().schema().fields().len();
if col_idx < input_field_count {
return None; // Filter is on an input column, not a window column
@@ -133,9 +154,7 @@ impl WindowTopN {
if window_expr_idx >= window_exprs.len() {
return None;
}
- if !is_row_number(&window_exprs[window_expr_idx]) {
- return None;
- }
+ let fn_kind = supported_window_fn(&window_exprs[window_expr_idx])?;
// Step 5: Verify child of window is SortExec
let sort_exec = window_exec.input().downcast_ref::<SortExec>()?;
@@ -151,12 +170,22 @@ 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()
+ {
+ return None;
+ }
+
// Step 7: Build PartitionedTopKExec using SortExec's expressions
let partitioned_topk = PartitionedTopKExec::try_new(
Arc::clone(sort_child),
sort_exec.expr().clone(),
partition_prefix_len,
limit_n,
+ fn_kind,
)
.ok()?;
@@ -287,20 +316,24 @@ fn scalar_to_usize(value: &ScalarValue) -> Option<usize> {
}
}
-/// Check if a window expression is `ROW_NUMBER`.
+/// Identify which supported ranking window function `expr` is.
///
/// Downcasts through `StandardWindowExpr` → `WindowUDFExpr` and checks
-/// that the UDF name is `"row_number"`. Returns `false` for all other
-/// window functions (e.g., `RANK`, `DENSE_RANK`, `SUM`).
-fn is_row_number(expr: &Arc<dyn datafusion_physical_expr::window::WindowExpr>) -> bool {
- let Some(swe) = expr.as_any().downcast_ref::<StandardWindowExpr>() else {
- return false;
- };
+/// the UDF name. Returns:
+/// - `Some(WindowFnKind::RowNumber)` for `"row_number"`
+/// - `Some(WindowFnKind::Rank)` for `"rank"`
+/// - `None` for everything else (e.g. `dense_rank`)
+fn supported_window_fn(
+ expr: &Arc<dyn datafusion_physical_expr::window::WindowExpr>,
+) -> Option<WindowFnKind> {
+ let swe = expr.as_any().downcast_ref::<StandardWindowExpr>()?;
let swfe = swe.get_standard_func_expr();
- let Some(udf) = swfe.as_any().downcast_ref::<WindowUDFExpr>() else {
- return false;
- };
- udf.fun().name() == "row_number"
+ let udf = swfe.as_any().downcast_ref::<WindowUDFExpr>()?;
+ match udf.fun().name() {
+ "row_number" => Some(WindowFnKind::RowNumber),
+ "rank" => Some(WindowFnKind::Rank),
+ _ => None,
+ }
}
/// Walk below a plan node looking for a [`BoundedWindowAggExec`].
diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs
index aee9e52568..730440a429 100644
--- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs
+++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs
@@ -23,10 +23,13 @@
//! FROM t WHERE rn <= N
//! ```
//!
-//! Instead of sorting the entire dataset, this operator maintains a
-//! [`TopK`](crate::topk::TopK) heap per partition (reusing the existing TopK implementation)
-//! and emits only the top-K rows per partition in sorted order
-//! `(partition_keys, order_keys)`.
+//! Instead of sorting the entire dataset, this operator delegates to a
+//! per-partition heap-of-K implementation (one variant for `ROW_NUMBER`
+//! and a sibling variant for `RANK`), both of which maintain one heap per
+//! distinct partition key while sharing a single [`arrow::row::RowConverter`],
+//! [`MemoryReservation`](datafusion_execution::memory_pool::MemoryReservation),
+//! and metrics set across all partitions, and emit only the top-K rows
+//! per partition in sorted order `(partition_keys, order_keys)`.
use std::fmt::{self, Formatter};
use std::sync::Arc;
@@ -43,12 +46,27 @@ use futures::TryStreamExt;
use crate::execution_plan::{Boundedness, EmissionType};
use crate::metrics::ExecutionPlanMetricsSet;
-use crate::topk::{PartitionedTopK, build_sort_fields};
+use crate::topk::{PartitionedTopK, PartitionedTopKRank, build_sort_fields};
use crate::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
PlanProperties, SendableRecordBatchStream, stream::RecordBatchStreamAdapter,
};
+/// Which window function `PartitionedTopKExec` is optimizing.
+///
+/// Different ranking functions have different per-partition retention rules:
+/// - [`RowNumber`](Self::RowNumber): exactly K rows per partition.
+/// - [`Rank`](Self::Rank): K rows plus any rows tied at the boundary
+/// ORDER BY value (RANK semantics — `WHERE rk <= K` may keep more
+/// than K rows when ties straddle the boundary).
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum WindowFnKind {
+ /// `ROW_NUMBER()` — keep exactly K rows per partition.
+ RowNumber,
+ /// `RANK()` — keep K rows plus any rows tied at the boundary.
+ Rank,
+}
+
/// Per-partition Top-K operator for window function queries.
///
/// # Background
@@ -89,9 +107,14 @@ use crate::{
/// DataSourceExec
/// ```
///
-/// Instead of sorting the entire dataset, this operator reads unsorted input,
-/// maintains a [`TopK`](crate::topk::TopK) heap per distinct partition key, and emits only the
-/// top-K rows per partition in sorted order `(partition_keys, order_keys)`.
+/// Instead of sorting the entire dataset, this operator reads unsorted input
+/// and delegates to a per-partition heap-of-K implementation (`PartitionedTopK`
+/// for `ROW_NUMBER` and `PartitionedTopKRank` for `RANK`), each maintaining
+/// one heap per distinct partition key while sharing a single
+/// [`arrow::row::RowConverter`] /
+/// [`MemoryReservation`](datafusion_execution::memory_pool::MemoryReservation)
+/// across all partitions, and emits only the top-K rows per partition in
+/// sorted order `(partition_keys, order_keys)`.
///
/// Cost: O(N log K) time instead of O(N log N), and O(K × P × row_size)
/// memory where K = fetch, P = number of distinct partitions.
@@ -139,9 +162,11 @@ use crate::{
///
/// # Limitations
///
-/// - Only activated when the window function is `ROW_NUMBER` with a
-/// `PARTITION BY` clause. Global top-K (no `PARTITION BY`) is already
-/// handled efficiently by `SortExec` with `fetch`.
+/// - Only activated when the window function is `ROW_NUMBER` or `RANK` with
+/// a `PARTITION BY` clause. `RANK` additionally requires a non-empty
+/// `ORDER BY` (with an empty `ORDER BY`, every row ties at rank 1 and the
+/// heap-of-K rewrite doesn't apply). Global top-K (no `PARTITION BY`) is
+/// already handled efficiently by `SortExec` with `fetch`.
/// - For very high cardinality partition keys (millions of distinct values),
/// both memory usage and runtime overhead can become significant. In such
/// cases, the sort-based plan is more robust. Therefore, this optimization
@@ -164,6 +189,9 @@ pub struct PartitionedTopKExec {
/// Derived from the filter predicate: `rn <= 3` → `fetch = 3`,
/// `rn < 3` → `fetch = 2`.
fetch: usize,
+ /// Which window function this operator is optimizing. Selects the
+ /// per-partition retention policy (see [`WindowFnKind`]).
+ fn_kind: WindowFnKind,
/// Execution metrics
metrics_set: ExecutionPlanMetricsSet,
/// Cached plan properties (output ordering, partitioning, etc.)
@@ -181,6 +209,8 @@ impl PartitionedTopKExec {
/// * `partition_prefix_len` - Number of leading expressions in `expr`
/// that form the partition key. Must be >= 1.
/// * `fetch` - Maximum rows to retain per partition (the K in "top-K").
+ /// * `fn_kind` - Which ranking window function this operator optimizes
+ /// ([`WindowFnKind::RowNumber`] or [`WindowFnKind::Rank`]).
///
/// # Example
///
@@ -191,6 +221,7 @@ impl PartitionedTopKExec {
/// LexOrdering([store ASC, revenue DESC]),
/// 1, // partition_prefix_len: 1 partition column (store)
/// 5, // fetch: keep top 5 per partition
+ /// WindowFnKind::RowNumber,
/// )
/// ```
pub fn try_new(
@@ -198,6 +229,7 @@ impl PartitionedTopKExec {
expr: LexOrdering,
partition_prefix_len: usize,
fetch: usize,
+ fn_kind: WindowFnKind,
) -> Result<Self> {
let cache = Self::compute_properties(&input, expr.clone())?;
Ok(Self {
@@ -205,6 +237,7 @@ impl PartitionedTopKExec {
expr,
partition_prefix_len,
fetch,
+ fn_kind,
metrics_set: ExecutionPlanMetricsSet::new(),
cache: Arc::new(cache),
})
@@ -231,6 +264,11 @@ impl PartitionedTopKExec {
self.fetch
}
+ /// Returns which window function this operator is optimizing.
+ pub fn fn_kind(&self) -> WindowFnKind {
+ self.fn_kind
+ }
+
/// Compute [`PlanProperties`] for this operator.
///
/// The output is sorted by `sort_exprs` (partition keys then order keys),
@@ -254,6 +292,10 @@ impl PartitionedTopKExec {
impl DisplayAs for PartitionedTopKExec {
fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result {
+ let fn_label = match self.fn_kind {
+ WindowFnKind::RowNumber => "row_number",
+ WindowFnKind::Rank => "rank",
+ };
match t {
DisplayFormatType::Default | DisplayFormatType::Verbose => {
let partition_exprs: Vec<String> = self.expr[..self.partition_prefix_len]
@@ -266,7 +308,8 @@ impl DisplayAs for PartitionedTopKExec {
.collect();
write!(
f,
- "PartitionedTopKExec: fetch={}, partition=[{}], order=[{}]",
+ "PartitionedTopKExec: fn={}, fetch={}, partition=[{}], order=[{}]",
+ fn_label,
self.fetch,
partition_exprs.join(", "),
order_exprs.join(", "),
@@ -281,6 +324,7 @@ impl DisplayAs for PartitionedTopKExec {
.iter()
.map(|e| format!("{e}"))
.collect();
+ writeln!(f, "fn={fn_label}")?;
writeln!(f, "fetch={}", self.fetch)?;
writeln!(f, "partition=[{}]", partition_exprs.join(", "))?;
writeln!(f, "order=[{}]", order_exprs.join(", "))
@@ -331,6 +375,7 @@ impl ExecutionPlan for PartitionedTopKExec {
self.expr.clone(),
self.partition_prefix_len,
self.fetch,
+ self.fn_kind,
)?))
}
@@ -354,6 +399,7 @@ impl ExecutionPlan for PartitionedTopKExec {
LexOrdering::new(self.expr[self.partition_prefix_len..].iter().cloned())
.expect("PartitionedTopKExec requires at least one order-by expression");
let fetch = self.fetch;
+ let fn_kind = self.fn_kind;
let batch_size = context.session_config().batch_size();
let runtime = Arc::clone(&context.runtime_env());
let metrics_set = self.metrics_set.clone();
@@ -367,6 +413,7 @@ impl ExecutionPlan for PartitionedTopKExec {
partition_sort_fields,
order_expr,
fetch,
+ fn_kind,
batch_size,
runtime,
metrics_set,
@@ -382,25 +429,29 @@ impl ExecutionPlan for PartitionedTopKExec {
}
}
-/// Read all input, feed each batch into a [`PartitionedTopK`] (which
-/// maintains one heap per distinct partition key), then emit results
-/// ordered by `(partition_keys, order_keys)`.
+/// Read all input, feed each batch into a per-partition top-K state
+/// (either [`PartitionedTopK`] for `ROW_NUMBER` or
+/// [`PartitionedTopKRank`] for `RANK`), then emit results ordered by
+/// `(partition_keys, order_keys)`.
///
/// # Phases
///
-/// 1. **Accumulation** — forward each input `RecordBatch` to
-/// [`PartitionedTopK::insert_batch`], which demultiplexes rows by
-/// partition key and dispatches them into the per-key heap. The
-/// `RowConverter` and `MemoryReservation` are shared across all
-/// partitions for this operator instance.
+/// 1. **Accumulation** — forward each input `RecordBatch` to the
+/// per-partition state's `insert_batch`. The `RowConverter` for
+/// ORDER BY columns, the operator's `MemoryReservation`, and the
+/// `TopKMetrics` are shared across all distinct partition keys for
+/// this operator instance.
///
-/// 2. **Emission** — [`PartitionedTopK::emit`] drains all heaps in
-/// sorted partition-key order, returning a coalesced batch stream.
+/// 2. **Emission** — `emit` drains all per-partition heaps in sorted
+/// partition-key order, returning a coalesced batch stream. For
+/// `RANK`, boundary-tied rows are materialized and emitted after
+/// each partition's heap rows.
///
/// # Cost
///
/// - Time: O(N log K) where N = total rows, K = fetch
/// - Memory: O(K × P × row_size) where P = number of distinct partitions
+/// plus, for RANK, the boundary ties' rows
#[expect(clippy::too_many_arguments)]
async fn do_partitioned_topk(
partition_id: usize,
@@ -410,26 +461,47 @@ async fn do_partitioned_topk(
partition_sort_fields: Vec<SortField>,
order_expr: LexOrdering,
fetch: usize,
+ fn_kind: WindowFnKind,
batch_size: usize,
runtime: Arc<RuntimeEnv>,
metrics_set: ExecutionPlanMetricsSet,
) -> Result<SendableRecordBatchStream> {
- let mut state = PartitionedTopK::try_new(
- partition_id,
- schema,
- partition_exprs,
- partition_sort_fields,
- order_expr,
- fetch,
- batch_size,
- &runtime,
- &metrics_set,
- )?;
-
- while let Some(batch) = input.next().await {
- state.insert_batch(&batch?)?;
+ match fn_kind {
+ WindowFnKind::RowNumber => {
+ let mut state = PartitionedTopK::try_new(
+ partition_id,
+ schema,
+ partition_exprs,
+ partition_sort_fields,
+ order_expr,
+ fetch,
+ batch_size,
+ &runtime,
+ &metrics_set,
+ )?;
+ while let Some(batch) = input.next().await {
+ state.insert_batch(&batch?)?;
+ }
+ drop(input);
+ state.emit()
+ }
+ WindowFnKind::Rank => {
+ let mut state = PartitionedTopKRank::try_new(
+ partition_id,
+ schema,
+ partition_exprs,
+ partition_sort_fields,
+ order_expr,
+ fetch,
+ batch_size,
+ &runtime,
+ &metrics_set,
+ )?;
+ while let Some(batch) = input.next().await {
+ state.insert_batch(&batch?)?;
+ }
+ drop(input);
+ state.emit()
+ }
}
- drop(input);
-
- state.emit()
}
diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs
index ee8675d718..1e3efff36b 100644
--- a/datafusion/physical-plan/src/topk/mod.rs
+++ b/datafusion/physical-plan/src/topk/mod.rs
@@ -307,6 +307,24 @@ impl TopKDynamicFilters {
// Guesstimate for memory allocation: estimated number of bytes used per row in the RowConverter
const ESTIMATED_BYTES_PER_ROW: usize = 20;
+/// Owned data of a row that was just evicted from a [`TopKHeap`].
+///
+/// Returned by [`TopKHeap::add`] so that callers (e.g. rank-aware
+/// wrappers that retain boundary ties) can decide whether to retain
+/// the evicted row externally. The underlying batch is captured
+/// before the heap's internal `RecordBatchStore` decrements the
+/// batch's use count, so the data remains accessible even if the
+/// heap drops its internal reference to the batch.
+#[derive(Debug, Clone)]
+pub(crate) struct EvictedRow {
+ /// The record batch the evicted row came from.
+ pub batch: RecordBatch,
+ /// Row index within `batch`.
+ pub index: usize,
+ /// Encoded ORDER BY tuple for the evicted row, in [`arrow::row`] format.
+ pub row_bytes: Vec<u8>,
+}
+
pub(crate) fn build_sort_fields(
ordering: &[PhysicalSortExpr],
schema: &SchemaRef,
@@ -895,12 +913,16 @@ impl TopKHeap {
/// Adds `row` to this heap. If inserting this new item would
/// increase the size past `k`, removes the previously smallest
/// item.
+ ///
+ /// Returns `Some(EvictedRow)` if an existing row was evicted to
+ /// make room for `row`, or `None` if the row was inserted into a
+ /// non-full heap.
fn add(
&mut self,
batch_entry: &mut RecordBatchEntry,
row: impl AsRef<[u8]>,
index: usize,
- ) {
+ ) -> Option<EvictedRow> {
let batch_id = batch_entry.id;
batch_entry.uses += 1;
@@ -911,6 +933,26 @@ impl TopKHeap {
if self.inner.len() == self.k {
let mut prev_min = self.inner.peek_mut().unwrap();
+ // Capture evicted row data before `unuse` (which may GC the
+ // batch from the store) and `replace_with` (which overwrites
+ // `prev_min` in place). The batch comes from `self.store` for
+ // cross-batch evictions, or directly from `batch_entry` when
+ // a row evicts another row from the same in-flight batch
+ // (entry not yet registered in the store).
+ let evicted_batch = if prev_min.batch_id == batch_entry.id {
+ batch_entry.batch.clone()
+ } else {
+ self.store
+ .get(prev_min.batch_id)
+ .map(|entry| entry.batch.clone())
+ .expect("evicted row's batch must be present in the store")
+ };
+ let evicted = EvictedRow {
+ batch: evicted_batch,
+ index: prev_min.index,
+ row_bytes: prev_min.row.clone(),
+ };
+
// Update batch use
if prev_min.batch_id == batch_entry.id {
batch_entry.uses -= 1;
@@ -924,12 +966,15 @@ impl TopKHeap {
prev_min.replace_with(row, batch_id, index);
self.owned_bytes += prev_min.owned_size();
+
+ Some(evicted)
} else {
let new_row = TopKRow::new(row, batch_id, index);
self.owned_bytes += new_row.owned_size();
// put the new row into the heap
self.inner.push(new_row);
- };
+ None
+ }
}
/// Returns the values stored in this heap, from values low to
@@ -1414,6 +1459,358 @@ impl PartitionedTopK {
}
}
+/// A run of rows from a single source [`RecordBatch`] that tied at the
+/// boundary when inserted. Stored as `(batch, indices)` and materialized
+/// at emit time via [`take_record_batch`].
+#[derive(Debug)]
+struct TieEntry {
+ batch: RecordBatch,
+ /// Indices into `batch` of the rows tied at the (then-current)
+ /// boundary. Always non-empty by construction.
+ row_indices: Vec<u32>,
+ /// `get_record_batch_memory_size(&batch)` captured at push time so
+ /// `RankPartitionState::size()` doesn't recurse through `batch`'s
+ /// columns on every `try_resize` call.
+ batch_bytes: usize,
+}
+
+/// Per-partition state for `RANK()` semantics.
+///
+/// Composes [`TopKHeap`] as the K-bounded core plus a sibling
+/// `Vec<TieEntry>` for boundary-tied rows. `RANK ≤ K` keeps the K
+/// best rows by ORDER BY plus every row tied at the K-th-best
+/// ORDER BY value — the boundary. So the total retained rows can
+/// exceed K when ties straddle the boundary.
+struct RankPartitionState {
+ heap: TopKHeap,
+ ties: Vec<TieEntry>,
+}
+
+impl RankPartitionState {
+ fn size(&self) -> usize {
+ let ties_buffer = self.ties.capacity() * size_of::<TieEntry>();
+ let ties_contents: usize = self
+ .ties
+ .iter()
+ .map(|t| t.row_indices.capacity() * size_of::<u32>() + t.batch_bytes)
+ .sum();
+ self.heap.size() + ties_buffer + ties_contents
+ }
+}
+
+/// Sibling to [`PartitionedTopK`] implementing `RANK()` semantics.
+///
+/// Per partition, retains the K-best rows plus every row tied at the
+/// K-th-best ORDER BY value (so `WHERE rk <= K` may keep more than K
+/// rows when ties straddle the boundary). Like [`PartitionedTopK`],
+/// the [`RowConverter`], [`MemoryReservation`], scratch [`Rows`]
+/// buffer, and [`TopKMetrics`] are shared across all partitions for
+/// this operator instance.
+///
+/// # Algorithm (per row)
+///
+/// For each incoming row, compare its encoded ORDER BY bytes against
+/// `heap.max()` — the K-th-best row, which is by definition the
+/// admission boundary. `heap.max()` is `None` until the heap fills
+/// to K rows:
+///
+/// - heap not full (`max() == None`) → forward to the heap
+/// - row's ob `==` max → push to ties (no heap call)
+/// - row's ob `>` max → drop
+/// - row's ob `<` max → forward to heap; on eviction, compare the
+/// new `heap.max()` to the evicted row's bytes: if equal, push
+/// evicted to ties (still tied at the new boundary's rank); else
+/// clear ties (boundary moved up, old ties no longer satisfy
+/// `rk ≤ K`)
+pub(crate) struct PartitionedTopKRank {
+ schema: SchemaRef,
+ metrics: TopKMetrics,
+ reservation: MemoryReservation,
+ /// ORDER BY expressions (excludes PARTITION BY).
+ expr: LexOrdering,
+ /// Encoder for ORDER BY columns. Reused across partitions.
+ row_converter: RowConverter,
+ /// Scratch row buffer reused across `insert_batch` calls.
+ scratch_rows: Rows,
+ /// PARTITION BY expressions.
+ partition_exprs: Vec<Arc<dyn PhysicalExpr>>,
+ /// Encoder for the partition key.
+ partition_converter: RowConverter,
+ /// Scratch row buffer for partition-key encoding. Reused across
+ /// `insert_batch` calls (cleared + appended each batch) so we
+ /// avoid allocating a fresh `Rows` buffer every batch.
+ partition_scratch_rows: Rows,
+ /// One rank state per distinct partition key seen so far.
+ states: HashMap<OwnedRow, RankPartitionState>,
+ k: usize,
+ batch_size: usize,
+}
+
+impl PartitionedTopKRank {
+ #[expect(clippy::too_many_arguments)]
+ pub(crate) fn try_new(
+ partition_id: usize,
+ schema: SchemaRef,
+ partition_exprs: Vec<Arc<dyn PhysicalExpr>>,
+ partition_sort_fields: Vec<SortField>,
+ order_expr: LexOrdering,
+ k: usize,
+ batch_size: usize,
+ runtime: &Arc<RuntimeEnv>,
+ metrics: &ExecutionPlanMetricsSet,
+ ) -> Result<Self> {
+ assert!(k > 0, "PartitionedTopKRank requires k > 0");
+ let reservation =
+ MemoryConsumer::new(format!("PartitionedTopKRank[{partition_id}]"))
+ .register(&runtime.memory_pool);
+
+ let order_sort_fields = build_sort_fields(&order_expr, &schema)?;
+ let row_converter = RowConverter::new(order_sort_fields)?;
+ let scratch_rows =
+ row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size);
+
+ let partition_converter = RowConverter::new(partition_sort_fields)?;
+ let partition_scratch_rows = partition_converter
+ .empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size);
+
+ Ok(Self {
+ schema,
+ metrics: TopKMetrics::new(metrics, partition_id),
+ reservation,
+ expr: order_expr,
+ row_converter,
+ scratch_rows,
+ partition_exprs,
+ partition_converter,
+ partition_scratch_rows,
+ states: HashMap::new(),
+ k,
+ batch_size,
+ })
+ }
+
+ /// Demultiplex `batch` rows by partition key, encode the ORDER BY
+ /// columns once for the whole batch, and feed each partition's
+ /// rows through the rank classifier into its dedicated heap and
+ /// ties Vec.
+ pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> {
+ let baseline = self.metrics.baseline.clone();
+ let _timer = baseline.elapsed_compute().timer();
+
+ let num_rows = batch.num_rows();
+ if num_rows == 0 {
+ return Ok(());
+ }
+
+ // Captured once so the per-tie push from this batch can reuse
+ // it (computing `get_record_batch_memory_size` is O(cols ×
+ // buffer walk) and we'd otherwise pay it per push and again
+ // per `try_resize` call).
+ let input_batch_bytes = get_record_batch_memory_size(batch);
+
+ // 1. Evaluate + encode partition columns into the reusable
+ // scratch (cleared then appended).
+ let pk_arrays: Vec<ArrayRef> = self
+ .partition_exprs
+ .iter()
+ .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows)))
+ .collect::<Result<_>>()?;
+ self.partition_scratch_rows.clear();
+ self.partition_converter
+ .append(&mut self.partition_scratch_rows, &pk_arrays)?;
+ let pk_rows = &self.partition_scratch_rows;
+
+ // 2. Demultiplex row indices by partition key (per-batch).
+ let mut groups: HashMap<OwnedRow, Vec<u32>> = HashMap::new();
+ for i in 0..num_rows {
+ groups
+ .entry(pk_rows.row(i).owned())
+ .or_default()
+ .push(i as u32);
+ }
+
+ // 3. Evaluate ORDER BY columns on the full batch and encode ONCE.
+ let ob_arrays: Vec<ArrayRef> = self
+ .expr
+ .iter()
+ .map(|e| e.expr.evaluate(batch).and_then(|v| v.into_array(num_rows)))
+ .collect::<Result<_>>()?;
+ self.scratch_rows.clear();
+ self.row_converter
+ .append(&mut self.scratch_rows, &ob_arrays)?;
+
+ // 4. Per-partition: classify each row and dispatch.
+ let k = self.k;
+ let mut replacements: usize = 0;
+
+ for (pk, indices) in groups {
+ let state = self.states.entry(pk).or_insert_with(|| RankPartitionState {
+ heap: TopKHeap::new(k),
+ ties: Vec::new(),
+ });
+
+ // Equal indices for THIS batch only. Coalesced into a single
+ // tie entry at the end of the partition's loop. Discarded if
+ // the boundary moves up mid-loop (those rows were tied to the
+ // old boundary, which is now strictly worse than the new K-th).
+ let mut equal_indices: Vec<u32> = Vec::new();
+ // Lazy-registered: only attached if at least one row reaches
+ // the heap from this batch in this partition.
+ let mut entry: Option<RecordBatchEntry> = None;
+
+ for &orig_idx in &indices {
+ let row = self.scratch_rows.row(orig_idx as usize);
+
+ // Classify against the current K-th-best (the heap top).
+ // `heap.max()` returns `None` while the heap is filling,
+ // so unclassified rows fall through to the heap path.
+ let classification = state
+ .heap
+ .max()
+ .map(|max_row| row.as_ref().cmp(max_row.row()));
+
+ match classification {
+ Some(Ordering::Equal) => {
+ equal_indices.push(orig_idx);
+ continue;
+ }
+ Some(Ordering::Greater) => continue,
+ Some(Ordering::Less) | None => {
+ // Heap path: heap not yet full, or row strictly
+ // better than the current boundary.
+ let entry_ref = entry.get_or_insert_with(|| {
+ state.heap.register_batch(batch.clone())
+ });
+ if let Some(EvictedRow {
+ batch: evicted_batch,
+ index: evicted_index,
+ row_bytes: evicted_bytes,
+ }) = state.heap.add(entry_ref, row, orig_idx as usize)
+ {
+ // Compare the new boundary (post-eviction heap
+ // top) against the evicted row's bytes — both
+ // already in encoded form, no clones needed.
+ let boundary_changed = state
+ .heap
+ .max()
+ .expect("heap was full to evict; must still be full")
+ .row()
+ != evicted_bytes.as_slice();
+ if boundary_changed {
+ // Boundary moved up — prior ties (across
+ // all prior batches) and equal_indices
+ // accumulated earlier in THIS batch were
+ // tied to the old boundary, now strictly
+ // worse than the new K-th-best. Discard.
+ state.ties.clear();
+ equal_indices.clear();
+ } else {
+ // Boundary unchanged — evicted row is tied
+ // at the (unchanged) boundary; push as a
+ // single-row entry.
+ let batch_bytes =
+ get_record_batch_memory_size(&evicted_batch);
+ state.ties.push(TieEntry {
+ batch: evicted_batch,
+ row_indices: vec![evicted_index as u32],
+ batch_bytes,
+ });
+ }
+ }
+ replacements += 1;
+ }
+ }
+ }
+
+ if let Some(e) = entry {
+ state.heap.insert_batch_entry(e);
+ state.heap.maybe_compact()?;
+ }
+
+ // Commit this batch's ties as a single entry.
+ if !equal_indices.is_empty() {
+ state.ties.push(TieEntry {
+ batch: batch.clone(),
+ row_indices: equal_indices,
+ batch_bytes: input_batch_bytes,
+ });
+ }
+ }
+
+ if replacements > 0 {
+ self.metrics.row_replacements.add(replacements);
+ }
+ self.reservation.try_resize(self.size())?;
+ Ok(())
+ }
+
+ /// Drain all heaps and ties in partition-key order and return the
+ /// rows as a stream of coalesced [`RecordBatch`]es ordered by
+ /// `(partition_keys, order_keys)`. Within a partition, heap rows
+ /// come first (sorted by ob), then tie rows (all sharing the
+ /// boundary ob).
+ pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> {
+ let Self {
+ schema,
+ metrics,
+ reservation: _,
+ expr: _,
+ row_converter: _,
+ scratch_rows: _,
+ partition_exprs: _,
+ partition_converter: _,
+ partition_scratch_rows: _,
+ mut states,
+ k: _,
+ batch_size,
+ } = self;
+ let _timer = metrics.baseline.elapsed_compute().timer();
+
+ let mut sorted_pks: Vec<OwnedRow> = states.keys().cloned().collect();
+ sorted_pks.sort();
+
+ let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size);
+
+ for pk in sorted_pks {
+ let RankPartitionState { mut heap, ties, .. } =
+ states.remove(&pk).expect("key from states.keys()");
+ if let Some(batch) = heap.emit()? {
+ (&batch).record_output(&metrics.baseline);
+ coalescer.push_batch(batch)?;
+ }
+ for tie in ties {
+ let indices = UInt32Array::from(tie.row_indices);
+ let tie_batch = take_record_batch(&tie.batch, &indices)?;
+ (&tie_batch).record_output(&metrics.baseline);
+ coalescer.push_batch(tie_batch)?;
+ }
+ }
+ coalescer.finish_buffered_batch()?;
+
+ let mut out: Vec<Result<RecordBatch>> = Vec::new();
+ while let Some(b) = coalescer.next_completed_batch() {
+ out.push(Ok(b));
+ }
+
+ Ok(Box::pin(RecordBatchStreamAdapter::new(
+ schema,
+ futures::stream::iter(out),
+ )))
+ }
+
+ /// Total memory currently held, including all per-partition states.
+ fn size(&self) -> usize {
+ size_of::<Self>()
+ + self.row_converter.size()
+ + self.partition_converter.size()
+ + self.scratch_rows.size()
+ + self.partition_scratch_rows.size()
+ + self.states.values().map(|s| s.size()).sum::<usize>()
+ + self.states.capacity()
+ * (size_of::<OwnedRow>() + size_of::<RankPartitionState>())
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -2376,4 +2773,397 @@ mod tests {
);
Ok(())
}
+
+ // ====================================================================
+ // PartitionedTopKRank operator tests
+ //
+ // These mirror the PartitionedTopK tests above plus three RANK-specific
+ // cases for the Equal / boundary-shift / boundary-unchanged-eviction
+ // arms in `PartitionedTopKRank::insert_batch`.
+ // ====================================================================
+
+ /// Builds a `(pk Int32, val Int32)` schema and a `PartitionedTopKRank`
+ /// keyed on `pk ASC` (partition) and `val ASC` (ORDER BY).
+ fn build_partitioned_topk_rank(
+ k: usize,
+ ) -> Result<(Arc<Schema>, PartitionedTopKRank)> {
+ build_partitioned_topk_rank_with_opts(k, SortOptions::default(), false)
+ }
+
+ /// Variant of [`build_partitioned_topk_rank`] that lets the test pick
+ /// the `val` column's `SortOptions` (direction, null ordering) and
+ /// nullability.
+ fn build_partitioned_topk_rank_with_opts(
+ k: usize,
+ val_sort_options: SortOptions,
+ val_nullable: bool,
+ ) -> Result<(Arc<Schema>, PartitionedTopKRank)> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("pk", DataType::Int32, false),
+ Field::new("val", DataType::Int32, val_nullable),
+ ]));
+
+ let pk_expr: Arc<dyn PhysicalExpr> = col("pk", schema.as_ref())?;
+ let pk_sort_expr = PhysicalSortExpr {
+ expr: Arc::clone(&pk_expr),
+ options: SortOptions::default(),
+ };
+ let val_sort_expr = PhysicalSortExpr {
+ expr: col("val", schema.as_ref())?,
+ options: val_sort_options,
+ };
+
+ let partition_sort_fields = build_sort_fields(&[pk_sort_expr], &schema)?;
+ let order_expr = LexOrdering::from([val_sort_expr]);
+
+ let state = PartitionedTopKRank::try_new(
+ 0,
+ Arc::clone(&schema),
+ vec![pk_expr],
+ partition_sort_fields,
+ order_expr,
+ k,
+ 8, // batch_size
+ &Arc::new(RuntimeEnv::default()),
+ &ExecutionPlanMetricsSet::new(),
+ )?;
+ Ok((schema, state))
+ }
+
+ /// Multiple distinct partition keys interleaved within a single
+ /// input batch — the per-batch demux, per-partition heap eviction,
+ /// and partition-key-ordered emit must all behave correctly. No
+ /// ties: result should match a `ROW_NUMBER` top-K under the same K.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_multi_partition_within_batch() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_rank(2)?;
+
+ // pk=1 vals: 10, 5, 8 → top-2 ASC = [5, 8]
+ // pk=2 vals: 20, 15 → top-2 ASC = [15, 20]
+ // pk=3 vals: 7 → top-2 ASC = [7]
+ let batch =
+ pk_val_batch(&schema, vec![1, 2, 1, 2, 1, 3], vec![10, 20, 5, 15, 8, 7])?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 5 |",
+ "| 1 | 8 |",
+ "| 2 | 15 |",
+ "| 2 | 20 |",
+ "| 3 | 7 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// State must accumulate across `insert_batch` calls. A row in
+ /// batch 2 that's strictly better than the existing K-th must
+ /// evict it; an evicted row whose bytes match the new boundary
+ /// becomes a `TieEntry` pinned to the prior batch.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_cross_batch_eviction() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_rank(2)?;
+
+ // Batch 1: pk=1 fills the heap with [50, 40].
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![50, 40])?)?;
+
+ // Batch 2: pk=1 sees a smaller value (10) — it must evict 50;
+ // 60 > 40 so it's dropped. pk=2 appears mid-stream.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 2, 1], vec![10, 99, 60])?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 10 |",
+ "| 1 | 40 |",
+ "| 2 | 99 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// Empty input must produce an empty output stream, not panic.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_empty_input() -> Result<()> {
+ let (_schema, state) = build_partitioned_topk_rank(3)?;
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert!(results.is_empty(), "empty input → empty output");
+ Ok(())
+ }
+
+ /// `fetch = 1` is a common case (rk = 1 filter) and exercises the
+ /// boundary-defined-immediately path: after the first admission per
+ /// partition, `heap.max()` is `Some`, so every subsequent row goes
+ /// through full Equal/Greater/Less classification.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_fetch_one() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_rank(1)?;
+ state.insert_batch(&pk_val_batch(
+ &schema,
+ vec![1, 1, 2, 2, 3],
+ vec![3, 1, 9, 4, 7],
+ )?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 1 |",
+ "| 2 | 4 |",
+ "| 3 | 7 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// `ORDER BY val DESC` exercises the shared encoder's sort-direction
+ /// handling: the row converter flips the sort sign for `val` so
+ /// larger values compare smaller in row-encoded form. Each
+ /// partition keeps its top-K *largest* values.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_desc_ordering() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_rank_with_opts(
+ 2,
+ SortOptions {
+ descending: true,
+ nulls_first: false,
+ },
+ false,
+ )?;
+
+ // pk=1 vals: 10, 5, 8, 12 → top-2 DESC = [12, 10]
+ // pk=2 vals: 20, 15, 25 → top-2 DESC = [25, 20]
+ let batch = pk_val_batch(
+ &schema,
+ vec![1, 2, 1, 2, 1, 1, 2],
+ vec![10, 20, 5, 15, 8, 12, 25],
+ )?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 12 |",
+ "| 1 | 10 |",
+ "| 2 | 25 |",
+ "| 2 | 20 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// NULL sort values exercise the shared encoder's null-ordering
+ /// handling. With `ASC NULLS LAST`, NULLs sort *after* every
+ /// non-NULL value, so a partition whose only non-NULL value beats
+ /// a NULL must evict the NULL when `K = 1`. A partition that holds
+ /// only NULLs must still emit them.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_nulls_last_ordering() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_rank_with_opts(
+ 1,
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ true,
+ )?;
+
+ // pk=1 vals: NULL, 7, NULL → top-1 ASC NULLS LAST = [7]
+ // pk=2 vals: NULL → top-1 = [NULL]
+ // pk=3 vals: NULL, 4, 2 → top-1 = [2]
+ let batch = nullable_pk_val_batch(
+ &schema,
+ vec![1, 2, 1, 1, 3, 3, 3],
+ vec![None, None, Some(7), None, None, Some(4), Some(2)],
+ )?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 7 |",
+ "| 2 | |",
+ "| 3 | 2 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// `ASC NULLS FIRST` (the `SortOptions::default()`) sorts NULLs
+ /// *before* every non-NULL value, so under `fetch = K` a partition's
+ /// NULLs are kept preferentially over larger non-NULL values.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_nulls_first_ordering() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_rank_with_opts(
+ 2,
+ SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ true,
+ )?;
+
+ // pk=1 vals: NULL, 5, NULL, 8 → top-2 ASC NULLS FIRST = [NULL, NULL]
+ // pk=2 vals: 7, NULL → top-2 = [NULL, 7]
+ // pk=3 vals: 3, 1 → top-2 = [1, 3]
+ let batch = nullable_pk_val_batch(
+ &schema,
+ vec![1, 2, 1, 3, 1, 2, 1, 3],
+ vec![
+ None,
+ Some(7),
+ Some(5),
+ Some(3),
+ None,
+ None,
+ Some(8),
+ Some(1),
+ ],
+ )?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | |",
+ "| 1 | |",
+ "| 2 | |",
+ "| 2 | 7 |",
+ "| 3 | 1 |",
+ "| 3 | 3 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// RANK-specific: heap fills with K rows tied at the same OB value,
+ /// then more rows at that same value arrive. They take the Equal arm
+ /// (heap is full, `heap.max() == row`) and accumulate as ties, while
+ /// strictly-greater rows are dropped. All retained rows have rank 1.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_boundary_ties_retained() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_rank(2)?;
+
+ // pk=1 vals: 5, 5, 10, 5
+ // - first two 5s fill the heap (max=None until heap reaches K=2)
+ // - third row 10 > 5 → drop (Greater)
+ // - fourth row 5 == 5 → push to ties (Equal)
+ // Sorted RANKs: 5→1, 5→1, 5→1, 10→4. WHERE rk ≤ 2 keeps the three 5s.
+ let batch = pk_val_batch(&schema, vec![1, 1, 1, 1], vec![5, 5, 10, 5])?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 5 |",
+ "| 1 | 5 |",
+ "| 1 | 5 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// RANK-specific: heap fills with K rows tied at value V, equal_indices
+ /// accumulate at V, then a strictly-better row arrives whose admission
+ /// shifts the boundary strictly below V. The boundary-changed branch
+ /// must clear both `state.ties` and the in-flight `equal_indices` —
+ /// otherwise the now-rank-> K rows at value V would leak into output.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_boundary_shifts_clears_ties() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_rank(2)?;
+
+ // pk=1 vals: 10, 10, 10, 5, 3
+ // - first two 10s fill heap (max=10)
+ // - third 10 → Equal → equal_indices=[2]
+ // - 5 < 10 → admit, evict 10 → heap={5,10}, max=10 (unchanged).
+ // Push evicted to ties: ties=[10@curr_batch[ev_idx]].
+ // - 3 < 10 → admit, evict 10 → heap={3,5}, max=5 (CHANGED).
+ // Clear ties AND equal_indices.
+ // Sorted RANKs: 3→1, 5→2, 10→3, 10→3, 10→3. WHERE rk ≤ 2 → [3, 5].
+ let batch = pk_val_batch(&schema, vec![1, 1, 1, 1, 1], vec![10, 10, 10, 5, 3])?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 3 |",
+ "| 1 | 5 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// RANK-specific: heap has multiple rows at boundary value V, then a
+ /// strictly-better row arrives. The heap evicts one V (popping
+ /// `prev_min`), but `heap.max()` is still V — boundary unchanged.
+ /// The evicted V row must be pushed as a `TieEntry`; without that
+ /// branch a `rk <= K` query would silently lose a tied row.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_eviction_at_unchanged_boundary() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_rank(2)?;
+
+ // pk=1 vals: 10, 10, 5
+ // - first two 10s fill the heap (max=10)
+ // - 5 < 10 → admit, evict 10. New heap={5,10}, max=10 (unchanged).
+ // Push the evicted 10 to ties.
+ // Sorted RANKs: 5→1, 10→2, 10→2. WHERE rk ≤ 2 → all 3 rows.
+ let batch = pk_val_batch(&schema, vec![1, 1, 1], vec![10, 10, 5])?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 5 |",
+ "| 1 | 10 |",
+ "| 1 | 10 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
}
diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt
index bf9ce26b35..2eb72f519b 100644
--- a/datafusion/sqllogictest/test_files/window_topn.slt
+++ b/datafusion/sqllogictest/test_files/window_topn.slt
@@ -64,7 +64,7 @@ logical_plan
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]
02)--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], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
+03)----PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
# Test 3: rn < 4 should give same results (fetch=3)
@@ -131,7 +131,7 @@ logical_plan
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]
02)--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], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
+03)----PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
# Test 7: Filter on data column (not window output) — should NOT optimize
@@ -164,7 +164,7 @@ EXPLAIN SELECT * FROM (
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]
02)--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], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
+03)----PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
@@ -236,19 +236,20 @@ physical_plan
33)│ PartitionedTopKExec │
34)│ -------------------- │
35)│ fetch: 3 │
-36)│ │
-37)│ order: │
-38)│ [val@2 ASC NULLS LAST] │
-39)│ │
-40)│ partition: [pk@1] │
-41)└─────────────┬─────────────┘
-42)┌─────────────┴─────────────┐
-43)│ DataSourceExec │
-44)│ -------------------- │
-45)│ bytes: 480 │
-46)│ format: memory │
-47)│ rows: 1 │
-48)└───────────────────────────┘
+36)│ fn: row_number │
+37)│ │
+38)│ order: │
+39)│ [val@2 ASC NULLS LAST] │
+40)│ │
+41)│ partition: [pk@1] │
+42)└─────────────┬─────────────┘
+43)┌─────────────┴─────────────┐
+44)│ DataSourceExec │
+45)│ -------------------- │
+46)│ bytes: 480 │
+47)│ format: memory │
+48)│ rows: 1 │
+49)└───────────────────────────┘
statement ok
SET datafusion.explain.format = indent;
@@ -308,10 +309,9 @@ EXPLAIN SELECT * FROM (
----
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, rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as rnk]
-02)--FilterExec: rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 <= 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, rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNB [...]
-04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[false]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--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, rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOU [...]
+03)----PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
# Test 14: Filter on rn AND rnk — compound predicate should NOT optimize
query TT
@@ -360,7 +360,7 @@ EXPLAIN SELECT * FROM (
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, window_topn_t.id] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk, window_topn_t.id] 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, window_topn_t.id] 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], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=3, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST]
+03)----PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
@@ -391,7 +391,7 @@ EXPLAIN SELECT * FROM (
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=3, partition=[id@0], order=[val@2 ASC NULLS LAST]
+03)----PartitionedTopKExec: fn=row_number, fetch=3, partition=[id@0], order=[val@2 ASC NULLS LAST]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
# Test 19: Overlapping keys correctness (each id is unique, so rn=1 for all)
@@ -426,7 +426,7 @@ EXPLAIN SELECT * FROM (
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.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 DESC]
+03)----PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 DESC]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
# Test 21: Correctness for PARTITION BY pk ORDER BY pk, val DESC
@@ -460,7 +460,7 @@ EXPLAIN SELECT * FROM (
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.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 DESC]
+03)----PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 DESC]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
@@ -494,7 +494,7 @@ QUALIFY 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]
02)--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], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
+03)----PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
# Test 30: QUALIFY with < operator
@@ -522,10 +522,9 @@ QUALIFY rnk <= 3;
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rnk]
-02)--FilterExec: rank() 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=[rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() 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], 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)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() 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], mode=[Sorted]
+03)----PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
SET datafusion.explain.physical_plan_only = false;
@@ -601,7 +600,7 @@ EXPLAIN SELECT * FROM (
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=2, partition=[pk@1], order=[val@2 ASC]
+03)----PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@1], order=[val@2 ASC]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
query TT
@@ -612,7 +611,7 @@ EXPLAIN SELECT * FROM (
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----PartitionedTopKExec: fetch=2, partition=[pk@1], order=[val@2 DESC NULLS LAST]
+03)----PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@1], order=[val@2 DESC NULLS LAST]
04)------DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
@@ -621,6 +620,427 @@ SET datafusion.explain.physical_plan_only = false;
statement ok
DROP TABLE window_topn_nulls;
+###############################################################################
+# RANK() tests
+###############################################################################
+#
+# RANK semantics differ from ROW_NUMBER in that ties at the boundary are
+# retained (`WHERE rk <= K` may keep more than K rows per partition). The
+# tests below exercise both the boundary-Equal case (incoming row tied
+# with current K-th-best) and the boundary-unchanged-after-eviction case
+# (PartitionedTopKRank: heap evicts a tied row → push to per-partition
+# `ties` Vec).
+
+# Table designed to produce ties at and around the rank-K boundary
+statement ok
+CREATE TABLE window_topn_rank_t (id INT, pk INT, val INT) AS VALUES
+ -- pk=1: ties at rank 2 (val=20 thrice), val=30 jumps to rank 5
+ (1, 1, 10),
+ (2, 1, 20),
+ (3, 1, 20),
+ (4, 1, 20),
+ (5, 1, 30),
+ -- pk=2: distinct values, no ties
+ (6, 2, 5),
+ (7, 2, 15),
+ (8, 2, 25),
+ -- pk=3: 100 then four 200s — exercises the boundary-unchanged-with-eviction
+ -- case from the design doc's worked example (heap fills with three 200s,
+ -- the fourth ties, then 100 evicts a 200 but new boundary is still 200,
+ -- so the evicted 200 must move to ties)
+ (9, 3, 100),
+ (10, 3, 200),
+ (11, 3, 200),
+ (12, 3, 200),
+ (13, 3, 200),
+ (14, 3, 300);
+
+# Test R1: Basic RANK correctness with ties at the boundary.
+# Expected per partition (RANK ASC, rk <= 3):
+# pk=1: 10 (rk=1), 20×3 (rk=2 each) → 4 rows
+# pk=2: 5 (rk=1), 15 (rk=2), 25 (rk=3) → 3 rows
+# pk=3: 100 (rk=1), 200×4 (rk=2 each) → 5 rows
+# Total: 12 rows kept, val=30 (pk=1, rk=5) and val=300 (pk=3, rk=6) dropped.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t
+) WHERE rk <= 3;
+----
+1 1 10
+10 3 200
+11 3 200
+12 3 200
+13 3 200
+2 1 20
+3 1 20
+4 1 20
+6 2 5
+7 2 15
+8 2 25
+9 3 100
+
+# Test R2: EXPLAIN shows PartitionedTopKExec with fn=rank
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t
+) WHERE rk <= 3;
+----
+logical_plan
+01)Projection: window_topn_rank_t.id, window_topn_rank_t.pk, window_topn_rank_t.val, rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rk
+02)--Filter: rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(3)
+03)----WindowAggr: windowExpr=[[rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+04)------TableScan: window_topn_rank_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_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
+02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+03)----PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Test R3: rk < 4 should give the same results (fetch = K-1 = 3)
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t
+) WHERE rk < 4;
+----
+1 1 10
+10 3 200
+11 3 200
+12 3 200
+13 3 200
+2 1 20
+3 1 20
+4 1 20
+6 2 5
+7 2 15
+8 2 25
+9 3 100
+
+# Test R4: Flipped predicate `3 >= rk` should also trigger optimization
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t
+) WHERE 3 >= rk;
+----
+1 1 10
+10 3 200
+11 3 200
+12 3 200
+13 3 200
+2 1 20
+3 1 20
+4 1 20
+6 2 5
+7 2 15
+8 2 25
+9 3 100
+
+# Test R5: Flipped strict `4 > rk` should also trigger optimization (fetch=3)
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t
+) WHERE 4 > rk;
+----
+1 1 10
+10 3 200
+11 3 200
+12 3 200
+13 3 200
+2 1 20
+3 1 20
+4 1 20
+6 2 5
+7 2 15
+8 2 25
+9 3 100
+
+# Test R6: RANK without PARTITION BY — should NOT trigger the optimization
+# (global top-K with ties; SortExec with fetch handles this without our rule).
+# Use window_topn_rank_t (still alive); window_topn_t was dropped earlier.
+query II rowsort
+SELECT id, val FROM (
+ SELECT *, RANK() OVER (ORDER BY val) as rk FROM window_topn_rank_t
+) WHERE rk <= 3;
+----
+1 10
+6 5
+7 15
+
+# Test R7: RANK with multi-column PARTITION BY
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk, id ORDER BY val) as rk FROM window_topn_rank_t
+) WHERE rk <= 1;
+----
+1 1 10
+10 3 200
+11 3 200
+12 3 200
+13 3 200
+14 3 300
+2 1 20
+3 1 20
+4 1 20
+5 1 30
+6 2 5
+7 2 15
+8 2 25
+9 3 100
+
+# Test R8: Verify multi-column partition plan still uses fn=rank
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk, id ORDER BY val) as rk FROM window_topn_rank_t
+) WHERE rk <= 1;
+----
+logical_plan
+01)Projection: window_topn_rank_t.id, window_topn_rank_t.pk, window_topn_rank_t.val, rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rk
+02)--Filter: rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1)
+03)----WindowAggr: windowExpr=[[rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+04)------TableScan: window_topn_rank_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_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
+02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+03)----PartitionedTopKExec: fn=rank, fetch=1, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Test R9: RANK with DESC ordering
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val DESC) as rk FROM window_topn_rank_t
+) WHERE rk <= 1;
+----
+14 3 300
+5 1 30
+8 2 25
+
+# Test R10: Mixed window functions — RANK + ROW_NUMBER in the same query.
+# Filter is on the RANK column; rule should still fire (matches by col_idx).
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *,
+ ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+ RANK() OVER (PARTITION BY pk ORDER BY val) as rk
+ FROM window_topn_rank_t
+) WHERE rk <= 1;
+----
+1 1 10
+6 2 5
+9 3 100
+
+# Test R11: QUALIFY form (parser desugars to the same plan)
+query IIII rowsort
+SELECT id, pk, val,
+ RANK() OVER (PARTITION BY pk ORDER BY val) as rk
+FROM window_topn_rank_t
+QUALIFY rk <= 1;
+----
+1 1 10 1
+6 2 5 1
+9 3 100 1
+
+statement ok
+DROP TABLE window_topn_rank_t;
+
+###############################################################################
+# RANK() — equality predicate (negative: rule supports only </<=/>=/>)
+###############################################################################
+#
+# `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.
+
+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),
+ (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).
+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
+4 2 5
+
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_eq_t
+) WHERE rk = 1;
+----
+logical_plan
+01)Projection: window_topn_rank_eq_t.id, window_topn_rank_eq_t.pk, window_topn_rank_eq_t.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 AS rk
+02)--Filter: 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(1)
+03)----WindowAggr: windowExpr=[[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]]
+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]
+
+statement ok
+DROP TABLE window_topn_rank_eq_t;
+
+###############################################################################
+# RANK() — dense-ties boundary preservation
+###############################################################################
+#
+# Heap fills with K=3 rows tied at the same value, then a strictly-better
+# row arrives. The heap evicts one of the tied rows, but the new
+# K-th-best is still tied with the evicted row (boundary unchanged).
+# PartitionedTopKRank must push the evicted row into `ties` rather than
+# discarding it. Without that branch, a `rk <= 3` query loses the
+# evicted tied row.
+
+statement ok
+CREATE TABLE window_topn_rank_dense_t (id INT, pk INT, val INT) AS VALUES
+ -- ten rows with the same val + one strictly-better row
+ (1, 1, 10), (2, 1, 10), (3, 1, 10), (4, 1, 10), (5, 1, 10),
+ (6, 1, 10), (7, 1, 10), (8, 1, 10), (9, 1, 10), (10, 1, 10),
+ (11, 1, 5);
+
+# Test R14: With `rk <= 3`, every row should be retained:
+# - val=5 → rk=1
+# - val=10 (×10) → rk=2 each
+# Total 11 rows. If the boundary-unchanged-eviction branch ever drops a
+# tied row, this query would return fewer than 11.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_dense_t
+) WHERE rk <= 3;
+----
+1 1 10
+10 1 10
+11 1 5
+2 1 10
+3 1 10
+4 1 10
+5 1 10
+6 1 10
+7 1 10
+8 1 10
+9 1 10
+
+# Test R15: rule fired (no FilterExec/SortExec)
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_dense_t
+) WHERE rk <= 3;
+----
+logical_plan
+01)Projection: window_topn_rank_dense_t.id, window_topn_rank_dense_t.pk, window_topn_rank_dense_t.val, rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rk
+02)--Filter: rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(3)
+03)----WindowAggr: windowExpr=[[rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+04)------TableScan: window_topn_rank_dense_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_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
+02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+03)----PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+DROP TABLE window_topn_rank_dense_t;
+
+###############################################################################
+# RANK() — NULL handling in ORDER BY
+###############################################################################
+#
+# RANK ASSIGNMENTS WITH NULLS:
+# ORDER BY val ASC NULLS LAST → non-NULLs ranked first, NULLs at the end
+# ORDER BY val DESC NULLS LAST → same shape, different non-NULL order
+# ORDER BY val ASC NULLS FIRST → NULLs all tie at rank 1
+# ORDER BY val DESC NULLS FIRST → NULLs all tie at rank 1
+#
+# Multiple NULLs in the same partition all share the same rank (they're
+# tied under the encoded ORDER BY).
+
+statement ok
+CREATE TABLE window_topn_rank_null_t (id INT, pk INT, val INT) AS VALUES
+ -- pk=1: distinct vals plus one NULL → ASC NULLS LAST → 1,2,3,NULL ranks 1,2,3,4
+ (1, 1, 1), (2, 1, 2), (3, 1, 3), (4, 1, NULL),
+ -- pk=2: one non-NULL plus two NULLs → ASC NULLS LAST → 5,NULL,NULL ranks 1,2,2
+ (5, 2, 5), (6, 2, NULL), (7, 2, NULL);
+
+# Test R16: ASC NULLS LAST, rk <= 4 covers everything in pk=1, only rk≤2
+# in pk=2 (since both NULLs tie at rank 2 and there's no rank 3 or 4).
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS LAST) as rk FROM window_topn_rank_null_t
+) WHERE rk <= 4;
+----
+1 1 1
+2 1 2
+3 1 3
+4 1 NULL
+5 2 5
+6 2 NULL
+7 2 NULL
+
+# Test R17: ASC NULLS LAST, rk <= 2 — pk=1's NULL (rk=4) drops out;
+# pk=2's NULLs (rk=2 each) are retained.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS LAST) as rk FROM window_topn_rank_null_t
+) WHERE rk <= 2;
+----
+1 1 1
+2 1 2
+5 2 5
+6 2 NULL
+7 2 NULL
+
+# Test R18: rule fires for NULLS LAST configuration
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS LAST) as rk FROM window_topn_rank_null_t
+) WHERE rk <= 2;
+----
+logical_plan
+01)Projection: window_topn_rank_null_t.id, window_topn_rank_null_t.pk, window_topn_rank_null_t.val, rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rk
+02)--Filter: rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(2)
+03)----WindowAggr: windowExpr=[[rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+04)------TableScan: window_topn_rank_null_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_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
+02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+03)----PartitionedTopKExec: fn=rank, fetch=2, partition=[pk@1], order=[val@2 ASC NULLS LAST]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Test R19: DESC NULLS LAST — pk=1: 3,2,1,NULL ranks 1,2,3,4; pk=2: 5,NULL,NULL ranks 1,2,2.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val DESC NULLS LAST) as rk FROM window_topn_rank_null_t
+) WHERE rk <= 4;
+----
+1 1 1
+2 1 2
+3 1 3
+4 1 NULL
+5 2 5
+6 2 NULL
+7 2 NULL
+
+# Test R20: ASC NULLS FIRST — pk=1: NULL,1,2,3 ranks 1,2,3,4;
+# pk=2: NULL,NULL,5 ranks 1,1,3. With rk <= 2, pk=2's NULLs are kept,
+# pk=1 keeps NULL and val=1.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS FIRST) as rk FROM window_topn_rank_null_t
+) WHERE rk <= 2;
+----
+1 1 1
+4 1 NULL
+6 2 NULL
+7 2 NULL
+
+statement ok
+DROP TABLE window_topn_rank_null_t;
+
# Reset config to default (false)
statement ok
SET datafusion.optimizer.enable_window_topn = false;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org