perf: Extend WindowTopN to `dense_rank` by SubhamSinghal · Pull Request #23869 · apache/datafusion

SubhamSinghal · 2026-07-24T15:28:36Z

Which issue does this PR close?

Related to Optimize "per partition" top-k : ROW_NUMBER < 5 / TopK #6899. Completes the ROW_NUMBER + RANK + DENSE_RANK trio requested by the issue (ROW_NUMBER shipped in Perf: Window topn optimisation #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

github-actions · 2026-07-24T15:50:42Z

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.1.0 (current)
       Built [ 104.507s] (current)
     Parsing datafusion v54.1.0 (current)
      Parsed [   0.035s] (current)
    Building datafusion v54.1.0 (baseline)
       Built [ 104.718s] (baseline)
     Parsing datafusion v54.1.0 (baseline)
      Parsed [   0.035s] (baseline)
    Checking datafusion v54.1.0 -> v54.1.0 (no change; assume patch)
     Checked [   0.632s] 223 checks: 223 pass, 30 skip
     Summary no semver update required
    Finished [ 211.639s] datafusion
    Building datafusion-physical-optimizer v54.1.0 (current)
       Built [  40.743s] (current)
     Parsing datafusion-physical-optimizer v54.1.0 (current)
      Parsed [   0.022s] (current)
    Building datafusion-physical-optimizer v54.1.0 (baseline)
       Built [  40.677s] (baseline)
     Parsing datafusion-physical-optimizer v54.1.0 (baseline)
      Parsed [   0.021s] (baseline)
    Checking datafusion-physical-optimizer v54.1.0 -> v54.1.0 (no change; assume patch)
     Checked [   0.118s] 223 checks: 223 pass, 30 skip
     Summary no semver update required
    Finished [  82.539s] datafusion-physical-optimizer
    Building datafusion-physical-plan v54.1.0 (current)
       Built [  37.625s] (current)
     Parsing datafusion-physical-plan v54.1.0 (current)
      Parsed [   0.144s] (current)
    Building datafusion-physical-plan v54.1.0 (baseline)
       Built [  37.077s] (baseline)
     Parsing datafusion-physical-plan v54.1.0 (baseline)
      Parsed [   0.145s] (baseline)
    Checking datafusion-physical-plan v54.1.0 -> v54.1.0 (no change; assume patch)
     Checked [   0.616s] 223 checks: 222 pass, 1 fail, 0 warn, 30 skip

--- failure enum_variant_added: enum variant added on exhaustive enum ---

Description:
A publicly-visible enum without #[non_exhaustive] has a new variant.
        ref: https://doc.rust-lang.org/cargo/reference/semver.html#enum-variant-new
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.49.0/src/lints/enum_variant_added.ron

Failed in:
  variant WindowFnKind:DenseRank in /home/runner/work/datafusion/datafusion/datafusion/physical-plan/src/sorts/partitioned_topk.rs:76

     Summary semver requires new major version: 1 major and 0 minor checks failed
    Finished [  76.689s] datafusion-physical-plan
    Building datafusion-sqllogictest v54.1.0 (current)
       Built [ 182.469s] (current)
     Parsing datafusion-sqllogictest v54.1.0 (current)
      Parsed [   0.021s] (current)
    Building datafusion-sqllogictest v54.1.0 (baseline)
       Built [ 182.258s] (baseline)
     Parsing datafusion-sqllogictest v54.1.0 (baseline)
      Parsed [   0.023s] (baseline)
    Checking datafusion-sqllogictest v54.1.0 -> v54.1.0 (no change; assume patch)
     Checked [   0.089s] 223 checks: 223 pass, 30 skip
     Summary no semver update required
    Finished [ 367.223s] datafusion-sqllogictest

codecov-commenter · 2026-07-24T16:10:42Z

Codecov Report

❌ Patch coverage is 88.21363% with 64 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.70%. Comparing base (582453b) to head (4530bba).
⚠️ Report is 42 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/topk/mod.rs	87.91%	13 Missing and 49 partials ⚠️
...fusion/physical-plan/src/sorts/partitioned_topk.rs	87.50%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #23869      +/-   ##
==========================================
- Coverage   80.72%   80.70%   -0.02%     
==========================================
  Files        1090     1095       +5     
  Lines      370384   373167    +2783     
  Branches   370384   373167    +2783     
==========================================
+ Hits       298975   301163    +2188     
- Misses      53590    54002     +412     
- Partials    17819    18002     +183

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

SubhamSinghal · 2026-07-28T05:53:57Z

@kumarUjjawal @kosiew can you please review this PR?

kumarUjjawal

Hi @SubhamSinghal I have left some comments please let me know what you think.

