[viiccwen]

Description

This PR adds a Top-N execution path for eligible Ray Data sort(...).limit(k) plans. With DataContext.shuffle_strategy = ShuffleStrategy.SHUFFLE_V2 and the private top_n_rewrite_enabled config enabled, eligible plans are rewritten as:

LocalTopN(limit=k) -> TopNMerge(limit=k)

• TopNRewriteRule rewrites Limit(k, Sort(...)) to a logical TopN. It runs after LimitPushdownRule and CombineShuffles, so it also handles eligible plans where row-preserving operators originally separate Sort and Limit, and collapses consecutive limits to the smaller value. Sorts with user-provided boundaries are left unchanged.
• LocalTopN reuses MapOperator. It sorts each input block and retains at most k rows before the final merge, avoiding sort sampling and the full range shuffle. The map is not fused into the final merge.
• TopNMergeOperator collects the locally sorted candidates, normalizes mixed Arrow and Pandas block types, merges the candidates, and returns one globally sorted Top-N block. Empty Arrow and Pandas inputs retain their schema.
• The final single-task merge uses actual candidate metadata to enforce configurable block, row, and byte limits. Missing metadata or candidates above a limit fail before the merge task is submitted. Candidate counts are also reported through operator extra metrics.
• Sort remote arguments and the execution label selector are propagated to the final merge task. Owned candidate bundles are released on successful, empty, and failing paths.

This path is private and default-off. It only supports SHUFFLE_V2; legacy sort shuffle is unchanged.

Benchmark

The workload is range -> map_batches(randomized key) -> sort(["key", "row_id"]) -> limit(k) -> materialize.

Input rows	Requested blocks	Limit	Row layout	Baseline median	Top-N median	Speedup	Final candidates
1,000,000	32	100	2x int64	1.843 s	0.317 s	5.81x	3,200 rows / 0.05 MiB
1,000,000	128	100	2x int64	3.577 s	1.320 s	2.71x	12,864 rows / 0.20 MiB
5,000,000	128	100	2x int64	3.529 s	1.162 s	3.04x	12,864 rows / 0.20 MiB
1,000,000	32	10,000	2x int64	1.977 s	0.344 s	5.75x	320,000 rows / 4.88 MiB
1,000,000	32	100	2x int64 + 128B payload	2.114 s	0.386 s	5.47x	3,200 rows / 0.45 MiB

Related Issues

Related to #66831.

Additional information

Follow-up work (Important)

• Replace the blocking AllToAllOperator wrapper with a native asynchronous PhysicalOperator lifecycle, including task tracking and resource accounting.
• Add an exact backend-level bounded-selection primitive and reuse it for local and final Top-N selection.
• Implement a semantics-preserving overflow path such as bounded hierarchical merge or candidate V2 sort before considering default-on rollout.
• Run release-sized benchmarks to establish selection thresholds, candidate budgets, and rollout criteria.

[gemini-code-assist]

Code Review

This pull request introduces a new TopN logical operator and an associated rewrite rule to optimize Sort + Limit operations into a more efficient TopN execution path. The implementation includes a local candidate reduction phase followed by a global merge, supported by a configurable budget to manage resource usage. My feedback suggests refactoring the rewrite function in top_n_rewrite.py to check for non-positive limits earlier, which improves code clarity and prevents invalid TopN operator creation.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 230bff4. Configure here.}

[owenowenisme]

High level feedbacks:

1. I'd suggest implementing a new operator based on PhysicalOperator (maybe on top of MapOperator, which might save a lot of work), since AllToAllOperator will be deprecated soon, once all our shuffles move to shuffle v2 and v1 is deprecated.
2. For top_n_rewrite_max_limit, I suggest we survey how other engines set this number, or sweep it ourselves. I'd also suggest defaulting top_n_rewrite_enabled to true.
3. For the final TopN merge, we could probably fetch the inputs in parts and maintain a top-N heap, instead of fetching all at once. With a native operator (1), this could even be done incrementally as candidates arrive.
4. Check pyarrow.compute.select_k_unstable to avoid sorting the whole block (though we probably need to check what problems the unstable ordering could cause).
5. I think LocalTopN should be able to fuse with upstream.

[rayo5o]

