andygrove opened a new pull request, #2722: URL: https://github.com/apache/datafusion-ballista/pull/2722
# 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 own `RangeExpr`, the function `RepartitionExec` applies, 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 nothing consumes a restricted reader's partitioning beyond its count. - **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 or the pipelined exchange operators. Submitted physical plans use the static planner today. #2682 would send them through the adaptive planner, whose exchange rule only turns hash repartitions into exchanges, so it will need the same treatment for ranges. 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. Not run locally: the tests that bind a TCP port (the client integration tests and a few in `ballista-core`); CI covers those. **Second commit.** `fix(scheduler): import ExecutionPlan in the execution graph tests` is unrelated to this change. #2613 deleted a block of test imports that included `ExecutionPlan`, which a test added by #2633 still uses. Each passed CI alone, but the scheduler's tests have not compiled on `main` since both landed (`cargo check`, Clippy and the test jobs are red there). It is a one-line fix, included so this PR can get a green run. Happy to drop it if a fix lands first. # 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
