minor: share common state between the three PartitionedTopK operators - #25534
jayzhan211 wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
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
What changed in this PR
Refactors the three partitioned Top-K operators to share common state and logic without intended behavior changes.
Changes:
- Adds
PartitionedTopKBasefor 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 Report❌ Patch coverage is
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. 🚀 New features to boost your workflow:
|
|
@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! |

Which issue does this PR close?
Rationale for this change
PartitionedTopK(ROW_NUMBER),PartitionedTopKRank(RANK) andPartitionedTopKDenseRank(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:Rows, thepartition_groupsscratch map,k,batch_size) and their doc comments,try_newbody,insert_batch(encode partition keys, group row indices by key, encode ORDER BY keys),emitskeleton (sort partition keys, coalesce, record output, wrap in a stream),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.
PartitionedTopKBaseholding the shared fields, withtry_new,encode_and_group,finish_batch,emitandsizehelpers.base+ its own per-partition map (plusob_runs/storeforDENSE_RANK). Only the per-partition classification loop, the per-partition emit closure and the operator-specific size terms remain in each type.PartitionedTopKExecand the public/crate-visible signatures of the three types are unchanged. Memory consumer names are unchanged.self.scratch_rows→self.base.scratch_rows).emitnow sorts(key, state)pairs directly instead of cloning every partition key into a separate sortedVecand 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
topkunit tests (including the exactsize()reconstruction test forDENSE_RANK, updated only for the new field paths), thewindow_topnphysical optimizer tests and the window Top-N sqllogictests.Are there any user-facing changes?
No.