apache / apache/datafusion-ballista

Propagate input ordering across pass-through (non-repartitioning) shuffles

Open
#2,006 0 comments 0 reactions 0 assignees View on GitHub
enhancement performance
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.