[Data] Add task-based shuffle v2 support for sort

[400Ping]

Description

This PR adds task-based shuffle v2 support for Dataset.sort().

Local benchmark

Benchmark:

release/nightly_tests/dataset/sort_benchmark.py

Environment:

OS: Ubuntu on a Microsoft hypervisor/WSL environment
CPU: Intel Core Ultra 7 265K
Physical cores: 20
Threads per core: 1
Memory: 62 GiB
Ray wheel: 2.58.0
Runs per configuration: 3

Input	Output partitions (legacy / v2)	Legacy median	v2 median	Runtime delta	Legacy object-store peak (median)	v2 object-store peak (median)	Memory delta	Legacy spill (median)	v2 spill (median)
5.96 GiB	64 / 200	6.040 s	8.891 s	+47.2%	12.149 GiB	5.983 GiB	-50.8%	0.000 GiB	0.000 GiB
11.92 GiB	128 / 200	30.634 s	16.927 s	-44.7%	15.727 GiB	11.944 GiB	-24.0%	25.997 GiB	0.000 GiB

Raw runtimes:

5.96 GiB legacy: 5.742 s, 6.040 s, 20.551 s
5.96 GiB v2:     8.891 s, 8.915 s, 8.352 s

11.92 GiB legacy: 28.371 s, 30.634 s, 68.221 s
11.92 GiB v2:     16.799 s, 16.974 s, 16.927 s

Additional information

[400Ping]

the microcheck error is unrelated to this PR

[gemini-code-assist]

Code Review

This pull request introduces the task-based shuffle v2 path for sorting, implementing SortShuffleMapOp with online range-boundary sampling and integrating it into the physical planner. It also updates ShuffleReduceOp to optionally preserve partition order and modifies the shuffle map task to handle compatible but non-identical schemas during concatenation. The review feedback highlights two important improvement opportunities: dynamically scaling the fallback tuple in _set_boundaries to support multi-column sort keys, and optimizing table concatenation in _shuffle_map_task by attempting a fast-path pa.concat_tables before falling back to type promotion.

[owenowenisme]

On second thought, I think we should sample all the blocks or we might get unexpected data skews. Could you change the pr to sample all blocks before doing range partitions?

[400Ping]

Hey @owenowenisme please take a look when you have time. Sorry for the late fix, I was at Beijing giving a talk in Community Over Code Asia.

[owenowenisme]

Have we thought about a separate sampling op instead of putting it inside of the same operator?

This could make sampling naturally support pipelining with upstream ops, get rid of a lot of things (e.g. max inflight sampling, _finish_map_inputs ) and fused with upstream map op if we implement some sampling op with the same map operator.

But this is just my hunch, can you dig into that direction and see if this is possible?

Also when you're doing this, could you survey how other frameworks sample input blocks before partitioning ?

[400Ping]

Have we thought about a separate sampling op instead of putting it inside of the same operator?

This could make sampling naturally support pipelining with upstream ops, get rid of a lot of things (e.g. max inflight sampling, _finish_map_inputs ) and fused with upstream map op if we implement some sampling op with the same map operator.

But this is just my hunch, can you dig into that direction and see if this is possible?

Also when you're doing this, could you survey how other frameworks sample input blocks before partitioning ?

Thanks, I dug into the separate sampling operator direction.

A separate zero-copy pass-through SortSamplingOp looks feasible. It could:

1. Retain each original RefBundle without copying its blocks.
2. Start a sampling task as each block arrives.
3. Collect the sample results and compute the global range boundaries.
4. Release the retained bundles only after the boundaries are ready.
5. Let the downstream SortShuffleMapOp receive those bundles through the
   normal executor admission path.

This would separate sampling from range partitioning, remove the custom
max-inflight sampling logic, and avoid replaying all buffered bundles into
ShuffleMapOp in one synchronous loop.

There are two limitations that I found:

First, the current MapOperator task protocol doesn't have a driver-side side
output for returning sampling metadata while forwarding the original block
references. A custom sampling operator can pipeline with upstream execution,
but fusing it into an upstream MapOperator without copying blocks would
require extending that protocol and the existing fusion rules.

Second, sort v1 calculates:

samples_per_block =
    num_partitions * SAMPLE_POINTS_PER_PARTITION / total_num_blocks

When the first input block arrives, the final number of input blocks may still
be unknown. Therefore, starting sampling immediately cannot exactly preserve
the current v1 sampling formula unless we introduce a mergeable per-block
sampling sketch or rely on an estimate.

This matches the general structure used by other frameworks:

• Spark runs a separate reservoir-sampling sketch before range partitioning,
  but it knows the number of RDD partitions in advance.
• Dask performs an initial approximate-quantile pass and then a second
  shuffle pass.
• Hadoop creates a partition file with InputSampler before running the
  total-order partitioner.

Ray Data cannot generally recompute an arbitrary upstream UDF safely, so the
sampling stage needs to retain the original object references while samples
are collected.

Would you be comfortable with changing the sampling implementation to a
bounded mergeable per-block sketch so sampling can start immediately, or do
you prefer preserving the exact v1 sampling budget even though that requires
waiting until the input block count is known?

For Ref:

• https://github.com/apache/spark/blob/master/core/src/main/scala/org/apache/spark/Partitioner.scala
• https://docs.dask.org/en/latest/generated/dask.dataframe.DataFrame.set_index.html
• https://apache.github.io/hadoop/api/org/apache/hadoop/mapreduce/lib/partition/InputSampler.html

[owenowenisme]

My take is choose the way that integrate with current the best, we don't have to follow everything v1 has, so even if you are only sampling a fixed 20 rows per block I think it's okay.

The current MapOperator task protocol doesn't have a driver-side side
output for returning sampling metadata while forwarding the original block
references.

About this I think we can have a custom op if map operator cannot do that. And we can disable fusing for this pr and figure ways to enable that as followup

[cursor]

Stale Bugbot comment from a previous run.

[400Ping]

The microcheck error is unrelated to this PR.

[400Ping]

@owenowenisme please take a look at this PR when you have time.

[owenowenisme]

Not blocker but why we have both this and ensure_boundaries?

[400Ping]

ShuffleMapOp needs a picklable partition_fn at construction, but sampled boundaries only exist after SortSamplingOp has seen every block. So the constructor gets a fail-loud placeholder and _ensure_boundaries swaps in the real one on the first input.

[cursor]

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

^{Reviewed by Cursor Bugbot for commit 5cede2f. Configure here.}