I ran some benchmarks on this branch to get memory numbers and to see how k and the number of blocks affect it. Setup: one node, 32 CPUs, Ray nightly a705161 + this PR. Both sides use SHUFFLE_V2 and only top_n_rewrite_enabled is toggled. Same pipeline as the PR description. Each number is the median of 5 runs, each run is a fresh Ray instance with a warm-up job first. Outputs matched in every run and nothing spilled.

rows	blocks	k	old	new	extra RSS old / new
1M	128	100	8.00s	1.39s	1.5 GiB / 20 MiB
10M	128	100	8.89s	1.89s	1.5 GiB / 236 MiB
100M	128	100	17.27s	4.19s	2.3 GiB / 1.0 GiB
100M (+128B payload)	128	100	21.60s	7.15s	15.9 GiB / 3.8 GiB
10M	128	10,000	8.06s	1.97s	1.1 GiB / 281 MiB
10M	128	100,000	8.28s	2.51s	1.1 GiB / 402 MiB
10M	2048	1,000	61.18s	9.26s	5.6 GiB / 284 MiB

A few things I noticed:

1. Small inputs with few blocks: for 1M rows / 32 blocks / k=100, with a warm-up job first I get old 0.78s vs new 1.08s. Without warm-up it is 3.32s vs 1.86s. So the 5.8x in the description is mostly worker startup, I think. Larger block counts show a real gap regardless of warm-up.
2. With the default configs, the merge budget can fail the query instead of falling back. k=10000 (the default max_limit) with 256 input blocks gives 2,560,128 candidate rows and raises ValueError. k=1000 with 2048 blocks does the same. Since max_limit is checked at plan time and the budget at run time, a k that passes the rewrite check can still fail. If this goes default-on, falling back to the regular sort might be safer.
3. On max_limit (@owenowenisme's point 2): speedup stayed around 4x up to k=30000, and k=100000 (every row becomes a candidate here) was still 3.3x. That's single node only, so the gather cost across nodes is not covered.
4. In ds.stats(), TopNMerge always reports about 3us remote wall time, even when merging ~2M candidate rows. It looks like the final merge task's time isn't recorded.

[viiccwen]

High level feedbacks:

1. I'd suggest implementing a new operator based on PhysicalOperator (maybe on top of MapOperator, which might save a lot of work), since AllToAllOperator will be deprecated soon, once all our shuffles move to shuffle v2 and v1 is deprecated.
2. For top_n_rewrite_max_limit, I suggest we survey how other engines set this number, or sweep it ourselves. I'd also suggest defaulting top_n_rewrite_enabled to true.
3. For the final TopN merge, we could probably fetch the inputs in parts and maintain a top-N heap, instead of fetching all at once. With a native operator (1), this could even be done incrementally as candidates arrive.
4. Check pyarrow.compute.select_k_unstable to avoid sorting the whole block (though we probably need to check what problems the unstable ordering could cause).
5. I think LocalTopN should be able to fuse with upstream.

Hello @owenowenisme. Sorry for lately replying and appreciate ur suggestion!

1. Yes, I already add TODO in relative code diff. Since the phase1 over 1200+ change lines, would you prefer having this included in the current PR, or keeping it separate to make reviewing easier?
2. After looking into other engines:

• Spark: Uses a configurable fixed threshold (topKSortFallbackThreshold), with a very large default.
• DuckDB: Uses a cardinality-aware heuristic, considering both the limit (5,000) and its ratio to estimated input rows (0.7%).
• DataFusion: Focuses on bounded Top-K execution and runtime pruning rather than relying solely on a fixed threshold.

Our benchmarks also show that Top-N remains beneficial at k=100,000. Since Ray's candidate volume depends on k, block count, and row width, I think a candidate-aware budget would be more appropriate than a fixed threshold alone. I'll run additional end-to-end benchmarks before deciding on max_limit and enabling the rewrite by default.

3. Great suggestion! I'll incorporate incremental candidate consumption with bounded memory usage into the plan.
4. I've also planned this as a follow-up PR. Given the current diff size, would you prefer including this in the current PR as well?
5. Agreed! I'll update the implementation to allow LocalTopN to fuse with upstream operators.

Thanks again! Let me know what you think about the PR scope, especially points 1 and 4.

Edit: After discussing, we should put point 1, 3, 4, 5 in this PR.