perf: Extend WindowTopN to support RANK by SubhamSinghal · Pull Request #22885 · apache/datafusion

SubhamSinghal · 2026-06-10T17:00:38Z

Which issue does this PR close?

Related to Optimize "per partition" top-k : ROW_NUMBER < 5 / TopK #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 perf: share encoder/reservation across PartitionedTopKExec partition … #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

SubhamSinghal · 2026-06-15T05:16:32Z

@kumarUjjawal @2010YOUY01 will you be able to review this PR.

comphead · 2026-06-19T17:38:17Z

might be related #23021

github-actions · 2026-06-30T17:06:57Z

Thank you for opening this pull request!

Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).

Details

     Cloning apache/main
    Building datafusion v54.0.0 (current)
       Built [ 107.234s] (current)
     Parsing datafusion v54.0.0 (current)
      Parsed [   0.038s] (current)
    Building datafusion v54.0.0 (baseline)
       Built [ 107.646s] (baseline)
     Parsing datafusion v54.0.0 (baseline)
      Parsed [   0.038s] (baseline)
    Checking datafusion v54.0.0 -> v54.0.0 (no change; assume patch)
     Checked [   0.956s] 223 checks: 223 pass, 30 skip
     Summary no semver update required
    Finished [ 218.224s] datafusion
    Building datafusion-physical-optimizer v54.0.0 (current)
       Built [  39.349s] (current)
     Parsing datafusion-physical-optimizer v54.0.0 (current)
      Parsed [   0.024s] (current)
    Building datafusion-physical-optimizer v54.0.0 (baseline)
       Built [  38.304s] (baseline)
     Parsing datafusion-physical-optimizer v54.0.0 (baseline)
      Parsed [   0.025s] (baseline)
    Checking datafusion-physical-optimizer v54.0.0 -> v54.0.0 (no change; assume patch)
     Checked [   0.187s] 223 checks: 223 pass, 30 skip
     Summary no semver update required
    Finished [  79.389s] datafusion-physical-optimizer
    Building datafusion-physical-plan v54.0.0 (current)
       Built [  35.598s] (current)
     Parsing datafusion-physical-plan v54.0.0 (current)
      Parsed [   0.148s] (current)
    Building datafusion-physical-plan v54.0.0 (baseline)
       Built [  36.255s] (baseline)
     Parsing datafusion-physical-plan v54.0.0 (baseline)
      Parsed [   0.149s] (baseline)
    Checking datafusion-physical-plan v54.0.0 -> v54.0.0 (no change; assume patch)
     Checked [   0.937s] 223 checks: 222 pass, 1 fail, 0 warn, 30 skip

--- failure method_parameter_count_changed: pub method parameter count changed ---

Description:
A publicly-visible method now takes a different number of parameters, not counting the receiver (self) parameter.
        ref: https://doc.rust-lang.org/cargo/reference/semver.html#fn-change-arity
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.48.0/src/lints/method_parameter_count_changed.ron

Failed in:
  datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec::try_new takes 4 parameters in /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/3f697da4e542e016b2a8f02a2a7f5d525c982695/datafusion/physical-plan/src/sorts/partitioned_topk.rs:196, but now takes 5 parameters in /home/runner/work/datafusion/datafusion/datafusion/physical-plan/src/sorts/partitioned_topk.rs:227

     Summary semver requires new major version: 1 major and 0 minor checks failed
    Finished [  74.965s] datafusion-physical-plan
    Building datafusion-sqllogictest v54.0.0 (current)
       Built [ 190.398s] (current)
     Parsing datafusion-sqllogictest v54.0.0 (current)
      Parsed [   0.023s] (current)
    Building datafusion-sqllogictest v54.0.0 (baseline)
       Built [ 189.344s] (baseline)
     Parsing datafusion-sqllogictest v54.0.0 (baseline)
      Parsed [   0.025s] (baseline)
    Checking datafusion-sqllogictest v54.0.0 -> v54.0.0 (no change; assume patch)
     Checked [   0.122s] 223 checks: 223 pass, 30 skip
     Summary no semver update required
    Finished [ 383.463s] datafusion-sqllogictest

SubhamSinghal · 2026-07-01T07:05:13Z

@kosiew can you help in reviewing this PR.

kosiew

@SubhamSinghal
Looks good overall.
I left a couple of suggestions to keep the docs accurate and to call out a possible memory-accounting follow-up.