kumarUjjawal · 2026-07-28T06:30:35Z

    match udf.fun().name() {
        "row_number" => Some(WindowFnKind::RowNumber),
        "rank" => Some(WindowFnKind::Rank),
+        "dense_rank" => Some(WindowFnKind::DenseRank),


Could we skip this rewrite when another expression in the same window operator needs following rows? The input is pruned before all window expressions run, so a sibling lead() returns the wrong value at the retained boundary. We can add a dense_rank with lead regression and only rewrite when every sibling expression is safe under tail pruning.

Fixed in 4530bba. The rewrite now bails unless every window expr is ROW_NUMBER/RANK/DENSE_RANK over the same PARTITION BY/ORDER BY. LEAD and aggregates are excluded. Added EXPLAIN + correctness SLT regressions for SUM and LEAD siblings.

kumarUjjawal · 2026-07-28T06:31:16Z

+        // For RANK / DENSE_RANK: an empty ORDER BY makes every row tie
+        // at rank 1 — the optimization is degenerate (we'd retain the
+        // entire input) and tie storage would be unbounded.
+        if matches!(fn_kind, WindowFnKind::Rank | WindowFnKind::DenseRank)


This guard misses orderings that become empty after partition keys are removed. DENSE_RANK() OVER (PARTITION BY pk ORDER BY pk) reaches PartitionedTopKExec without order expressions and panics. Please validate the effective order keys and add a regression.

Fixed in 4530bba. Replaced order_by().is_empty() with a structural guard: sort_exec.expr().len() <= partition_prefix_len → bail. Catches PARTITION BY pk ORDER BY pk (and no-ORDER-BY ROW_NUMBER). Added SLT regression.

kumarUjjawal · 2026-07-28T06:33:13Z

+                    entries.push(TieEntry {
+                        batch: batch.clone(),
+                        row_indices: run_indices,
+                        batch_bytes: input_batch_bytes,


records the full source-batch size for every retained group even though the batch clones share Arrow buffers. With many partitions, memory reservation is heavily inflated and can fail with ResourcesExhausted. This expands #23326 to the normal dense-rank path; please account each source batch once or compact the retained rows.

This is the DENSE_RANK instance of #23326 (same TieEntry/size() shared-buffer over-count as RANK). Keeping storage consistent so one fix in #23326 covers both operators.

kumarUjjawal · 2026-07-28T06:33:44Z

+                        .groups
+                        .keys()
+                        .map(|key| key.as_slice())
+                        .max()


Finding the maximum scans all K keys for every new distinct value. On mostly distinct input this makes the rewrite O(N × K) and can be much slower for large K than the sort it replaces.

I had kept it unoptimised thinking K will be small. Fixed in 4530bba. Cache the admission boundary (max_key); the O(K) scan runs only when stale (after eviction/fill), so the common "worse than boundary → drop" path is O(1).

perf: Extend WindowTopN to dense rank

d0c0314

github-actions Bot added optimizer Optimizer rules core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) physical-plan Changes to the physical-plan crate labels Jul 24, 2026

github-actions Bot added the auto detected api change Auto detected API change label Jul 24, 2026

comphead mentioned this pull request Jul 24, 2026

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

Open

comphead changed the title ~~perf: Extend WindowTopN to dense rank~~ perf: Extend WindowTopN to dense_rank Jul 24, 2026

Merge branch 'main' into window-topn-dense-rank

32b3b5c

kumarUjjawal reviewed Jul 28, 2026

View reviewed changes

address review comments

4530bba

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

perf: Extend WindowTopN to `dense_rank` - #23869

perf: Extend WindowTopN to `dense_rank`#23869
SubhamSinghal wants to merge 3 commits into
apache:mainfrom
SubhamSinghal:window-topn-dense-rank

SubhamSinghal commented Jul 24, 2026

Uh oh!

github-actions Bot commented Jul 24, 2026 •

edited

Loading

Uh oh!

codecov-commenter commented Jul 24, 2026 •

edited

Loading

Uh oh!

SubhamSinghal commented Jul 28, 2026

Uh oh!

kumarUjjawal left a comment

Uh oh!

kumarUjjawal Jul 28, 2026

Uh oh!

SubhamSinghal Jul 28, 2026

Uh oh!

kumarUjjawal Jul 28, 2026

Uh oh!

SubhamSinghal Jul 28, 2026

Uh oh!

kumarUjjawal Jul 28, 2026

Uh oh!

SubhamSinghal Jul 28, 2026

Uh oh!

kumarUjjawal Jul 28, 2026

Uh oh!

SubhamSinghal Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

SubhamSinghal commented Jul 24, 2026

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!

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

Uh oh!

Uh oh!

codecov-commenter commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

SubhamSinghal commented Jul 28, 2026

Uh oh!

kumarUjjawal left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

github-actions Bot commented Jul 24, 2026 •

edited

Loading

codecov-commenter commented Jul 24, 2026 •

edited

Loading