[viiccwen]

Description

Add an opt-in Top-N path for Ray Data sort(...).limit(k) under SHUFFLE_V2. It selects local candidates and merges them incrementally to avoid the full range shuffle.

Before: Sort samples boundaries, sorts and partitions all input rows, then merges partitions. Limit can stop consumption before all reducers finish.

After: TopNRewriteRule replaces eligible Sort + Limit plans. Each block contributes at most k candidates. One active merge task maintains the running Top-N. The operator emits the sorted result after input completion.

flowchart TD
    subgraph Before
        subgraph BE["Execution / shuffle: V2 sort"]
            B1[SortSamplingOp] --> B2[SortShuffleMapOp]
            B2 --> B3[ShuffleReduceOp] --> B4[LimitOperator]
        end
    end
    subgraph After
        subgraph AE["Execution: Top-N operators"]
            A1[LocalTopN] --> A2["TopNMergeOperator: batch + running Top-N"]
            A2 -->|"next batch"| A2
            A2 --> A3["Sorted result after input completion"]
        end
    end
    classDef removed fill:#ffcdd2,stroke:#c62828,color:#000
    classDef added fill:#c8e6c9,stroke:#2e7d32,color:#000
    class B1,B2,B3,B4 removed
    class A1,A2,A3 added

Loading

Arrow uses select_k_unstable where supported. The fallback sorts keys and gathers selected rows. Pandas preserves its sort semantics. The native operator uses DataOpTask for resource accounting, backpressure, and input cleanup.

Batch thresholds are 32 blocks or 128 MiB. Partial batches can flush after 100 ms. Admission can exceed a threshold by one bundle. These are not hard memory limits; the running Top-N must fit in one worker.

The feature stays private and default-off. Eligibility requires positive k, no custom boundaries, and a configurable maximum k of 10,000 by default. LocalTopN stays unfused. Boundary ties can select different row identities.

Benchmark

Workload: range -> map_batches(SplitMix64 key) -> sort(["key", "row_id"]) -> limit(k) -> materialize.

Narrow rows contain two int64 columns. Wide rows add a 128-byte payload. Block counts are requested counts; 128 produces 129 candidate blocks.

Rows / blocks / k	Layout	Time (s)	Speedup	Worker USS peak (MiB)	Object-store peak (MiB)	Memory reduction (USS / store)
1M / 32 / 100	narrow	0.518 / 0.592	0.87x	1004.9 / 867.7	27.5 / 12.6	13.6% / 54.3%
1M / 128 / 100	narrow	3.838 / 1.429	2.69x	1778.4 / 866.1	18.9 / 11.8	51.3% / 37.7%
5M / 128 / 100	narrow	4.949 / 1.602	3.09x	1744.1 / 877.0	136.3 / 55.4	49.7% / 59.3%
1M / 32 / 10,000	narrow	0.520 / 0.698	0.74x	1040.3 / 869.0	27.6 / 19.5	16.5% / 29.2%
1M / 32 / 100	wide	0.648 / 0.760	0.85x	1172.8 / 966.6	161.3 / 141.8	17.6% / 12.1%

Related issue

closes #66831.

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

Stale Bugbot comment from a previous run.

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

[cursor]

Stale Bugbot comment from a previous run.

[owenowenisme]

Hey @viiccwen, I made a few changes to your pr:

1. Revert the operator fusion in this pr for now, as I am thinking to merge different map operators on physical layer instead of logical layer and right now it's kinda hacky to fuse ops that aren't "AbstractMap" but physical map. If we somehow fuse on physical layer this should be support automatically.
2. Support the streaming top N operator so that the logic of checking candidate budget can be easier, and also this should be something we wanna support eventually.

[cursor]

Stale Bugbot comment from a previous run.

[rayo5o]

Re-ran my earlier benchmarks on e705cd6 (nightly ed1d83a + this PR, same 32-CPU machine, both paths on the same base, median of 5 warmed-up runs). Point 2 from my last comment no longer applies: with default config, k=10000 over 256 blocks (2.56M candidates) and k=1000 over 2048 blocks both finish and return the right rows. Point 4 is fixed too, TopNMerge now reports real time (0.88s at k=100000).

The numbers got better. At 128+ blocks the speedup is now 4-8x (was 3-7x). For 100M rows with a 128-byte payload, sort+limit takes 21.4s and 16.3 GiB of extra worker memory; Top-N now takes 3.5s and 0.5 GiB (was 7.2s and 3.9 GiB on the earlier version). The 1M rows / 32 blocks case is now 0.88s vs 0.80s for sort+limit, so the earlier slowdown there is mostly gone.

One thing that moved the other way: in the 100M-row case the object store peak for Top-N went from 6.2 GiB to 9.0 GiB (7.6-9.2 GiB across runs; sort+limit stays around 14 GiB). My guess is it comes from disabling upstream backpressure until all candidates arrive, but I haven't checked. Still single-node only.

[viiccwen]

Re-ran my earlier benchmarks on e705cd6 (nightly ed1d83a + this PR, same 32-CPU machine, both paths on the same base, median of 5 warmed-up runs). Point 2 from my last comment no longer applies: with default config, k=10000 over 256 blocks (2.56M candidates) and k=1000 over 2048 blocks both finish and return the right rows. Point 4 is fixed too, TopNMerge now reports real time (0.88s at k=100000).

The numbers got better. At 128+ blocks the speedup is now 4-8x (was 3-7x). For 100M rows with a 128-byte payload, sort+limit takes 21.4s and 16.3 GiB of extra worker memory; Top-N now takes 3.5s and 0.5 GiB (was 7.2s and 3.9 GiB on the earlier version). The 1M rows / 32 blocks case is now 0.88s vs 0.80s for sort+limit, so the earlier slowdown there is mostly gone.

One thing that moved the other way: in the 100M-row case the object store peak for Top-N went from 6.2 GiB to 9.0 GiB (7.6-9.2 GiB across runs; sort+limit stays around 14 GiB). My guess is it comes from disabling upstream backpressure until all candidates arrive, but I haven't checked. Still single-node only.

Thanks for the detailed benchmarks!

This PR is still a work in progress, and I'm making changes to the implementation. It might be best to hold off on further benchmarking until it's ready for review, to avoid spending time testing intermediate versions.

Appreciate your help!

[cursor]

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

^{Reviewed by Cursor Bugbot for commit 6d3a166. Configure here.}

[cursor]

Bundle threshold can stall LocalTopN

Medium Severity

When object_store_memory is below the LocalTopN bundle target, min_bundle_bytes is set to that entire limit. EstimateBytes will not submit until the input queue holds that much data, but those queued bundles already count against the same budget, so LocalTopN can sit below the threshold, refuse more input, and never flush until upstream is fully consumed. Tight limits such as the existing 1024-byte merge test can hang a finite sort().limit() job.

 

^{Reviewed by Cursor Bugbot for commit 6d3a166. Configure here.}