kosiew · 2026-07-02T10:37:48Z

Small docs nit: the Limitations section still says this operator is only activated for ROW_NUMBER. Could you update this sentence to mention RANK and the non-empty ORDER BY requirement?

kosiew · 2026-07-02T10:39:26Z

+
+            // Commit this batch's ties as a single entry.
+            if !equal_indices.is_empty() {
+                state.ties.push(TieEntry {


This looks correct for preserving tied rows. One possible follow-up: TieEntry keeps a clone of the full input RecordBatch and charges input_batch_bytes per partition tie entry, which could over-reserve memory when a single batch has boundary ties across many partitions.

It may be worth looking at compact tie batches or shared batch accounting later.

I will file a followup issue for this.

created issue for this: #23326

kosiew · 2026-07-03T11:54:19Z

@SubhamSinghal
Thanks for working on this.

Leaving this open for a few days for more reviews as this is a big PR.

SubhamSinghal · 2026-07-14T04:42:41Z

@kosiew can we merge this PR?

kumarUjjawal · 2026-07-14T07:28:03Z

Thank you @SubhamSinghal and @kosiew

## Which issue does this PR close? - Related to apache#6899 (DENSE_RANK to follow in a separate PR will close it). ## Rationale for this change PR apache#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 apache#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

## Which issue does this PR close?  - Part of apache#21706 ## Rationale for this change Updated sql queries to match existing benchmark queries, mostly apache#22885 and apache#24050 ## What changes are included in this PR? sql benchmark query updates. ## Are these changes tested? Yes ## Are there any user-facing changes? No

github-actions Bot added optimizer Optimizer rules core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) physical-plan Changes to the physical-plan crate auto detected api change Auto detected API change labels Jun 10, 2026

Extend WindowTopN to support RANK

173b53c

SubhamSinghal force-pushed the window-topn-rank branch from a2585c8 to 173b53c Compare June 30, 2026 16:52

github-actions Bot removed optimizer Optimizer rules core Core DataFusion crate physical-plan Changes to the physical-plan crate labels Jun 30, 2026

Extend WindowTopN to support RANK

a274a1b

github-actions Bot added optimizer Optimizer rules core Core DataFusion crate physical-plan Changes to the physical-plan crate and removed auto detected api change Auto detected API change labels Jun 30, 2026

github-actions Bot added the auto detected api change Auto detected API change label Jun 30, 2026

Fix build failure

78c0c93

kosiew approved these changes Jul 2, 2026

View reviewed changes

Update doc

d894dd8

SubhamSinghal mentioned this pull request Jul 5, 2026

PartitionedTopKRank over-accounts memory when a single batch has boundary ties across many partitions #23326

Open

Merge branch 'main' into window-topn-rank

3dba92f

kumarUjjawal added this pull request to the merge queue Jul 14, 2026

Merged via the queue into apache:main with commit 79692c1 Jul 14, 2026
40 checks passed

SubhamSinghal deleted the window-topn-rank branch July 14, 2026 07:42

comphead mentioned this pull request Jul 24, 2026

Support WindowGroupLimitExec (window-based limit pushdown) apache/datafusion-comet#4837

Closed

Omega359 mentioned this pull request Aug 22, 2026

Sql benchmark sync updated queries #24583

Merged

alamb mentioned this pull request Aug 24, 2026

Enable TopK per Partition (datafusion.optimizer.enable_window_topn) by default #24608

Open

7 tasks

Uh oh!

Conversation

SubhamSinghal commented Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

SubhamSinghal commented Jun 15, 2026

Uh oh!

comphead commented Jun 19, 2026

Uh oh!

github-actions Bot commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

SubhamSinghal commented Jul 1, 2026

Uh oh!

kosiew left a comment

Choose a reason for hiding this comment

Uh oh!

kosiew Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

SubhamSinghal Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

kosiew Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

SubhamSinghal Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

SubhamSinghal Jul 5, 2026

Choose a reason for hiding this comment

Uh oh!

kosiew commented Jul 3, 2026

Uh oh!

SubhamSinghal commented Jul 14, 2026

Uh oh!

kumarUjjawal commented Jul 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

SubhamSinghal commented Jun 10, 2026 •

edited

Loading

github-actions Bot commented Jun 30, 2026 •

edited

Loading