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 only in the per-partition state and in how one partition's rows are classified, but every copy also repeats:

• the same 12 fields (schema, metrics, reservation, both row converters, both scratch Rows, the partition_groups scratch map, k, batch_size) and their doc comments,
• the same try_new body,
• the same first three steps of insert_batch (encode partition keys, group row indices by key, encode ORDER BY keys),
• the same emit skeleton (sort partition keys, coalesce, record output, wrap in a stream),
• the same size() skeleton.

That is roughly 150 lines per operator that must be kept in sync by hand, and a fix to one (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_and_group, finish_batch, emit and size helpers.
• Each operator becomes base + its own per-partition map (plus ob_runs / store for DENSE_RANK). Only the per-partition classification loop, the per-partition emit closure and the operator-specific size terms remain in each type.
• PartitionedTopKExec and the public/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).
• emit now sorts (key, state) pairs directly instead of cloning every partition key into a separate sorted Vec and looking each one up again.

Net: +241 / −459 lines in topk/mod.rs.

What is the testing strategy for this PR?

Covered by existing tests: the 90 topk unit tests (including the exact size() reconstruction test for DENSE_RANK, 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.12121% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.38%. Comparing base (d20936c) to head (7d29c1c).

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/topk/mod.rs	92.12%	4 Missing and 9 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25534      +/-   ##
==========================================
- Coverage   82.38%   82.38%   -0.01%     
==========================================
  Files        1138     1138              
  Lines      434491   434391     -100     
  Branches   434491   434391     -100     
==========================================
- Hits       357969   357873      -96     
- Misses      54875    54881       +6     
+ Partials    21647    21637      -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.