minor: share common state between the three PartitionedTopK operators

[jayzhan211]

Which issue does this PR close?

• N/A — refactor, no issue filed.

Rationale for this change

PartitionedTopK (ROW_NUMBER), PartitionedTopKRank (RANK) and PartitionedTopKDenseRank (DENSE_RANK) were added one after another, each as a copy of the previous one. They differ in the per-partition state and in how one partition's rows are classified, but every copy also repeats:

• the same 11 fields (schema, metrics, reservation, both expression lists, both row converters, both scratch Rows, k, batch_size) and their doc comments,
• the same try_new body,
• the same encode step at the top of insert_batch (evaluate and row-encode the partition and ORDER BY columns), plus the group-rows-by-partition loop in RANK and DENSE_RANK,
• the same tail of insert_batch (record replacements, resize the reservation),
• the same shared terms in size(), and in RANK / DENSE_RANK the same per-partition map accounting.

A fix to one copy (for example the partition-key byte accounting in size()) has to be remembered for the other two.

What changes are included in this PR?

No behaviour change.

• Add a private PartitionedTopKBase holding the shared fields, with try_new, encode, group_by_partition, finish_batch and size helpers, plus a partition_states_size helper for the per-partition map that RANK and DENSE_RANK both account the same way.
• Each operator becomes base + its own state. ROW_NUMBER keeps its single-pass admission, shared store and O(1) size(); RANK and DENSE_RANK keep their partition_groups scratch map (now drained in place) and their classification loops.
• PartitionedTopKExec and the crate-visible signatures of the three types are unchanged. Memory consumer names are unchanged.
• The per-partition loops are moved verbatim (only self.scratch_rows → self.base.scratch_rows, groups.drain() → self.partition_groups.drain()).
• RANK / DENSE_RANK emit now sorts (key, state) pairs directly instead of cloning every partition key into a separate sorted Vec and looking each one up again, as ROW_NUMBER already does.

Net: +271 / −365 lines in topk/mod.rs.

What is the testing strategy for this PR?

Covered by existing tests: the 111 topk unit tests (including the exact size() reconstruction test for DENSE_RANK and the reservation-vs-size() invariant check for ROW_NUMBER, updated only for the new field paths), the window_topn physical optimizer tests and the window Top-N sqllogictests.

Are there any user-facing changes?

No.

[Copilot]

Copilot review overview

🟡 Changes recommended

Release the reservation before emission and correct the shared-state documentation.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1

Open (1)

• Release reservation before emission to avoid memory pool exhaustion · New

What changed in this PR

Refactors the three partitioned Top-K operators to share common state and logic without intended behavior changes.

Changes:

• Adds PartitionedTopKBase for shared construction, grouping, emission, and memory accounting.
• Retains ranking-specific state and classification logic.
• Sorts partition key/state pairs directly and updates dense-rank tests.

File	Summary
datafusion/​physical-plan/​src/​topk/​mod.rs	Shared implementation and operator refactor. Review findings: line 1392—moderate (2 votes), release the reservation before emission; line 1258—nit (1 vote), correct documentation of the separate converters and scratch buffers.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 92.69663% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.66%. Comparing base (5741773) to head (d618fea).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/topk/mod.rs	92.69%	3 Missing and 10 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25534      +/-   ##
==========================================
- Coverage   82.66%   82.66%   -0.01%     
==========================================
  Files        1147     1147              
  Lines      446522   446470      -52     
  Branches   446522   446470      -52     
==========================================
- Hits       369134   369079      -55     
- Misses      54988    55001      +13     
+ Partials    22400    22390      -10     

☔ 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]

@jayzhan211 can we wait for more days. I wanted to fix perf issues first then do a final refactor to avoid any potential merge conflict.

[jayzhan211]

@jayzhan211 can we wait for more days. I wanted to fix perf issues first then do a final refactor to avoid any potential merge conflict.

Sure!

[kosiew]

@jayzhan211, thanks for working on this. The refactor looks good overall. Extracting the shared partitioned Top-K machinery into PartitionedTopKBase removes a fair amount of duplication while keeping the ranking-specific logic nicely separated. I only have one optional testing suggestion below.

[kosiew]

Nice refactor. One optional suggestion: could we extend test_partitioned_topk_output_batches_metric_counts_emitted_batches to assert the emitted row order as well? It already exercises multiple partitions, a small batch_size, coalescing, and the output-batch metrics through the shared emit path, so checking row order would round out the coverage here. Not blocking.