apache / apache/datafusion-ballista
Propagate input ordering across pass-through (non-repartitioning) shuffles
- Dominant language
- Rust
- Stars
- 2.1k
- Forks
- 320
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 66
Description
**Is your feature request related to a problem or challenge? Please describe what you are trying to do.**
`ShuffleReaderExec` always reports an empty output ordering, even when the shuffle it reads is a **pass-through** (no repartitioning). This forces downstream stages to re-sort data that is already sorted on disk.
Both sides of a shuffle rebuild their `PlanProperties` with a fresh, empty `EquivalenceProperties::new(schema)`:
- Writer: `ballista/core/src/execution_plans/shuffle_writer.rs:164`
- Reader (one-to-one, broadcast, coalesced): `ballista/core/src/execution_plans/shuffle_reader.rs:139,169,220`
So any `output_ordering()` the input plan had (e.g. from a `SortExec`, or the sorted output of a `SortMergeJoinExec`) is never serialized, never reconstructed, and never advertised to the child stage. When the child stage needs that order, DataFusion's planner cannot prove the input is sorted and re-inserts a `SortExec`. This redundant sorting is a real contributor to the shuffle+sort volume in TPC-H plans that use sort-merge joins.
**Describe the solution you'd like**
Propagate the input ordering across the shuffle boundary **for the pass-through case only** — i.e. when `shuffle_output_partitioning == None`, so the writer does not repartition.
In that case the ordering is preserved end-to-end and re-advertising it is provably sound with no execution-time cost:
- The writer streams each input partition straight to disk **in execution order** as a single file (`shuffle_writer.rs:221-268`, `write_stream_to_disk`), emitting one `ShuffleWritePartition { partition_id: input_partition }`.
- On the read side, output partition `p` therefore maps to **exactly one** upstream location — `self.partition[p]` is a single-element vec — so the per-partition randomized concatenation (`shuffle_reader.rs:385-392`) is a no-op and on-disk order is preserved.
Concretely:
1. Carry the input plan's `LexOrdering` from the writer to the reader (via `PartitionLocation` / the shuffle proto), gated on `shuffle_output_partitioning.is_none()`.
2. In the pass-through reader constructor, build `EquivalenceProperties::new(schema).with_reorder(ordering)` instead of the empty one, so `output_ordering()` reflects the preserved order.
This lets the child stage's planner skip a redundant `SortExec` whenever a pass-through shuffle sits below an order-dependent operator.
**Describe alternatives you've considered**
- **Also re-advertising ordering for hash / broadcast / coalesced shuffles.** Out of scope here, because those readers fan in `M` upstream runs per output partition and *concatenate* them — and `shuffle_reader.rs:392` deliberately randomizes the fetch order for load balancing. Sound re-advertising there would additionally require a sort-preserving k-way merge on read (replacing the concatenation) plus guaranteeing each run is individually sorted. That is a much larger change with an execution-time cost and a trade-off against fetch-load balancing; it should be tracked separately if desired.
- **Sort-based shuffle (`sort_shuffle/`)** does not help: it sorts by *output partition ID* (Spark-style, to cut file count), not by any user ordering, so it cannot carry an `O` ordering without folding `O` into its sort key.
- **Doing nothing** leaves redundant `SortExec`s in plans below pass-through shuffles.
**Additional context**
Scope is intentionally narrow (pass-through shuffles) to keep the change sound and zero-cost at execution time. Verified present on `upstream/main` at the time of filing.
Contributor guide
Research direction
Start with the referenced PlanProperties construction in ballista/core/src/execution_plans/shuffle_writer.rs:164 and shuffle_reader.rs:139,169,220, then trace PartitionLocation and the shuffle proto for ordering metadata. Verify the pass-through path preserves the input LexOrdering and that downstream planning no longer inserts a redundant SortExec, while non-pass-through shuffle paths remain unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- distributed-systems, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 64/100