[andygrove]

Which issue does this PR close?

Closes #2680.

Rationale for this change

The static planner kept only hash repartitions as stage boundaries and dropped every other RepartitionExec. Dropping a round-robin repartition only costs parallelism, but a range repartition is part of what a plan means. A plan that range-partitions and then sorts each partition (a distributed global sort) came back as partitions that were each sorted but not range-disjoint, with no error. Comet Native Executors need this for Spark's range exchanges (apache/datafusion-comet#7288).

What changes are included in this PR?

• Planner. A RepartitionExec with Partitioning::Range is a stage boundary, like a hash one, so partition i of the consuming stage reads exactly range i. Round-robin repartitions are still dropped. A range repartition that must preserve its input order (preserve_order()) returns an error: a shuffle cannot merge sorted inputs, and planning it anyway would return partitions that are range-disjoint but no longer sorted.
• Writer. SortShuffleWriterExec accepts Partitioning::Range and routes rows with DataFusion's RangeExpr, which shares its routing function with the BatchPartitioner behind RepartitionExec, so the shuffle puts every row where the in-process operator would have. Everything after routing (buffering, spilling, the consolidated file and its index) is shared with the hash shuffle.
• Per-task restriction. A ShuffleReaderExec with a range partitioning slices its split points along with its partitions, replacing the TODO and the error in task_builder.rs. A slice that does not ascend (a retried partition queued ahead of a later one) or is empty reports UnknownPartitioning rather than failing the task: no split points describe it, and UnknownPartitioning claims nothing about which rows each partition holds. (DataFusion's partitioned hash-join dynamic filters do read the probe side's reported partitioning, but Ballista turns dynamic filter pushdown off by default.)
• Wire format. A new partitioning field on SortShuffleWriterExecNode carries the range partitioning. Hash stages keep using output_partitioning, so a hash plan is encoded as before. UnresolvedShuffleExec and ShuffleReaderExec already carried a general Partitioning.

Two things this PR deliberately does not do:

• It does not use RangeShuffleWriterExec / RangeShuffleReaderExec, which the issue mentions. RangeShuffleWriterExec requires an ordered producer (try_new rejects a child with no output ordering) and its reader merges sorted runs, while the issue's own example sorts after the repartition. The range edge reuses the sort-shuffle writer, as the hash edge does.
• It does not change the adaptive planner's own rules or the pipelined exchange operators. Since feat: run submitted physical plans through the adaptive planner #2682, submitted physical plans go through the adaptive planner by default, and its PinSubmittedShufflesRule already pins a range repartition as an exchange. On main that job fails with unsupported shuffle output partitioning, because only hash stages get a writer. The writer change here makes it run, so the adaptive path is fixed too. The pipelined reader still refuses to restrict a range partitioning, and nothing plans pipelined jobs yet.

