Skip to content

minor: share common state between the three PartitionedTopK operators - #25534

Draft
jayzhan211 wants to merge 1 commit into
apache:mainfrom
jayzhan211:simplify/partitioned-topk-base
Draft

jayzhan211 wants to merge 1 commit into
apache:mainfrom
jayzhan211:simplify/partitioned-topk-base

Conversation

@jayzhan211

Copy link
Copy Markdown
Contributor

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_rowsself.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.

@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Sep 20, 2026
@jayzhan211
jayzhan211 marked this pull request as ready for review September 20, 2026 08:45
@jayzhan211
jayzhan211 requested review from kosiew and a lite review from Copilot September 20, 2026 08:46

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 Medium severity

Open (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 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.

Comment thread datafusion/physical-plan/src/topk/mod.rs
@codecov-commenter

Copy link
Copy Markdown

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

Copy link
Copy Markdown
Contributor

@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

Copy link
Copy Markdown
Contributor Author

@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!

@jayzhan211
jayzhan211 marked this pull request as draft September 20, 2026 12:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants