This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 6eaca8bfe8 perf: Extend WindowTopN to `dense_rank` (#23869)
6eaca8bfe8 is described below
commit 6eaca8bfe8421faca1413f29f71fcaf0ccaf631a
Author: Subham Singhal <[email protected]>
AuthorDate: Tue Aug 18 03:44:06 2026 +0000
perf: Extend WindowTopN to `dense_rank` (#23869)
## Which issue does this PR close?
- Related to #6899. Completes the ROW_NUMBER + RANK + DENSE_RANK trio
requested by the issue (ROW_NUMBER shipped in #21479; RANK is in flight
as PR A).
## Rationale for this change
Naive plan for `Filter(dr <= K) → BoundedWindowAggExec(DENSE_RANK) →
SortExec(pk, ob) → input` sorts the full input even though only rows at
the K distinct-smallest ob values per partition are needed. DENSE_RANK
is stable under tail-pruning, so the rewrite preserves ranks for every
retained row.
## What changes are included in this PR?
## Are these changes tested?
Yes
**h2o bench** — Q24–Q29 added; top-2 on 10M-row `large`, 3-iter avg vs
`enable_window_topn=false`:
| # | Cardinality | OFF (ms) | ON (ms) | Δ |
|-----|------------------|---------:|--------:|:-------------|
| Q24 | ~100 parts | 308.19 | 114.45 | 2.20× faster |
| Q25 | ~1K, low ties | 261.24 | 113.85 | 2.13× faster |
| Q26 | ~1K, heavy ties | 296.54 | 124.42 | 2.09× faster |
| Q27 | ~10K, low ties | 250.83 | 131.85 | 1.70× faster |
| Q28 | ~10K, heavy ties | 281.00 | 149.00 | 1.76× faster |
| Q29 | ~100K parts | 264.32 | 304.37 | 1.40× slower |
## Are there any user-facing changes?
No breaking changes. Rule fires on DENSE_RANK when
`datafusion.optimizer.enable_window_topn = true`; flag stays
`false` by default
---
.../core/tests/physical_optimizer/window_topn.rs | 103 +-
datafusion/physical-optimizer/src/window_topn.rs | 95 +-
.../physical-plan/src/sorts/partitioned_topk.rs | 68 +-
datafusion/physical-plan/src/topk/mod.rs | 1165 +++++++++++++++++++-
datafusion/sqllogictest/test_files/window_topn.slt | 453 ++++++++
5 files changed, 1828 insertions(+), 56 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/window_topn.rs
b/datafusion/core/tests/physical_optimizer/window_topn.rs
index be78b77b32..a31e7a45e5 100644
--- a/datafusion/core/tests/physical_optimizer/window_topn.rs
+++ b/datafusion/core/tests/physical_optimizer/window_topn.rs
@@ -454,8 +454,12 @@ fn build_ranking_topn_plan(
Ok(filter)
}
-/// Build a RANK plan with NO ORDER BY: every row ties at rank 1 — degenerate.
-fn build_rank_no_order_by_plan(limit_value: i64) -> Result<Arc<dyn
ExecutionPlan>> {
+/// Build a RANK / DENSE_RANK plan with NO ORDER BY: every row ties at rank 1
— degenerate.
+fn build_no_order_by_plan(
+ udwf_factory: fn() -> Arc<datafusion_expr::WindowUDF>,
+ udwf_name: &str,
+ limit_value: i64,
+) -> Result<Arc<dyn ExecutionPlan>> {
let s = schema();
let input: Arc<dyn ExecutionPlan> =
Arc::new(PlaceholderRowExec::new(Arc::clone(&s)));
@@ -469,7 +473,7 @@ fn build_rank_no_order_by_plan(limit_value: i64) ->
Result<Arc<dyn ExecutionPlan
let partition_by = vec![col("pk", &s)?];
let window_expr = Arc::new(StandardWindowExpr::new(
- create_udwf_window_expr(&rank_udwf(), &[], &s, "rank".to_string(),
false)?,
+ create_udwf_window_expr(&udwf_factory(), &[], &s,
udwf_name.to_string(), false)?,
&partition_by,
&[], // empty ORDER BY
Arc::new(WindowFrame::new_bounds(
@@ -486,7 +490,7 @@ fn build_rank_no_order_by_plan(limit_value: i64) ->
Result<Arc<dyn ExecutionPlan
true,
)?);
- let rk_col = Arc::new(Column::new("rank", 2));
+ let rk_col = Arc::new(Column::new(udwf_name, 2));
let limit_lit = lit(ScalarValue::UInt64(Some(limit_value as u64)));
let predicate = Arc::new(BinaryExpr::new(rk_col, Operator::LtEq,
limit_lit));
let filter: Arc<dyn ExecutionPlan> =
@@ -548,7 +552,7 @@ fn rank_no_order_by_no_change() -> Result<()> {
// Without ORDER BY, every row ties at rank 1 — the optimization is
// degenerate (entire input would be retained, ties storage unbounded).
// The rule must skip.
- let plan = build_rank_no_order_by_plan(3)?;
+ let plan = build_no_order_by_plan(rank_udwf, "rank", 3)?;
let before = plan_str(plan.as_ref());
let optimized = optimize(plan)?;
let after = plan_str(optimized.as_ref());
@@ -559,17 +563,98 @@ fn rank_no_order_by_no_change() -> Result<()> {
Ok(())
}
+// ----------------------------------------------------------------------
+// DENSE_RANK rule tests
+// ----------------------------------------------------------------------
+
#[test]
-fn dense_rank_no_change() -> Result<()> {
- // DENSE_RANK is not yet supported by the rule. The plan must pass
- // through unchanged.
+fn basic_dense_rank_dr_lteq_3() -> Result<()> {
let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 3,
Operator::LtEq)?;
+ let optimized = optimize(plan)?;
+ assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+ BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 },
frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0],
order=[val@1 ASC]
+ PlaceholderRowExec
+ "#);
+ Ok(())
+}
+
+#[test]
+fn dense_rank_dr_lt_4_becomes_fetch_3() -> Result<()> {
+ let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 4,
Operator::Lt)?;
+ let optimized = optimize(plan)?;
+ assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+ BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 },
frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0],
order=[val@1 ASC]
+ PlaceholderRowExec
+ "#);
+ Ok(())
+}
+
+#[test]
+fn dense_rank_flipped_3_gteq_dr() -> Result<()> {
+ let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 3,
Operator::GtEq)?;
+ let optimized = optimize(plan)?;
+ assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+ BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 },
frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0],
order=[val@1 ASC]
+ PlaceholderRowExec
+ "#);
+ Ok(())
+}
+
+#[test]
+fn dense_rank_flipped_4_gt_dr_becomes_fetch_3() -> Result<()> {
+ let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 4,
Operator::Gt)?;
+ let optimized = optimize(plan)?;
+ assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+ BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 },
frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0],
order=[val@1 ASC]
+ PlaceholderRowExec
+ "#);
+ Ok(())
+}
+
+#[test]
+fn dense_rank_no_order_by_no_change() -> Result<()> {
+ // Without ORDER BY, every row ties at dense_rank 1 — the optimization
+ // is degenerate (entire input would be retained). The rule must skip.
+ let plan = build_no_order_by_plan(dense_rank_udwf, "dense_rank", 3)?;
let before = plan_str(plan.as_ref());
let optimized = optimize(plan)?;
let after = plan_str(optimized.as_ref());
assert_eq!(
before, after,
- "DENSE_RANK is unsupported and must not be rewritten"
+ "DENSE_RANK with empty ORDER BY must not be rewritten"
);
Ok(())
}
+
+// ----------------------------------------------------------------------
+// Shared guard: `fn < 1` keeps nothing
+// ----------------------------------------------------------------------
+
+#[test]
+fn predicate_lt_1_no_change() -> Result<()> {
+ // `fn < 1` (and the flipped `1 > fn`) yields limit_n = 0. Since
+ // ROW_NUMBER / RANK / DENSE_RANK are always >= 1, the predicate keeps
+ // nothing and the rule must skip — a fetch=0 PartitionedTopK* would
+ // otherwise panic on its `k > 0` assertion at execution time.
+ type UdwfFactory = fn() -> Arc<datafusion_expr::WindowUDF>;
+ let cases: [(UdwfFactory, &str); 3] = [
+ (row_number_udwf, "row_number"),
+ (rank_udwf, "rank"),
+ (dense_rank_udwf, "dense_rank"),
+ ];
+ for (factory, name) in cases {
+ let plan = build_ranking_topn_plan(factory, name, 1, Operator::Lt)?;
+ let before = plan_str(plan.as_ref());
+ let optimized = optimize(plan)?;
+ let after = plan_str(optimized.as_ref());
+ assert_eq!(
+ before, after,
+ "`{name} < 1` (limit 0) must not be rewritten"
+ );
+ }
+ Ok(())
+}
diff --git a/datafusion/physical-optimizer/src/window_topn.rs
b/datafusion/physical-optimizer/src/window_topn.rs
index aa46a5a9a7..758ad66de3 100644
--- a/datafusion/physical-optimizer/src/window_topn.rs
+++ b/datafusion/physical-optimizer/src/window_topn.rs
@@ -26,7 +26,7 @@
//! ) WHERE rn <= K;
//! ```
//!
-//! or with `RANK()` in place of `ROW_NUMBER()`:
+//! or with `RANK()` / `DENSE_RANK()` in place of `ROW_NUMBER()`:
//!
//! ```sql
//! SELECT * FROM (
@@ -40,8 +40,8 @@
//! `FilterExec` and inserting `PartitionedTopKExec` under the window.
//!
//! The appropriate [`WindowFnKind`] is forwarded to `PartitionedTopKExec`.
-//! RANK requires a non-empty `ORDER BY` clause (otherwise all rows tie at
-//! rank 1 and the optimization is degenerate).
+//! `RANK` and `DENSE_RANK` require a non-empty `ORDER BY` clause (otherwise
+//! all rows tie at rank 1 and the optimization is degenerate).
//!
//! See [`PartitionedTopKExec`] for details on the replacement operator.
//!
@@ -69,9 +69,9 @@ use datafusion_physical_plan::sorts::partitioned_topk::{
};
use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
-/// Physical optimizer rule that converts per-partition `ROW_NUMBER` and
-/// `RANK` top-K queries into a more efficient plan using
-/// [`PartitionedTopKExec`].
+/// Physical optimizer rule that converts per-partition `ROW_NUMBER`,
+/// `RANK`, and `DENSE_RANK` top-K queries into a more efficient plan
+/// using [`PartitionedTopKExec`].
///
/// # Pattern Detected
///
@@ -86,12 +86,14 @@ use
datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
/// ```text
/// [optional ProjectionExec]
/// BoundedWindowAggExec(<ranking fn> PARTITION BY ... ORDER BY ...)
-/// PartitionedTopKExec(fn=<row_number|rank>, partition_keys, order_keys,
fetch=K)
+/// PartitionedTopKExec(fn=<row_number|rank|dense_rank>, partition_keys,
order_keys, fetch=K)
/// ```
///
/// The `FilterExec` is removed entirely. The child of `BoundedWindowAggExec`
is now
-/// `PartitionedTopKExec`, which maintains a per-partition top-K heap (and,
-/// for `RANK`, a sibling ties `Vec`) instead of sorting the whole dataset.
+/// `PartitionedTopKExec`, which maintains per-partition top-K state (a
+/// heap for `ROW_NUMBER`, a heap plus boundary ties for `RANK`, a
+/// K-bounded distinct-ob map for `DENSE_RANK`) instead of sorting the
+/// whole dataset.
///
/// # Supported Predicates
///
@@ -105,12 +107,19 @@ use
datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
/// All of the following must be true:
/// - Config flag `enable_window_topn` is `true`
/// - The plan matches `FilterExec → [ProjectionExec] → BoundedWindowAggExec`
-/// - The window function is `ROW_NUMBER` or `RANK` (not `DENSE_RANK`)
+/// - The window function is `ROW_NUMBER`, `RANK`, or `DENSE_RANK`
+/// - Every window expression in the `BoundedWindowAggExec` is `ROW_NUMBER`,
+/// `RANK`, or `DENSE_RANK` over the same `PARTITION BY` / `ORDER BY`. A
+/// sibling that reads pruned rows (e.g. `LEAD`, or an aggregate whose
+/// frame is not strictly backward-looking) would be computed over the
+/// pruned input and give wrong results.
/// - The window function has a `PARTITION BY` clause (global top-K is
/// already handled by `SortExec` with `fetch`)
-/// - For `RANK`: a non-empty `ORDER BY` clause (otherwise all rows tie
-/// at rank 1 — the optimization is useless and the boundary-tie storage
-/// would be unbounded)
+/// - At least one `ORDER BY` key survives past the `PARTITION BY` prefix
+/// (so the operator has a non-empty ORDER BY). This rejects both a
+/// missing `ORDER BY` and one fully covered by the partition prefix
+/// such as `PARTITION BY pk ORDER BY pk`; for `RANK` / `DENSE_RANK`
+/// such orderings also make every row tie at rank 1 (degenerate).
/// - The filter predicate compares the window output column to an integer
/// literal using `<=`, `<`, `>=`, or `>`
///
@@ -142,10 +151,10 @@ impl WindowTopN {
let (col_idx, limit_n) = extract_window_limit(filter.predicate())?;
// A predicate such as `rn < 1` (or the flipped `1 > rn`) yields a
fetch of
- // 0. `ROW_NUMBER`/`RANK` are always >= 1, so no row can satisfy it
and the
- // correct result is empty. `PartitionedTopKExec` requires `k > 0` and
would
- // panic on `k = 0`, so bail out here and let the regular `FilterExec`
produce
- // the (empty) result instead of rewriting.
+ // 0. `ROW_NUMBER`/`RANK`/`DENSE_RANK` are always >= 1, so no row can
satisfy
+ // it and the correct result is empty. `PartitionedTopKExec` requires
`k > 0`
+ // and would panic on `k = 0`, so bail out here and let the regular
+ // `FilterExec` produce the (empty) result instead of rewriting.
if limit_n == 0 {
return None;
}
@@ -167,6 +176,28 @@ impl WindowTopN {
}
let fn_kind = supported_window_fn(&window_exprs[window_expr_idx])?;
+ // Tail-pruning drops the rows that rank after the retained top-K,
+ // and every window expression in this `BoundedWindowAggExec` is
+ // then evaluated over the *pruned* input. The rewrite is only valid
+ // if each expression's value for a *retained* row is unaffected by
+ // the dropped rows. ROW_NUMBER / RANK / DENSE_RANK over the same
+ // PARTITION BY / ORDER BY satisfy this — each depends only on rows
+ // at or before the current row in that order, all of which are
+ // retained. A sibling like `LEAD(x)` reads following (pruned) rows,
+ // so at the retained boundary it would resolve to a dropped row and
+ // give a wrong result. Bail out unless every window expression is a
+ // supported ranking function sharing the matched expression's
+ // partition/order keys.
+ let matched_expr = &window_exprs[window_expr_idx];
+ let all_prune_safe = window_exprs.iter().all(|e| {
+ supported_window_fn(e).is_some()
+ && e.partition_by() == matched_expr.partition_by()
+ && e.order_by() == matched_expr.order_by()
+ });
+ if !all_prune_safe {
+ return None;
+ }
+
// Step 5: Validate PARTITION BY / ORDER BY and collect sort keys from
the window expr
let partition_by = window_exprs[window_expr_idx].partition_by();
let partition_prefix_len = partition_by.len();
@@ -177,21 +208,31 @@ impl WindowTopN {
return None;
}
- // For RANK: an empty ORDER BY makes every row tie at rank 1 —
- // the optimization is degenerate (we'd retain the entire input)
- // and tie storage would be unbounded.
- let order_by = window_exprs[window_expr_idx].order_by();
- if matches!(fn_kind, WindowFnKind::Rank) && order_by.is_empty() {
- return None;
- }
-
// Step 6: Build PartitionedTopKExec from the window's partition/order
keys
+ let order_by = window_exprs[window_expr_idx].order_by();
let expr_iterator = partition_by
.iter()
.map(|e| PhysicalSortExpr::new_default(Arc::clone(e)))
.chain(order_by.iter().cloned());
let expr = LexOrdering::new(expr_iterator)?;
+ // `PartitionedTopKExec` derives its ORDER BY keys from the ordering
+ // *beyond* the partition prefix (`expr[partition_prefix_len..]`).
+ // That slice is empty in two cases:
+ // * no ORDER BY at all (e.g. `ROW_NUMBER() OVER (PARTITION BY pk)`);
+ // * ORDER BY keys fully covered by the partition prefix (e.g.
+ // `DENSE_RANK() OVER (PARTITION BY pk ORDER BY pk)`, whose
+ // deduplicated ordering is just `[pk]`).
+ // With zero order keys the operator panics on execution (it
+ // requires at least one), and for RANK / DENSE_RANK every row
+ // would tie at rank 1 (a degenerate, unbounded retained set).
+ // `order_by()` alone does not catch the second case — it reports
+ // `[pk]` even though no order key survives past the partition
+ // prefix — so guard on the effective order-key count instead.
+ if expr.len() <= partition_prefix_len {
+ return None;
+ }
+
let partitioned_topk = PartitionedTopKExec::try_new(
Arc::clone(window_exec_typed.input()),
expr,
@@ -325,7 +366,8 @@ fn scalar_to_usize(value: &ScalarValue) -> Option<usize> {
/// the UDF name. Returns:
/// - `Some(WindowFnKind::RowNumber)` for `"row_number"`
/// - `Some(WindowFnKind::Rank)` for `"rank"`
-/// - `None` for everything else (e.g. `dense_rank`)
+/// - `Some(WindowFnKind::DenseRank)` for `"dense_rank"`
+/// - `None` for everything else
fn supported_window_fn(
expr: &Arc<dyn datafusion_physical_expr::window::WindowExpr>,
) -> Option<WindowFnKind> {
@@ -335,6 +377,7 @@ fn supported_window_fn(
match udf.fun().name() {
"row_number" => Some(WindowFnKind::RowNumber),
"rank" => Some(WindowFnKind::Rank),
+ "dense_rank" => Some(WindowFnKind::DenseRank),
_ => None,
}
}
diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs
b/datafusion/physical-plan/src/sorts/partitioned_topk.rs
index 41ccfab683..cd8bd340d2 100644
--- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs
+++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs
@@ -24,9 +24,9 @@
//! ```
//!
//! 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`],
+//! per-partition top-K implementation — one variant each for `ROW_NUMBER`,
+//! `RANK`, and `DENSE_RANK` — all of which keep per-partition state 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)`.
@@ -47,7 +47,9 @@ use futures::TryStreamExt;
use crate::execution_plan::{Boundedness, EmissionType};
use crate::metrics::ExecutionPlanMetricsSet;
-use crate::topk::{PartitionedTopK, PartitionedTopKRank, build_sort_fields};
+use crate::topk::{
+ PartitionedTopK, PartitionedTopKDenseRank, PartitionedTopKRank,
build_sort_fields,
+};
use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions};
use crate::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
ExecutionPlanProperties,
@@ -61,12 +63,19 @@ use crate::{
/// - [`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).
+/// - [`DenseRank`](Self::DenseRank): every row whose ORDER BY value is
+/// among the K distinct-smallest ORDER BY values in the partition
+/// (DENSE_RANK semantics — total kept rows is unbounded in
+/// rows-per-distinct-value).
#[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,
+ /// `DENSE_RANK()` — keep every row whose ob value is among the K
+ /// distinct-smallest ob values seen in the partition.
+ DenseRank,
}
/// Per-partition Top-K operator for window function queries.
@@ -110,9 +119,10 @@ pub enum WindowFnKind {
/// ```
///
/// 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
+/// and delegates to a per-partition top-K implementation (`PartitionedTopK`
+/// for `ROW_NUMBER`, `PartitionedTopKRank` for `RANK`, and
+/// `PartitionedTopKDenseRank` for `DENSE_RANK`), each maintaining
+/// per-partition state 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
@@ -164,11 +174,12 @@ pub enum WindowFnKind {
///
/// # Limitations
///
-/// - 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`.
+/// - Only activated when the window function is `ROW_NUMBER`, `RANK`, or
+/// `DENSE_RANK` with a `PARTITION BY` clause. `RANK` and `DENSE_RANK`
+/// additionally require a non-empty `ORDER BY` (with an empty `ORDER BY`
+/// every row ties at rank 1 and the 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
@@ -212,7 +223,8 @@ impl PartitionedTopKExec {
/// 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`]).
+ /// ([`WindowFnKind::RowNumber`], [`WindowFnKind::Rank`], or
+ /// [`WindowFnKind::DenseRank`]).
///
/// # Example
///
@@ -297,6 +309,7 @@ impl DisplayAs for PartitionedTopKExec {
let fn_label = match self.fn_kind {
WindowFnKind::RowNumber => "row_number",
WindowFnKind::Rank => "rank",
+ WindowFnKind::DenseRank => "dense_rank",
};
match t {
DisplayFormatType::Default | DisplayFormatType::Verbose => {
@@ -453,9 +466,9 @@ impl ExecutionPlan for PartitionedTopKExec {
}
/// 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)`.
+/// ([`PartitionedTopK`] for `ROW_NUMBER`, [`PartitionedTopKRank`] for
+/// `RANK`, or [`PartitionedTopKDenseRank`] for `DENSE_RANK`), then emit
+/// results ordered by `(partition_keys, order_keys)`.
///
/// # Phases
///
@@ -465,10 +478,11 @@ impl ExecutionPlan for PartitionedTopKExec {
/// `TopKMetrics` are shared across all distinct partition keys for
/// this operator instance.
///
-/// 2. **Emission** — `emit` drains all per-partition heaps in sorted
+/// 2. **Emission** — `emit` drains all per-partition state 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.
+/// each partition's heap rows. For `DENSE_RANK`, rows are emitted
+/// from a K-bounded map of distinct ob keys, sorted ascending.
///
/// # Cost
///
@@ -526,5 +540,23 @@ async fn do_partitioned_topk(
drop(input);
state.emit()
}
+ WindowFnKind::DenseRank => {
+ let mut state = PartitionedTopKDenseRank::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()
+ }
}
}
diff --git a/datafusion/physical-plan/src/topk/mod.rs
b/datafusion/physical-plan/src/topk/mod.rs
index 1e3efff36b..1d02397b92 100644
--- a/datafusion/physical-plan/src/topk/mod.rs
+++ b/datafusion/physical-plan/src/topk/mod.rs
@@ -1811,10 +1811,454 @@ impl PartitionedTopKRank {
}
}
+/// A run of rows from a single source [`RecordBatch`] sharing one
+/// distinct ORDER BY value. Materialized at emit time via
+/// [`take_record_batch`].
+///
+/// The batch is referenced by `batch_id` rather than held directly, so the
+/// operator-scoped [`RecordBatchStore`] can charge each distinct source
+/// batch once however many entries reference it. Holding a batch per entry
+/// and charging its bytes per entry would inflate the reservation by a
+/// factor of (partitions × K), since a single batch contributes an entry to
+/// every partition and ob group it touches.
+#[derive(Debug)]
+struct GroupEntry {
+ /// Indices into the batch identified by `batch_id`. Always non-empty
+ /// by construction.
+ row_indices: Vec<u32>,
+ /// Key into `PartitionedTopKDenseRank::store`.
+ batch_id: u32,
+}
+
+/// Per-partition state for `DENSE_RANK()` semantics.
+///
+/// A `HashMap<Vec<u8>, Vec<GroupEntry>>` keyed by the row-encoded ORDER
+/// BY bytes, capped at `k` distinct keys. Each key's `Vec<GroupEntry>`
+/// holds every row seen at that ob value, one entry per contributing
+/// source `RecordBatch`.
+#[derive(Default)]
+struct DenseRankPartitionState {
+ groups: HashMap<Vec<u8>, Vec<GroupEntry>>,
+ /// The same keys as `groups`, in a max-heap: the admission boundary
+ /// (the largest tracked ob value) is an O(1) `peek()` check, and
+ /// admission / removal are O(log K).
+ ///
+ /// INVARIANT: `keys` and `groups.keys()` hold the same set. Every
+ /// insertion into / removal from `groups` must mirror into `keys`.
+ keys: BinaryHeap<Vec<u8>>,
+}
+
+impl DenseRankPartitionState {
+ fn size(&self) -> usize {
+ let table_overhead = self.groups.capacity()
+ * (size_of::<Vec<u8>>() + size_of::<Vec<GroupEntry>>());
+ let contents: usize = self
+ .groups
+ .iter()
+ .map(|(key, entries)| {
+ key.capacity()
+ + entries.capacity() * size_of::<GroupEntry>()
+ + entries
+ .iter()
+ .map(|e| e.row_indices.capacity() * size_of::<u32>())
+ .sum::<usize>()
+ })
+ .sum();
+ // `keys` duplicates every key's bytes; charge for them plus the
+ // heap's backing Vec (one `Vec<u8>` slot per reserved element).
+ let keys_overhead: usize = self.keys.capacity() * size_of::<Vec<u8>>()
+ + self.keys.iter().map(|k| k.capacity()).sum::<usize>();
+ table_overhead + contents + keys_overhead
+ }
+}
+
+/// Sibling to [`PartitionedTopK`] / [`PartitionedTopKRank`] implementing
+/// `DENSE_RANK()` semantics.
+///
+/// Per partition, retains every row whose ORDER BY value is among the K
+/// distinct-smallest ob values seen for that partition. The total row
+/// count kept per partition is unbounded in `rows_per_distinct_value`
+/// (unlike `RANK`, which is bounded above by K + boundary ties).
+///
+/// Like [`PartitionedTopK`], the [`RowConverter`], [`MemoryReservation`],
+/// scratch [`Rows`] buffer, and [`TopKMetrics`] are shared across all
+/// partitions for this operator instance. So is the
+/// [`RecordBatchStore`]: retained rows are held as `(batch_id, indices)`
+/// so each source batch is charged once for the whole operator, however
+/// many partitions and ob groups reference it.
+///
+/// # Algorithm (per batch)
+///
+/// Evaluate + encode partition-by and order-by columns once, then group
+/// the batch's row indices by partition key. For each partition, bucket
+/// that partition's rows by distinct ob value (a within-call
+/// accumulation), then merge each bucket into the partition state. Every
+/// bucket is built from the current batch's rows, so each `GroupEntry` is
+/// pinned to the batch its `row_indices` point into.
+///
+/// For each partition, for each distinct `ob_key` run in this batch:
+/// - `ob_key` already in `state.groups` → push this batch's run as a
+/// new `GroupEntry` (one entry per contributing batch).
+/// - `ob_key` new, `state.groups.len() < k` → insert the run as a new
+/// group.
+/// - `ob_key` new, `state.groups.len() == k` → the largest tracked ob
+/// value is the admission boundary, read from the `state.keys` max-heap
+/// in O(1):
+/// - `ob_key < max` → remove the max key (evict the entire max-key
+/// group — up to many rows) and insert the run. The evicted group's
+/// row count is added to the `row_replacements` metric.
+/// - `ob_key >= max` → drop the whole run; no map mutation.
+pub(crate) struct PartitionedTopKDenseRank {
+ 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).
+ partition_scratch_rows: Rows,
+ /// One state per distinct partition key seen so far. Keyed by the
+ /// row-encoded PARTITION BY bytes (byte-comparable encoding, so the
+ /// `Vec<u8>` hashes, compares, and sorts identically to an
+ /// `OwnedRow`) which lets `insert_batch` look partitions up with
+ /// `entry_ref` — allocating a key only on first sight of a partition
+ /// rather than once per row.
+ states: HashMap<Vec<u8>, DenseRankPartitionState>,
+ /// Scratch map reused across `insert_batch` calls to group a batch's
+ /// row indices by partition key. Drained (not reallocated) each batch
+ /// so its backing table is allocated once, not per batch.
+ partition_groups: HashMap<Vec<u8>, Vec<u32>>,
+ /// Scratch map reused across partitions within a batch to bucket a
+ /// partition's rows by distinct ORDER BY value. Drained (not
+ /// reallocated) per partition so its backing table is allocated once,
+ /// not once per distinct partition key.
+ ob_runs: HashMap<Vec<u8>, Vec<u32>>,
+ /// Source batches referenced by the retained `GroupEntry`s, held once
+ /// for the whole operator with a use count per batch. This is what
+ /// keeps the reservation proportional to the batches actually pinned
+ /// rather than to the number of entries pointing at them.
+ store: RecordBatchStore,
+ k: usize,
+ batch_size: usize,
+}
+
+impl PartitionedTopKDenseRank {
+ #[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, "PartitionedTopKDenseRank requires k > 0");
+ let reservation =
+
MemoryConsumer::new(format!("PartitionedTopKDenseRank[{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(),
+ partition_groups: HashMap::new(),
+ ob_runs: HashMap::new(),
+ store: RecordBatchStore::new(),
+ k,
+ batch_size,
+ })
+ }
+
+ /// Encode PARTITION BY and ORDER BY columns once, demultiplex the
+ /// batch's rows by partition key, then per partition bucket the rows
+ /// by distinct ob value and merge each bucket into the partition
+ /// state as one [`GroupEntry`].
+ 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(());
+ }
+
+ // Register this batch with the store up front. `uses` is bumped
+ // once per `GroupEntry` created below but lives on this local
+ // entry, so the store sees a single insert per batch rather than
+ // one per group — and `insert` drops batches that retained
+ // nothing without ever charging for them.
+ let mut batch_entry = self.store.register(batch.clone());
+ let batch_id = batch_entry.id;
+
+ // 1. Encode partition columns.
+ 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)?;
+
+ // 2. Group this batch's row indices by partition key.
+ // `partition_groups` is a reused scratch map: taken out here
+ // and drained below, so its backing table is allocated once
+ // for the operator, not once per batch. `entry_ref` owns the
+ // key only on Vacant, so it allocates one `Vec<u8>` per
+ // distinct partition rather than one per row.
+ let mut groups = std::mem::take(&mut self.partition_groups);
+ groups.clear();
+ {
+ let pk_rows = &self.partition_scratch_rows;
+ for i in 0..num_rows {
+ groups
+ .entry_ref(pk_rows.row(i).as_ref())
+ .or_default()
+ .push(i as u32);
+ }
+ }
+
+ // 3. Evaluate ORDER BY columns 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)?;
+
+ let k = self.k;
+ let mut replacements: usize = 0;
+
+ // 4. Per-partition: bucket this batch's rows by distinct ob value
+ // (within-call accumulation), then merge each bucket into the
+ // partition state as a single `GroupEntry`.
+ for (pk, indices) in groups.drain() {
+ let state = self
+ .states
+ .entry(pk)
+ .or_insert_with(DenseRankPartitionState::default);
+
+ // Bucket by ob key. `ob_runs` is a reused scratch map (taken
+ // out and drained below) so its backing table is allocated
+ // once, not once per distinct partition key. `entry_ref` owns
+ // the key only on Vacant, so repeated rows of the same ob
+ // value don't re-allocate.
+ let mut runs = std::mem::take(&mut self.ob_runs);
+ runs.clear();
+ for &orig_idx in &indices {
+ let ob_row = self.scratch_rows.row(orig_idx as usize);
+ runs.entry_ref(ob_row.as_ref()).or_default().push(orig_idx);
+ }
+
+ for (ob_key, run_indices) in runs.drain() {
+ // Case A: ob already tracked — push this batch's run as a
+ // new `GroupEntry` (one entry per contributing batch).
+ if let Some(entries) = state.groups.get_mut(&ob_key) {
+ batch_entry.uses += 1;
+ entries.push(GroupEntry {
+ row_indices: run_indices,
+ batch_id,
+ });
+ continue;
+ }
+
+ // Case B: new ob, room available.
+ if state.groups.len() < k {
+ batch_entry.uses += 1;
+ state.keys.push(ob_key.clone());
+ state.groups.insert(
+ ob_key,
+ vec![GroupEntry {
+ row_indices: run_indices,
+ batch_id,
+ }],
+ );
+ continue;
+ }
+
+ // Case C: new ob, at K distinct keys. The largest tracked
+ // ob value is the admission boundary.
+ let max_key = state.keys.peek().expect("state.groups has k >=
1 keys");
+ if ob_key.as_slice() < max_key.as_slice() {
+ // Evict the entire max-key group, from both the map
+ // and its ordered mirror.
+ let evicted_key = state.keys.pop().expect("max key
present");
+ let evicted = state
+ .groups
+ .remove(&evicted_key)
+ .expect("keys mirrors groups");
+ for e in &evicted {
+ replacements += e.row_indices.len();
+ if e.batch_id == batch_id {
+ // Admitted earlier in this same call, so the
+ // store has not seen `batch_entry` yet — drop
+ // the pending use rather than calling `unuse`,
+ // which panics on an unregistered id.
+ batch_entry.uses -= 1;
+ } else {
+ self.store.unuse(e.batch_id);
+ }
+ }
+ batch_entry.uses += 1;
+ state.keys.push(ob_key.clone());
+ state.groups.insert(
+ ob_key,
+ vec![GroupEntry {
+ row_indices: run_indices,
+ batch_id,
+ }],
+ );
+ }
+ // else: ob >= max — drop the whole run.
+ }
+
+ // Return the drained scratch map (capacity retained) for the
+ // next partition to reuse.
+ self.ob_runs = runs;
+ }
+
+ // Return the drained scratch map (capacity retained) for the next
+ // batch to reuse.
+ self.partition_groups = groups;
+
+ // Charges `batch` once if any group retained rows from it.
+ self.store.insert(batch_entry);
+
+ if replacements > 0 {
+ self.metrics.row_replacements.add(replacements);
+ }
+ self.reservation.try_resize(self.size())?;
+ Ok(())
+ }
+
+ /// Drain all per-partition maps in partition-key order and return
+ /// the rows as a stream of coalesced [`RecordBatch`]es ordered by
+ /// `(partition_keys, order_keys)`. Within a partition the distinct
+ /// ob keys are sorted (byte-comparable encoding == sort order) so
+ /// emitted rows are in ob-sorted order.
+ 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,
+ partition_groups: _,
+ ob_runs: _,
+ store,
+ k: _,
+ batch_size,
+ } = self;
+ let _timer = metrics.baseline.elapsed_compute().timer();
+
+ let mut sorted_pks: Vec<Vec<u8>> = states.keys().cloned().collect();
+ sorted_pks.sort();
+
+ let mut coalescer = BatchCoalescer::new(Arc::clone(&schema),
batch_size);
+
+ for pk in sorted_pks {
+ let DenseRankPartitionState { groups, keys: _ } =
+ states.remove(&pk).expect("key from states.keys()");
+ // Sort the <= K distinct ob keys so rows emit ascending
+ // (byte-comparable encoding == sort order).
+ let mut sorted_obs: Vec<(Vec<u8>, Vec<GroupEntry>)> =
+ groups.into_iter().collect();
+ sorted_obs.sort_by(|a, b| a.0.cmp(&b.0));
+ for (_ob, entries) in sorted_obs {
+ for entry in entries {
+ let batch = &store
+ .get(entry.batch_id)
+ .expect("retained batch_id present in store")
+ .batch;
+ let indices = UInt32Array::from(entry.row_indices);
+ let sub = take_record_batch(batch, &indices)?;
+ (&sub).record_output(&metrics.baseline);
+ coalescer.push_batch(sub)?;
+ }
+ }
+ }
+ 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 {
+ // Per partition: the state itself plus the encoded partition key
+ // owned by the map. The key bytes are a heap allocation the table
+ // slot doesn't cover, and with wide or numerous partition keys
+ // they dominate the fixed-size slots.
+ let states_contents: usize = self
+ .states
+ .iter()
+ .map(|(pk, state)| pk.capacity() + state.size())
+ .sum();
+ // `partition_groups` and `ob_runs` are drained, not dropped, so
+ // their backing tables outlive every `insert_batch` call. Both are
+ // empty by the time `size()` runs (drained above), so only the
+ // retained capacity is charged.
+ let scratch_tables = self.partition_groups.capacity()
+ * (size_of::<Vec<u8>>() + size_of::<Vec<u32>>())
+ + self.ob_runs.capacity() * (size_of::<Vec<u8>>() +
size_of::<Vec<u32>>());
+ size_of::<Self>()
+ + self.row_converter.size()
+ + self.partition_converter.size()
+ + self.scratch_rows.size()
+ + self.partition_scratch_rows.size()
+ + states_contents
+ + self.states.capacity()
+ * (size_of::<Vec<u8>>() + size_of::<DenseRankPartitionState>())
+ + scratch_tables
+ + self.store.size()
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
- use arrow::array::{BooleanArray, Float64Array, Int32Array};
+ use arrow::array::{BooleanArray, Float64Array, Int32Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use arrow_schema::SortOptions;
use datafusion_common::assert_batches_eq;
@@ -2537,7 +2981,7 @@ mod tests {
}
/// Multiple distinct partition keys interleaved within a single
- /// input batch — the per-batch demux, per-partition heap eviction,
+ /// input batch — grouping rows by partition key, per-partition heap
eviction,
/// and partition-key-ordered emit must all behave correctly.
#[tokio::test]
async fn test_partitioned_topk_multi_partition_within_batch() ->
Result<()> {
@@ -2831,7 +3275,7 @@ mod tests {
}
/// Multiple distinct partition keys interleaved within a single
- /// input batch — the per-batch demux, per-partition heap eviction,
+ /// input batch — grouping rows by partition key, 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]
@@ -3166,4 +3610,719 @@ mod tests {
);
Ok(())
}
+
+ // ====================================================================
+ // PartitionedTopKDenseRank operator tests
+ //
+ // These mirror the RANK tests plus DENSE_RANK-specific cases: rows
+ // sharing an ob key coalesce into one `GroupEntry`, unbounded
+ // rows-per-distinct-key, and eviction removes the entire max group
+ // when a strictly-smaller distinct ob arrives.
+ // ====================================================================
+
+ /// Builds a `(pk Int32, val Int32)` schema and a
+ /// `PartitionedTopKDenseRank` keyed on `pk ASC` (partition) and
+ /// `val ASC` (ORDER BY).
+ fn build_partitioned_topk_dense_rank(
+ k: usize,
+ ) -> Result<(Arc<Schema>, PartitionedTopKDenseRank)> {
+ build_partitioned_topk_dense_rank_with_opts(k, SortOptions::default(),
false)
+ }
+
+ fn build_partitioned_topk_dense_rank_with_opts(
+ k: usize,
+ val_sort_options: SortOptions,
+ val_nullable: bool,
+ ) -> Result<(Arc<Schema>, PartitionedTopKDenseRank)> {
+ 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 = PartitionedTopKDenseRank::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))
+ }
+
+ /// Single-batch DENSE_RANK top-2 across multiple partitions with
+ /// distinct ob values only — should behave identically to a
+ /// ROW_NUMBER top-2. Exercises per-partition grouping + emit order.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_multi_partition_within_batch()
-> Result<()>
+ {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
+
+ // pk=1 vals: 10, 5, 8 → distinct-top-2 ASC = {5, 8}
+ // pk=2 vals: 20, 15 → distinct-top-2 ASC = {15, 20}
+ // pk=3 vals: 7 → distinct-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(())
+ }
+
+ /// DENSE_RANK-specific: heavy ties within a batch. All rows at each
+ /// distinct ob value must be kept — within-call bucketing groups them
+ /// into one `GroupEntry` per distinct ob.
+ ///
+ /// vals per partition (sorted logically):
+ /// pk=1: 1, 1, 1, 2, 2, 3, 3, 3, 4
+ /// distinct-top-2 = {1, 2} → all 5 rows at those values retained.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_heavy_ties_coalesced() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
+
+ let batch = pk_val_batch(
+ &schema,
+ vec![1, 1, 1, 1, 1, 1, 1, 1, 1],
+ vec![1, 3, 1, 2, 3, 1, 2, 3, 4],
+ )?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 1 |",
+ "| 1 | 1 |",
+ "| 1 | 1 |",
+ "| 1 | 2 |",
+ "| 1 | 2 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// Rows tied at the same ob across two source batches must both
+ /// land under the same map key as separate `GroupEntry`s — one per
+ /// source batch — but emit as a single contiguous run.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_cross_batch_same_key() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
+
+ // Batch 1: pk=1 with ob values {5, 5, 8}. groups after: {5→[..],
8→[..]}.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1, 1], vec![5, 5,
8])?)?;
+
+ // Batch 2: pk=1 with more 5s and an 8, plus a 20 that's dropped.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1, 1], vec![5, 8,
20])?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 5 |",
+ "| 1 | 5 |",
+ "| 1 | 5 |",
+ "| 1 | 8 |",
+ "| 1 | 8 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// Refactor guard: the full RANK-style path in one run — multi-partition
+ /// per-batch grouping, within-batch bucketing of scattered same-ob rows,
+ /// cross-batch append to an existing group, cross-batch new-key insert,
+ /// and cross-batch eviction of a whole max group. Every `GroupEntry` is
+ /// built from its own source batch (no cross-batch coalescing), so the
+ /// retained rows must be exactly the K=2 smallest distinct ob values
+ /// per partition with all their rows, regardless of arrival order.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_multi_batch_multi_partition() ->
Result<()>
+ {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
+
+ // Batch 1 interleaves pk=1 and pk=2, with same-ob rows scattered:
+ // pk=1 vals: 10, 20, 10, 20, 10 → {10:[×3], 20:[×2]}
+ // pk=2 vals: 100, 100 → {100:[×2]}
+ state.insert_batch(&pk_val_batch(
+ &schema,
+ vec![1, 2, 1, 1, 2, 1, 1],
+ vec![10, 100, 20, 20, 100, 10, 10],
+ )?)?;
+
+ // Batch 2:
+ // pk=1 vals: 20, 5, 10 → append a 20, insert 5 (evicts the whole
+ // 20 group), append a 10 → retained distinct {5, 10}.
+ // pk=2 vals: 50 → insert 5th... new key, room → {50, 100}.
+ state.insert_batch(&pk_val_batch(
+ &schema,
+ vec![1, 2, 1, 1],
+ vec![20, 50, 5, 10],
+ )?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ // pk=1: val=5 (×1 from batch 2), val=10 (×3 batch 1 + ×1 batch 2 =
×4).
+ // All 20s dropped (evicted). pk=2: val=50 (×1), val=100 (×2).
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 5 |",
+ "| 1 | 10 |",
+ "| 1 | 10 |",
+ "| 1 | 10 |",
+ "| 1 | 10 |",
+ "| 2 | 50 |",
+ "| 2 | 100 |",
+ "| 2 | 100 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// DENSE_RANK-specific: eviction removes the entire max group when
+ /// a strictly-smaller distinct ob arrives. Multiple rows at the
+ /// evicted key all disappear.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_max_group_eviction() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
+
+ // Batch 1: pk=1 with {10, 10, 20, 20}. groups={10→[..], 20→[..]}, at
K.
+ state.insert_batch(&pk_val_batch(
+ &schema,
+ vec![1, 1, 1, 1],
+ vec![10, 10, 20, 20],
+ )?)?;
+
+ // Batch 2: pk=1 with 5 — strictly smaller than max=20, evict entire
+ // 20 group; now groups={10, 5}. Then a 30 comes in and is dropped.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![5, 30])?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 5 |",
+ "| 1 | 10 |",
+ "| 1 | 10 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// Empty input must produce an empty output stream, not panic.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_empty_input() -> Result<()> {
+ let (_schema, state) = build_partitioned_topk_dense_rank(3)?;
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert!(results.is_empty(), "empty input → empty output");
+ Ok(())
+ }
+
+ /// `fetch = 1` retains only the smallest distinct ob per partition,
+ /// with all rows at that value kept.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_fetch_one() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(1)?;
+
+ // pk=1 vals: 5, 3, 5, 3, 7 → distinct-top-1 = {3} → both 3s kept.
+ // pk=2 vals: 9, 4 → distinct-top-1 = {4} → single 4.
+ let batch = pk_val_batch(
+ &schema,
+ vec![1, 1, 1, 2, 1, 2, 1],
+ vec![5, 3, 5, 9, 3, 4, 7],
+ )?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 3 |",
+ "| 1 | 3 |",
+ "| 2 | 4 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// `K > distinct_ob_count` — nothing should be dropped.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_k_exceeds_distinct() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(10)?;
+
+ // Only 3 distinct ob values under pk=1; all rows must be retained.
+ let batch = pk_val_batch(&schema, vec![1, 1, 1, 1], vec![5, 3, 3, 7])?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 3 |",
+ "| 1 | 3 |",
+ "| 1 | 5 |",
+ "| 1 | 7 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// `ORDER BY val DESC` — the row-encoded key ordering must reflect
+ /// the direction so the "distinct-K best" set is the K *largest*
+ /// distinct ob values.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_desc_ordering() -> Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank_with_opts(
+ 2,
+ SortOptions {
+ descending: true,
+ nulls_first: false,
+ },
+ false,
+ )?;
+
+ // pk=1 vals: 10, 5, 8, 12, 10 → distinct-top-2 DESC = {12, 10}
+ // → keep both 10s and 12.
+ let batch = pk_val_batch(&schema, vec![1, 1, 1, 1, 1], vec![10, 5, 8,
12, 10])?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 12 |",
+ "| 1 | 10 |",
+ "| 1 | 10 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// Cross-partition eviction independence — Case-C eviction in one
+ /// partition must not affect another partition's state.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_partition_independence() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
+
+ // Batch 1: pk=1 fills {10, 20}; pk=2 fills {30, 40}.
+ state.insert_batch(&pk_val_batch(
+ &schema,
+ vec![1, 1, 2, 2],
+ vec![10, 20, 30, 40],
+ )?)?;
+
+ // Batch 2: pk=1 sees 5 (evicts 20). pk=2 sees 25 (evicts 40).
+ // Each partition's Case-C branch is independent.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 2], vec![5, 25])?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 5 |",
+ "| 1 | 10 |",
+ "| 2 | 25 |",
+ "| 2 | 30 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// NULL sort values exercise the shared encoder's null-ordering
+ /// through the row-encoded key byte order. With `ASC NULLS
+ /// LAST`, a NULL is the *largest* distinct ob, so a partition with
+ /// >= K non-NULL distinct values evicts its NULLs, while a partition
+ /// whose only distinct value is NULL still emits it.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_nulls_last_ordering() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank_with_opts(
+ 2,
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ true,
+ )?;
+
+ // pk=1 vals: NULL, 10, 20, NULL → distinct-top-2 NULLS LAST = {10, 20}
+ // pk=2 vals: NULL → distinct-top-2 = {NULL}
+ // pk=3 vals: 3, 3 → distinct-top-2 = {3}
+ let batch = nullable_pk_val_batch(
+ &schema,
+ vec![1, 1, 1, 1, 2, 3, 3],
+ vec![None, Some(10), Some(20), None, None, Some(3), Some(3)],
+ )?;
+ state.insert_batch(&batch)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 10 |",
+ "| 1 | 20 |",
+ "| 2 | |",
+ "| 3 | 3 |",
+ "| 3 | 3 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// `ASC NULLS FIRST` sorts NULLs *before* every non-NULL value, so a
+ /// NULL is the smallest distinct ob and is kept preferentially. Every
+ /// row at a retained distinct ob — including all tied NULLs — emits.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_nulls_first_ordering() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank_with_opts(
+ 2,
+ SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ true,
+ )?;
+
+ // pk=1 vals: NULL, 5, NULL, 8 → distinct-top-2 NULLS FIRST = {NULL, 5}
+ // pk=2 vals: 7, NULL → distinct-top-2 = {NULL, 7}
+ // pk=3 vals: 3, 1 → distinct-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 | |",
+ "| 1 | 5 |",
+ "| 2 | |",
+ "| 2 | 7 |",
+ "| 3 | 1 |",
+ "| 3 | 3 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// Total `GroupEntry` count across all partitions.
+ fn dense_rank_entry_count(state: &PartitionedTopKDenseRank) -> usize {
+ state
+ .states
+ .values()
+ .flat_map(|s| s.groups.values())
+ .map(|entries| entries.len())
+ .sum()
+ }
+
+ /// One source batch feeding many retained groups must be charged
+ /// once, not once per group.
+ ///
+ /// Dense-rank retains up to K distinct-ob groups per partition and
+ /// each can draw rows from the same batch, so charging per entry
+ /// inflates the reservation by (partitions × K) — here 6× — and can
+ /// trip a spurious `ResourcesExhausted`.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_charges_batch_once() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(3)?;
+
+ // pk=1 retains {1,2,3}, pk=2 retains {10,20,30}: 6 groups, all
+ // from this one batch.
+ let batch = pk_val_batch(
+ &schema,
+ vec![1, 1, 1, 1, 2, 2, 2, 2],
+ vec![1, 2, 3, 4, 10, 20, 30, 40],
+ )?;
+ let batch_bytes = get_record_batch_memory_size(&batch);
+ state.insert_batch(&batch)?;
+
+ assert_eq!(dense_rank_entry_count(&state), 6);
+ assert_eq!(state.store.len(), 1);
+ assert_eq!(state.store.batches_size, batch_bytes);
+ Ok(())
+ }
+
+ /// Evicting the last group referencing a batch must release the
+ /// batch's bytes, or the reservation only ever grows.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_releases_evicted_batch() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
+
+ let first = pk_val_batch(&schema, vec![1, 1], vec![50, 60])?;
+ state.insert_batch(&first)?;
+ assert_eq!(state.store.len(), 1);
+
+ // Both values beat {50, 60}, so every group from `first` is
+ // evicted and only `second` remains charged.
+ let second = pk_val_batch(&schema, vec![1, 1], vec![5, 6])?;
+ let second_bytes = get_record_batch_memory_size(&second);
+ state.insert_batch(&second)?;
+
+ assert_eq!(state.store.len(), 1);
+ assert_eq!(state.store.batches_size, second_bytes);
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 5 |",
+ "| 1 | 6 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// The mirror of the above: a batch whose every run is rejected must
+ /// not be charged at all.
+ ///
+ /// Nothing references it, so nothing would ever release it — charging
+ /// it would pin both the bytes and the batch for the operator's
+ /// lifetime.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_ignores_fully_rejected_batch()
-> Result<()>
+ {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
+
+ let first = pk_val_batch(&schema, vec![1, 1], vec![5, 6])?;
+ let first_bytes = get_record_batch_memory_size(&first);
+ state.insert_batch(&first)?;
+ assert_eq!(state.store.len(), 1);
+
+ // At K=2 with {5, 6} tracked, both values lose to the boundary, so
+ // no `GroupEntry` points at `second`.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![50, 60])?)?;
+
+ assert_eq!(dense_rank_entry_count(&state), 2);
+ assert_eq!(state.store.len(), 1);
+ assert_eq!(state.store.batches_size, first_bytes);
+ Ok(())
+ }
+
+ /// `keys` must grow on demand rather than reserve K slots when a
+ /// partition is first seen. `size()` charges `keys.capacity()`, so
+ /// eager sizing reserves O(partitions * K) for slots that never hold
+ /// a key — enough to fail a memory limit on a high-cardinality input
+ /// whose partitions each keep a handful of distinct values.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_heap_grows_on_demand() ->
Result<()> {
+ const K: usize = 512;
+ const PARTITIONS: usize = 64;
+ let (schema, mut state) = build_partitioned_topk_dense_rank(K)?;
+
+ // One row per partition: every partition holds exactly one
+ // distinct ob value, K - 1 slots short of capacity.
+ let pks: Vec<i32> = (0..PARTITIONS as i32).collect();
+ let vals = pks.clone();
+ state.insert_batch(&pk_val_batch(&schema, pks, vals)?)?;
+
+ assert_eq!(state.states.len(), PARTITIONS);
+ let heap_slots: usize = state.states.values().map(|s|
s.keys.capacity()).sum();
+ // Eager `with_capacity(K)` would reserve PARTITIONS * K = 32768.
+ assert!(
+ heap_slots <= PARTITIONS * 8,
+ "reserved {heap_slots} heap slots to hold {PARTITIONS} keys"
+ );
+ Ok(())
+ }
+
+ /// The encoded partition keys owned by `states`, and the backing
+ /// tables the drained scratch maps keep, are long-lived heap
+ /// allocations `size()` must charge: they persist for the operator's
+ /// life yet belong to no `GroupEntry`, so per-entry accounting can't
+ /// see them. With numerous or wide partition keys the key bytes are
+ /// the larger term.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_size_covers_keys_and_scratch()
-> Result<()>
+ {
+ const PARTITIONS: usize = 64;
+ const KEY_WIDTH: usize = 1024;
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("pk", DataType::Utf8, false),
+ Field::new("val", DataType::Int32, false),
+ ]));
+ let pk_expr: Arc<dyn PhysicalExpr> = col("pk", schema.as_ref())?;
+ let partition_sort_fields = build_sort_fields(
+ &[PhysicalSortExpr {
+ expr: Arc::clone(&pk_expr),
+ options: SortOptions::default(),
+ }],
+ &schema,
+ )?;
+ let order_expr = LexOrdering::from([PhysicalSortExpr {
+ expr: col("val", schema.as_ref())?,
+ options: SortOptions::default(),
+ }]);
+ let mut state = PartitionedTopKDenseRank::try_new(
+ 0,
+ Arc::clone(&schema),
+ vec![pk_expr],
+ partition_sort_fields,
+ order_expr,
+ 4,
+ 8, // batch_size
+ &Arc::new(RuntimeEnv::default()),
+ &ExecutionPlanMetricsSet::new(),
+ )?;
+
+ // One row per partition, each with a wide key.
+ let pks: Vec<String> = (0..PARTITIONS)
+ .map(|i| format!("{}{i:04}", "p".repeat(KEY_WIDTH - 4)))
+ .collect();
+ let vals: Vec<i32> = (0..PARTITIONS as i32).collect();
+ state.insert_batch(&RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(StringArray::from(pks)),
+ Arc::new(Int32Array::from(vals)),
+ ],
+ )?)?;
+ assert_eq!(state.states.len(), PARTITIONS);
+
+ let key_bytes: usize = state.states.keys().map(|pk|
pk.capacity()).sum();
+ let scratch_bytes = state.partition_groups.capacity()
+ * (size_of::<Vec<u8>>() + size_of::<Vec<u32>>())
+ + state.ob_runs.capacity() * (size_of::<Vec<u8>>() +
size_of::<Vec<u32>>());
+ assert!(key_bytes >= PARTITIONS * KEY_WIDTH, "key bytes {key_bytes}");
+ assert!(scratch_bytes > 0, "scratch tables never allocated");
+
+ // Reconstruct the total from its parts. Both terms above have to
+ // appear for this to balance, so dropping either from `size()`
+ // fails here rather than being absorbed by the slack in some
+ // other term.
+ let expected = size_of::<PartitionedTopKDenseRank>()
+ + state.row_converter.size()
+ + state.partition_converter.size()
+ + state.scratch_rows.size()
+ + state.partition_scratch_rows.size()
+ + key_bytes
+ + state.states.values().map(|s| s.size()).sum::<usize>()
+ + state.states.capacity()
+ * (size_of::<Vec<u8>>() + size_of::<DenseRankPartitionState>())
+ + scratch_bytes
+ + state.store.size();
+ assert_eq!(state.size(), expected);
+ Ok(())
+ }
+
+ /// A group admitted and then evicted within the *same*
+ /// `insert_batch` call: the batch is still pending (not yet handed to
+ /// the store), so releasing it must decrement the in-flight use count
+ /// rather than call `unuse` on an unregistered id.
+ ///
+ /// `ob_runs` drains in hash order, so with K=1 and many distinct
+ /// values the minimum is almost never seen first and the run
+ /// admit-then-evict path is taken repeatedly.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_evicts_same_call_group() ->
Result<()> {
+ let (schema, mut state) = build_partitioned_topk_dense_rank(1)?;
+
+ let pks = vec![1; 32];
+ let vals: Vec<i32> = (0..32).rev().collect();
+ let batch = pk_val_batch(&schema, pks, vals)?;
+ let batch_bytes = get_record_batch_memory_size(&batch);
+ state.insert_batch(&batch)?;
+
+ assert_eq!(dense_rank_entry_count(&state), 1);
+ assert_eq!(state.store.len(), 1);
+ assert_eq!(state.store.batches_size, batch_bytes);
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+",
+ "| pk | val |",
+ "+----+-----+",
+ "| 1 | 0 |",
+ "+----+-----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
}
diff --git a/datafusion/sqllogictest/test_files/window_topn.slt
b/datafusion/sqllogictest/test_files/window_topn.slt
index a9a52a654f..8f65e496b5 100644
--- a/datafusion/sqllogictest/test_files/window_topn.slt
+++ b/datafusion/sqllogictest/test_files/window_topn.slt
@@ -365,9 +365,115 @@ physical_plan
04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST],
preserve_partitioning=[false]
05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+# Sibling safety: an aggregate sibling (SUM) is computed over the pruned
+# input, so the rule must NOT fire even though the filter is on ROW_NUMBER.
+# SUM over the default RANGE frame includes tie-peers that pruning would
+# drop, giving a wrong sum. The plan keeps SortExec + FilterExec — no
+# PartitionedTopKExec.
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *,
+ ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+ SUM(val) OVER (PARTITION BY pk ORDER BY val) as running_sum
+ FROM window_topn_t
+) WHERE rn <= 3;
+----
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number()
PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST]
RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn,
sum(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY
[window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND
CURRENT ROW@4 as running_sum]
+02)--FilterExec: row_number() PARTITION BY [window_topn_t.pk] ORDER BY
[window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND
CURRENT ROW@3 <= 3
+03)----BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk]
ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY
[window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND
CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT
ROW, sum(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY
[window_topn_t.val ASC NULLS LAST] R [...]
+04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST],
preserve_partitioning=[false]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Sibling safety: a LEAD sibling reads following (pruned) rows, so the
+# rule must NOT fire — the boundary row's LEAD would resolve to a pruned
+# row. The plan keeps SortExec + FilterExec.
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *,
+ ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+ LEAD(val) OVER (PARTITION BY pk ORDER BY val) as next_val
+ FROM window_topn_t
+) WHERE rn <= 3;
+----
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number()
PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST]
RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn,
lead(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY
[window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND
CURRENT ROW@4 as next_val]
+02)--FilterExec: row_number() PARTITION BY [window_topn_t.pk] ORDER BY
[window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND
CURRENT ROW@3 <= 3
+03)----BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk]
ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY
[window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND
CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT
ROW, lead(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY
[window_topn_t.val ASC NULLS LAST] [...]
+04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST],
preserve_partitioning=[false]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+
statement ok
SET datafusion.explain.physical_plan_only = false;
+# Sibling safety (correctness): LEAD reads following rows, so if the rule
+# fired, the pruned input would give a wrong LEAD at the retained boundary
+# — each rn=2 row's next_val would become NULL instead of the (pruned)
+# rn=3 row's value. The guard keeps the normal plan, so LEAD sees the full
+# partition and next_val is correct (30 / 25 / 100 for the rn=2 rows).
+query IIIII rowsort
+SELECT id, pk, val, rn, next_val FROM (
+ SELECT *,
+ ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+ LEAD(val) OVER (PARTITION BY pk ORDER BY val) as next_val
+ FROM window_topn_t
+) WHERE rn <= 2;
+----
+1 1 10 1 20
+10 3 75 2 100
+2 1 20 2 30
+5 2 5 1 15
+6 2 15 2 25
+9 3 50 1 75
+
+# A tie table for sibling-safety correctness: pk=1 has three peers at
+# val=5, one 10, one 20 (row_number 1..5; rank 1,1,1,4,5; dense_rank
+# 1,1,1,2,3).
+statement ok
+CREATE TABLE window_topn_sib_t (id INT, pk INT, val INT) AS VALUES
+ (1, 1, 5),
+ (2, 1, 5),
+ (3, 1, 5),
+ (4, 1, 10),
+ (5, 1, 20);
+
+# SUM sibling (correctness): fetch=2 but there are three peers at val=5,
+# so one peer would be pruned. The guard keeps the normal plan, so the
+# RANGE frame sees all three 5s and running_sum = 15. If the rule fired,
+# the pruned 2-row input would sum to 10. (id excluded — which two of the
+# three tied rows get row_number 1/2 is not deterministic.)
+query III rowsort
+SELECT pk, val, running_sum FROM (
+ SELECT *,
+ ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+ SUM(val) OVER (PARTITION BY pk ORDER BY val) as running_sum
+ FROM window_topn_sib_t
+) WHERE rn <= 2;
+----
+1 5 15
+1 5 15
+
+# All three ranking siblings (ROW_NUMBER + RANK + DENSE_RANK) over the same
+# window are prune-safe, so the rule DOES fire (filter on rn). Verify the
+# rank/dense_rank values stay correct under the optimized plan — including
+# the val=10 row where row_number=4, rank=4, dense_rank=2 all differ.
+query IIII rowsort
+SELECT pk, val, rnk, dr FROM (
+ SELECT *,
+ ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+ RANK() OVER (PARTITION BY pk ORDER BY val) as rnk,
+ DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr
+ FROM window_topn_sib_t
+) WHERE rn <= 4;
+----
+1 10 4 2
+1 5 1 1
+1 5 1 1
+1 5 1 1
+
+statement ok
+DROP TABLE window_topn_sib_t;
+
# Test 15: ROW_NUMBER with DESC ordering — correctness
query III rowsort
SELECT id, pk, val FROM (
@@ -1089,6 +1195,353 @@ SELECT id, pk, val FROM (
statement ok
DROP TABLE window_topn_rank_null_t;
+###############################################################################
+# DENSE_RANK() tests
+###############################################################################
+#
+# DENSE_RANK semantics: `WHERE dr <= K` keeps every row whose ORDER BY
+# value is among the K distinct-smallest ORDER BY values in the
+# partition. The total kept per partition is unbounded in
+# rows-per-distinct-value (unlike RANK, which is bounded above by
+# `K + ties at the boundary`).
+#
+# The tests below exercise:
+# - heavy ties within a distinct ob key (all rows retained per key)
+# - cross-batch appends under the same key
+# - eviction of the entire max-key group when a strictly-smaller
+# distinct ob arrives
+# - the empty-ORDER-BY degenerate case (rule must NOT fire)
+
+statement ok
+SET datafusion.optimizer.enable_window_topn = true;
+
+statement ok
+CREATE TABLE window_topn_dense_rank_t (id INT, pk INT, val INT) AS VALUES
+ -- pk=1: distinct top-2 = {10, 20}. Every row at those keeps dr <= 2.
+ -- 10 appears once, 20 appears three times (heavy tie), 30/40 dropped.
+ (1, 1, 10),
+ (2, 1, 20),
+ (3, 1, 20),
+ (4, 1, 20),
+ (5, 1, 30),
+ (6, 1, 40),
+ -- pk=2: distinct top-2 = {5, 15}. Two rows at 5, one at 15.
+ (7, 2, 5),
+ (8, 2, 5),
+ (9, 2, 15),
+ (10, 2, 25),
+ -- pk=3: 100 then four 200s — 200 group is the "boundary-max" but with
+ -- dense_rank <= 2 both distinct values (100, 200) are retained.
+ (11, 3, 100),
+ (12, 3, 200),
+ (13, 3, 200),
+ (14, 3, 200),
+ (15, 3, 200),
+ (16, 3, 300);
+
+# Test DR1: Basic DENSE_RANK correctness — every row at the K distinct
+# smallest values is retained.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM
window_topn_dense_rank_t
+) WHERE dr <= 2;
+----
+1 1 10
+11 3 100
+12 3 200
+13 3 200
+14 3 200
+15 3 200
+2 1 20
+3 1 20
+4 1 20
+7 2 5
+8 2 5
+9 2 15
+
+# Test DR2: EXPLAIN shows PartitionedTopKExec with fn=dense_rank
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM
window_topn_dense_rank_t
+) WHERE dr <= 2;
+----
+logical_plan
+01)Projection: window_topn_dense_rank_t.id, window_topn_dense_rank_t.pk,
window_topn_dense_rank_t.val, dense_rank() PARTITION BY
[window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS dr
+02)--Filter: dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY
[window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW <= UInt64(2)
+03)----WindowAggr: windowExpr=[[dense_rank() PARTITION BY
[window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+04)------TableScan: window_topn_dense_rank_t projection=[id, pk, val]
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, dense_rank()
PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY
[window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW@3 as dr]
+02)--BoundedWindowAggExec: wdw=[dense_rank() PARTITION BY
[window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "dense_rank()
PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY
[window_topn_dense_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)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1,
maintains_sort_order=true
+04)------PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1],
order=[val@2 ASC NULLS LAST]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Test DR3: dr < 3 should give the same results (fetch = K-1 = 2)
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM
window_topn_dense_rank_t
+) WHERE dr < 3;
+----
+1 1 10
+11 3 100
+12 3 200
+13 3 200
+14 3 200
+15 3 200
+2 1 20
+3 1 20
+4 1 20
+7 2 5
+8 2 5
+9 2 15
+
+# Test DR4: flipped predicate (2 >= dr) also fires the rule
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM
window_topn_dense_rank_t
+) WHERE 2 >= dr;
+----
+1 1 10
+11 3 100
+12 3 200
+13 3 200
+14 3 200
+15 3 200
+2 1 20
+3 1 20
+4 1 20
+7 2 5
+8 2 5
+9 2 15
+
+# Test DR5: K exceeds every partition's distinct count — nothing dropped.
+# pk=1: 4 distinct, pk=2: 3 distinct, pk=3: 3 distinct. dr <= 100 retains all.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM
window_topn_dense_rank_t
+) WHERE dr <= 100;
+----
+1 1 10
+10 2 25
+11 3 100
+12 3 200
+13 3 200
+14 3 200
+15 3 200
+16 3 300
+2 1 20
+3 1 20
+4 1 20
+5 1 30
+6 1 40
+7 2 5
+8 2 5
+9 2 15
+
+# Test DR6: DENSE_RANK without PARTITION BY — should NOT trigger the
optimization
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *, DENSE_RANK() OVER (ORDER BY val) as dr FROM
window_topn_dense_rank_t
+) WHERE dr <= 2;
+----
+logical_plan
+01)Projection: window_topn_dense_rank_t.id, window_topn_dense_rank_t.pk,
window_topn_dense_rank_t.val, dense_rank() ORDER BY
[window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW AS dr
+02)--Filter: dense_rank() ORDER BY [window_topn_dense_rank_t.val ASC NULLS
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(2)
+03)----WindowAggr: windowExpr=[[dense_rank() ORDER BY
[window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW]]
+04)------TableScan: window_topn_dense_rank_t projection=[id, pk, val]
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, dense_rank()
ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED
PRECEDING AND CURRENT ROW@3 as dr]
+02)--FilterExec: dense_rank() ORDER BY [window_topn_dense_rank_t.val ASC NULLS
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 <= 2
+03)----BoundedWindowAggExec: wdw=[dense_rank() ORDER BY
[window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW: Field { "dense_rank() ORDER BY [window_topn_dense_rank_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=[val@2 ASC NULLS LAST], preserve_partitioning=[false]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Test DR7: DENSE_RANK with empty ORDER BY — degenerate (every row at
+# dense_rank 1), rule must NOT fire.
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk) as dr FROM
window_topn_dense_rank_t
+) WHERE dr <= 3;
+----
+logical_plan
+01)Projection: window_topn_dense_rank_t.id, window_topn_dense_rank_t.pk,
window_topn_dense_rank_t.val, dense_rank() PARTITION BY
[window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED
FOLLOWING AS dr
+02)--Filter: dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ROWS
BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING <= UInt64(3)
+03)----WindowAggr: windowExpr=[[dense_rank() PARTITION BY
[window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED
FOLLOWING]]
+04)------TableScan: window_topn_dense_rank_t projection=[id, pk, val]
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, dense_rank()
PARTITION BY [window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND
UNBOUNDED FOLLOWING@3 as dr]
+02)--FilterExec: dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ROWS
BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@3 <= 3
+03)----BoundedWindowAggExec: wdw=[dense_rank() PARTITION BY
[window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED
FOLLOWING: Field { "dense_rank() PARTITION BY [window_topn_dense_rank_t.pk]
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING": UInt64 }, frame:
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING], mode=[Sorted]
+04)------SortExec: expr=[pk@1 ASC NULLS LAST], preserve_partitioning=[false]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Test DR8: DESC ordering — distinct-top-2 DESC per partition.
+# pk=1 DESC {40, 30, 20, 10}: top-2 = {40, 30} → 2 rows
+# pk=2 DESC {25, 15, 5}: top-2 = {25, 15} → 2 rows
+# pk=3 DESC {300, 200, 100}: top-2 = {300, 200} → 5 rows (200 appears 4x)
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val DESC) as dr FROM
window_topn_dense_rank_t
+) WHERE dr <= 2;
+----
+10 2 25
+12 3 200
+13 3 200
+14 3 200
+15 3 200
+16 3 300
+5 1 30
+6 1 40
+9 2 15
+
+# Test DR9: multi-column PARTITION BY
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk, id ORDER BY val) as dr FROM
window_topn_dense_rank_t
+) WHERE dr <= 1;
+----
+1 1 10
+10 2 25
+11 3 100
+12 3 200
+13 3 200
+14 3 200
+15 3 200
+16 3 300
+2 1 20
+3 1 20
+4 1 20
+5 1 30
+6 1 40
+7 2 5
+8 2 5
+9 2 15
+
+# Test DR10: mixed ROW_NUMBER + DENSE_RANK — filter on DR should optimize
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT *,
+ ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn,
+ DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr
+ FROM window_topn_dense_rank_t
+) WHERE dr <= 2;
+----
+logical_plan
+01)Projection: window_topn_dense_rank_t.id, window_topn_dense_rank_t.pk,
window_topn_dense_rank_t.val, row_number() PARTITION BY
[window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn, dense_rank()
PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY
[window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW AS dr
+02)--Filter: dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY
[window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW <= UInt64(2)
+03)----WindowAggr: windowExpr=[[row_number() PARTITION BY
[window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, dense_rank() PARTITION
BY [window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC
NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+04)------TableScan: window_topn_dense_rank_t projection=[id, pk, val]
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number()
PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY
[window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW@3 as rn, dense_rank() PARTITION BY
[window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as dr]
+02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY
[window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number()
PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY
[window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND
CURRENT ROW, dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] OR [...]
+03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1,
maintains_sort_order=true
+04)------PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1],
order=[val@2 ASC NULLS LAST]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Test DR11: `dr < 1` keeps nothing (limit_n = 0). The rule must NOT
+# fire (a fetch=0 PartitionedTopKExec would panic); the ordinary
+# FilterExec returns the empty result.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM
window_topn_dense_rank_t
+) WHERE dr < 1;
+----
+
+# Test DR11b: ORDER BY keys fully covered by the PARTITION BY prefix
+# (`PARTITION BY pk ORDER BY pk`). The deduplicated sort ordering is just
+# `[pk]`, so no order key survives past the partition prefix. The rule
+# must NOT fire — a PartitionedTopKExec built with zero order expressions
+# panics on execution. Every row ties at dense_rank 1, so `dr <= 2` keeps
+# the whole table (returned via the ordinary window + filter plan).
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY pk) as dr FROM
window_topn_dense_rank_t
+) WHERE dr <= 2;
+----
+1 1 10
+10 2 25
+11 3 100
+12 3 200
+13 3 200
+14 3 200
+15 3 200
+16 3 300
+2 1 20
+3 1 20
+4 1 20
+5 1 30
+6 1 40
+7 2 5
+8 2 5
+9 2 15
+
+statement ok
+DROP TABLE window_topn_dense_rank_t;
+
+# ---------------------------------------------------------------------------
+# DENSE_RANK NULL-in-ORDER-BY ordering
+#
+# Under DENSE_RANK a NULL is a distinct ob value occupying its own rank
+# slot; whether it lands among the K distinct-smallest depends on
+# NULLS FIRST / NULLS LAST. Correctness rests entirely on the
+# byte-comparable row encoding driving the distinct-ob key ordering, so
+# exercise both null placements explicitly.
+# ---------------------------------------------------------------------------
+
+statement ok
+CREATE TABLE window_topn_dr_null_t (id INT, pk INT, val INT) AS VALUES
+ (1, 1, NULL),
+ (2, 1, NULL),
+ (3, 1, 10),
+ (4, 1, 20),
+ (5, 1, 30);
+
+# Test DR12: NULLS FIRST — NULL is the smallest distinct ob (dr 1), then
+# 10 (dr 2). `dr <= 2` keeps both NULL rows and val=10.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS FIRST)
as dr
+ FROM window_topn_dr_null_t
+) WHERE dr <= 2;
+----
+1 1 NULL
+2 1 NULL
+3 1 10
+
+# Test DR13: EXPLAIN confirms the NULLS FIRST case uses PartitionedTopKExec.
+query TT
+EXPLAIN SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS FIRST)
as dr
+ FROM window_topn_dr_null_t
+) WHERE dr <= 2;
+----
+logical_plan
+01)Projection: window_topn_dr_null_t.id, window_topn_dr_null_t.pk,
window_topn_dr_null_t.val
+02)--Filter: dense_rank() PARTITION BY [window_topn_dr_null_t.pk] ORDER BY
[window_topn_dr_null_t.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING
AND CURRENT ROW <= UInt64(2)
+03)----WindowAggr: windowExpr=[[dense_rank() PARTITION BY
[window_topn_dr_null_t.pk] ORDER BY [window_topn_dr_null_t.val ASC NULLS FIRST]
RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+04)------TableScan: window_topn_dr_null_t projection=[id, pk, val]
+physical_plan
+01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val]
+02)--BoundedWindowAggExec: wdw=[dense_rank() PARTITION BY
[window_topn_dr_null_t.pk] ORDER BY [window_topn_dr_null_t.val ASC NULLS FIRST]
RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "dense_rank()
PARTITION BY [window_topn_dr_null_t.pk] ORDER BY [window_topn_dr_null_t.val ASC
NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 },
frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1,
maintains_sort_order=true
+04)------PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1],
order=[val@2 ASC]
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Test DR14: NULLS LAST — NULL is the largest distinct ob (dr 4). 10 (dr
+# 1) and 20 (dr 2) are the two smallest, so `dr <= 2` excludes the NULLs.
+query III rowsort
+SELECT id, pk, val FROM (
+ SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS LAST) as
dr
+ FROM window_topn_dr_null_t
+) WHERE dr <= 2;
+----
+3 1 10
+4 1 20
+
+statement ok
+DROP TABLE window_topn_dr_null_t;
+
# Reset config to default (false)
statement ok
SET datafusion.optimizer.enable_window_topn = false;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