RangeExpr extracts a ScalarValue per row, so range routing will be slower than hash routing on large inputs (the same cost as DataFusion's own RepartitionExec). I have not benchmarked it.

Tests. Besides unit tests for routing (with BatchPartitioner as an oracle), serde and reader restriction, planner/range_shuffle_tests.rs plans a range-partition-then-sort and runs the stages for real, without sockets: it restricts each task's plan, ships it through the codec, and runs the writer and reader. It checks that the final partitions are range-disjoint, sorted, and hold every row exactly once, with one task per partition, with tasks owning several partitions, and with a slice out of key order. With the planner change reverted, those tests fail with the symptom in the issue. Two tests in ballista/client/tests/physical_plan_submission.rs submit the same shape to a standalone cluster with max_partitions_per_task = 1, one per planner. The job's output partitions come back in the order their tasks finished, so they check every key exactly once, each partition sorted, and no two partitions' key ranges overlapping. The static one fails with overlapping partitions when the planner change is reverted, and the adaptive one fails with unsupported shuffle output partitioning when the writer change is reverted.

Other commits. Besides the fix, three commits add entries to the 56.0.0 upgrade guide, move the new proto field to number 11 (see the overlap with #2723 below), and simplify the routing code and trim the tests and comments after a review pass. The last commit adds the standalone-cluster tests.

Overlap with other open PRs.

• feat: let custom ExecutionPlan nodes restrict themselves to a task's partitions #2721 gives restrict_plan_to_partitions a third resolver argument, and both PRs append tests to task_builder.rs. Whichever lands second has a conflict in that test module, and the calls to restrict_plan_to_partitions in this PR's tests need None added.
• feat: let hash shuffles place rows by a client-computed partition id #2723 also changes how SortShuffleWriterExec routes rows, with a ShuffleRouting enum in place of the writer's Partitioning field. The two conflict in the writer, serde, the planner and the proto, and whichever lands second has to be reworked onto the other's routing; a range partitioning would fit as a third ShuffleRouting variant. This PR numbers its new proto field 11 because feat: let hash shuffles place rows by a client-computed partition id #2723 uses 10 on the same message.

Are there any user-facing changes?

No SQL-visible changes. A physical plan submitted with a range RepartitionExec now returns correct results instead of range-overlapping partitions, and a range repartition that must preserve its order is rejected with an error.

The generated SortShuffleWriterExecNode struct gains a public partitioning field, so code that builds it with a struct literal needs partitioning: None. The only such use in this repo is updated. Both changes are described in the 56.0.0 upgrade guide, and the PR carries the api-change label for the struct field.

[coderfender]

Looks good. I like that the routing test checks against DataFusion's BatchPartitioner. Since the end-to-end tests run without sockets, could we add one standalone-cluster test in ballista/client/tests that submits a range-repartition plan with a split stage (e.g. max_partitions_per_task = 1) and checks the output is globally sorted? A rough perf number for range vs hash routing would be nice too, but not blocking.

[andygrove]

@coderfender Added in fa5fe79, under both the static and the adaptive planner, with max_partitions_per_task = 1. One wrinkle: the job's output partitions come back in the order their tasks finished, so the concatenated stream isn't something we can check for global order. The test checks what makes it a global sort instead: every key exactly once, each partition sorted, and no two partitions' key ranges overlapping. With the planner change reverted, the static test fails with overlapping partitions.

Writing it turned up something useful. Since #2682 landed, submitted plans go through the adaptive planner by default, and it pins a range repartition as an exchange. On main that job fails with unsupported shuffle output partitioning, and this PR's writer change is what makes it work, so the adaptive test covers that path too. I've rebased onto main as well. No perf numbers yet, I'll leave that for a follow-up.

[avantgardnerio]

This seems complimentary to the existing AQE range partitioning, and necessary as you say since DataFusion plans now have Partitioning::Range, which Ballista needs to support.

[andygrove]

Good question. Ballista never computes the split points for Comet. Spark does, in the driver, before the plan reaches Ballista.

When Spark builds the shuffle for a range exchange (ShuffleExchangeExec.prepareShuffleDependency), it constructs a RangePartitioner. That constructor runs a separate small Spark job over the child. The job reservoir-samples the sort keys, sorts the sample, and picks numPartitions - 1 boundary rows, which are the split points. They come from sampling the data and not from static statistics. Comet's native shuffle already does this: CometShuffleExchangeExec builds the same RangePartitioner over the sampled sort keys and reads rangeBounds out of it. They go into the boundary_rows field of Comet's range partitioning proto.

Ballista is not called repeatedly. The sampling is a separate Spark job, and the Ballista job is submitted once.

Comet Native Executors doesn't hand Ballista a range partitioning yet. Today the driver sends none of that, and the cluster side turns a range exchange into a coalesce, so the consuming stage reads it as one partition. That sorts correctly but runs in a single task. This PR is what lets us stop doing that, because the planner no longer drops Partitioning::Range. Getting the bounds into the plan is the driver's job. It would run the sampling job and put the result into the RepartitionExec's Partitioning::Range expressions. We haven't built that part, so I can't say yet how the sampling job gets its rows.

This also means Ballista only needs to route against fixed split points here. It doesn't need RuntimeStatsExec for this path. I think that's why it complements your AQE work. Your AQE path derives the split points from runtime stats. This path accepts a plan whose split points are already decided.