perf(core): keep small sort-shuffle output in memory instead of on disk - #2317
Closed
Dandandan wants to merge 7 commits into
Closed
perf(core): keep small sort-shuffle output in memory instead of on disk#2317Dandandan wants to merge 7 commits into
Dandandan wants to merge 7 commits into
Conversation
Adaptive query planning re-optimises the physical plan after every stage completion. Four defects in that path made the TPC-DS suite hang at q23 and return wrong results for two queries. InterleaveExec requires all children to share one Hash/Range partitioning and asserts on rebuild. AQE rewrites union branches independently, and TreeNode rebuilds a parent as soon as a child changes, so the assertion fired before any rule could see the parent and aborted the replan (apache#2047). NormalizeInterleaveRule rewrites interleaves to UnionExec at the head of each replan; EnforceDistribution re-forms them later through its can_interleave-guarded constructor. A replan can drop a stage that still has tasks in flight. Their completions hit an error that the task-status handler only logs, so the stage never advanced and the job hung forever. Late completions for superseded stages are now ignored. DataFusion's join_selection promotes a small build side to CollectLeft without restricting by join type. That is safe in one process, but Ballista runs one task per probe partition, each with a full copy of the build side, so a Left join emitted its unmatched rows once per task. TPC-DS q77 summed a 16x inflated column as a result. The static planner already guards this in maybe_promote_to_broadcast; DemoteUnsafeBroadcastJoinRule applies the same guard to AQE, between join_selection and EnsureRequirements. ExchangeExec inherited its input's equivalence properties, including constants that only hold within a partition. A UNION ALL branch projecting a literal marks that column constant, and keeping it across a hash repartition on the same column let EnforceSorting drop it from a sort whose SortPreservingMergeExec still merged on it, so q76 returned the wrong rows. RepartitionExec clears these for the same reason. Finally, LateCollectLeft broadcast a build side after both inputs had already been shuffled, paying for both. It is now restricted to null-aware anti joins, which require single-task CollectLeft. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
q31 and q71 were skipped because ORDER BY leaves the order of rows with equal sort keys unspecified, so two correct engines can emit tied rows in different orders and a positional diff is unstable. Compare positionally first, which keeps ordering under test, and only fall back to a canonical row-sorted comparison once that has failed. A wrong value still fails; a pure permutation passes and is reported as such. The skip list is now empty. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
AQE chooses joins and partition counts from measured runtime statistics rather than planning-time estimates. With the correctness fixes in place it verifies clean against single-process DataFusion across all 99 TPC-DS queries, and cuts the TPC-H SF10 suite from 62.2s to 29.6s (-52%). Physical-plan submissions now fall back to the static planner instead of erroring, since AQE plans from the logical plan and that path would otherwise break on the default configuration. Adds SessionConfigExt::with_ballista_adaptive_query_planner so the static planner can be selected in code. Scheduler tests that assert the static planner's stage and partition layout pin it explicitly. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
AQE wrapped both inputs of every join in a fresh exchange before resolving it, because repartitioning is what gives the resolver measured statistics. An equi-join already leaves its output partitioned on both sides of each join key, so a downstream join on the other key is co-partitioned and has nothing to wait for. TPC-H q18 moved lineitem twice as a result, 166.5M rows shuffled against the static planner's 106.5M, and shuffle-write CPU doubled from 12.3s to 24.4s. When both inputs already satisfy the required distribution, treat the join as repartitioned and resolve it in the same pass. Partitioning::satisfaction consults equivalence classes, so it sees that a join on l_orderkey = o_orderkey leaves the output partitioned on both. Skipping the exchange without resolving in the same pass does not work: the exchange is also the stage boundary that keeps an unresolved DynamicJoinSelectionExec out of a dispatched stage, and the stage plan then fails to serialize. TPC-H SF10 drops from 29.6s to 23.8s (-20%), q18 from 8.0s to 3.5s (-56%), and TPC-DS still verifies 99/99 against single-process DataFusion. The EXPLAIN tests assert the static planner's stage layout, so they now pin it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…default Two configuration defaults change, both aimed at the per-stage overhead that dominates short distributed queries. `ballista.scheduler.max_partitions_per_task` moves from `1` to `0` (unbounded). At `1` every input partition became its own task, so a 16-partition stage dispatched 16 tasks, each carrying its own restricted plan through protobuf encode, dispatch and decode. At `0` the scheduler fills each task up to the assigned executor's free vcore count instead, and operators such as sort and hash join can work across partitions inside one plan invocation. `ballista.optimizer.broadcast_join_threshold_bytes` moves from 10 MB to 128 MB, promoting more joins to a broadcast `CollectLeft` join. A broadcast join removes a shuffle outright: no hash-repartition, no shuffle files written or read back. The cost is that the build side is replicated into every concurrent probe task, so peak memory per executor scales with the number of tasks it runs at once. Measured on TPC-H SF10 (2 executors x 4 vcores, target_partitions=16), running each configuration once per round with the order rotated between rounds so drift is not attributed to a config. Against that baseline the two settings together were the best combination, and they compose — they attack different costs. Enabling the coalesce rule or forcing hash joins on top added little. Both defaults are documented in the 55.0.0 upgrade guide along with how to restore the previous behavior. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…nput
A task that owns several input partitions wrote one `data.arrow` plus index
per input partition. Every downstream reader fetching partition k therefore
opened one file per input partition, and a stage left M files behind where M
is the stage's input partition count.
The writer now emits a single file per task. Each input partition still
buckets and spills concurrently, and — importantly — still encodes its own
buckets to IPC bytes on its own task, so the interleave, framing and
compression stay parallel. The coordinator, which already awaited every input
before responding, then concatenates the finished buffers into one file and
writes one index.
The file is laid out partition-major: the schema header, then output
partition 0's bytes from every input in turn, then partition 1's, and so on.
Keeping each output partition contiguous is what lets the index stay one
offset per partition, so the reader is unchanged — `create_shuffle_path`
already resolves a sort-shuffle summary to `{stage_id}/{file_id}/data.arrow`,
and `MultiStreamPartitionStream` already crosses concatenated IPC streams
inside a byte range.
An earlier revision of this change did the encoding in the coordinator, which
collapsed write parallelism from P to 1 and cost 29% on TPC-H SF10. Encoding
per input is what makes the file-count reduction free.
Spill directories move from `{stage}/{file_id}/spill` to
`{stage}/{task_id}/spill-{input}`, so a task owns exactly one directory under
the stage and cleanup no longer strands an empty directory per input.
TPC-H SF10, 2 executors x 4 vcores, `--partitions 16`,
`max_partitions_per_task=0`. Old and new binaries alternate within each round
so drift is shared, 8 runs each:
| | median | sort-shuffle files |
|---|---|---|
| one file per input partition | 21.37 s | 1762 |
| one file per task | 20.05 s | 442 |
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Writing shuffle bytes to the work directory and reading them back is the dominant cost of a distributed query here. Pointing --work-dir at a RAM disk, changing nothing else, takes TPC-H SF10 from 23.27s to 13.81s (-40.6%, medians of 5 alternating rounds) — essentially single-process DataFusion's 13.30s. This adds an executor-wide store that holds sort-shuffle output in memory and skips the file write. It composes with the per-task consolidation: the writer already holds each partition's finished IPC bytes, and those bytes are byte-identical to the range the on-disk reader would address through the index, so storing them is a matter of not writing. The budget is enforced at admission, not by eviction: a task that skipped its write has nowhere else to serve from, so an entry that did not fit would lose data a downstream stage still needs. A task whose output does not fit writes to disk exactly as before, which makes the store a fast path rather than a new failure mode. Entries are released when a job's data is reclaimed, which for intermediate stages already happens immediately on job completion. All three read paths consult the store before building a path: the co-located local read, the Flight do_get, and the IO_BLOCK_TRANSPORT do_action used for remote reads by default. The block path prepends the schema-header stream the way the on-disk path does, so the receiver cannot tell them apart, which is why the store carries the header bytes. `ballista.shuffle.memory_store_limit_bytes` bounds the executor, 1 GiB by default; 0 disables the store. Measured on TPC-H SF10, 2 executors x 4 vcores, max_partitions_per_task=0, store on vs off within alternating rounds: | budget | median | sort-shuffle files left | |---|---|---| | off | 22.12 s | 443 | | 1 GiB (default) | 22.54 s | ~270 | | 12 GiB | 20.98 s | 16 | So the store is currently neutral at its default and worth ~5% when sized to hold everything — well short of the 40% the RAM disk shows. The gap is coverage: only sort-shuffle goes through this path. The passthrough ShuffleWriterExec still writes ~855 files per run, including every query's final result, and those bytes were in RAM in the RAM-disk experiment. Covering that writer is the follow-up that would realise the rest. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
the
documentation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The measurement that motivates this
Pointing
--work-dirat a RAM disk, identical binaries and nothing elsechanged, on TPC-H SF10:
-40.6%, with non-overlapping distributions. Single-process DataFusion on the
same data is 13.30 s.
So shuffle disk I/O is essentially the whole remaining gap to single-node. It
is not the scheduler (planning is 1-2 ms per query, dead time 10-20 ms), not
join strategy, and not the IPC format — the RAM disk still pays full IPC
encode/decode and lz4, and only file I/O is removed. It also explains why
disabling lz4 makes things 2.2x worse: bytes written is the whole game.
What this adds
An executor-wide store holding sort-shuffle output in memory, skipping the
file write. It composes with #2316: after that change the writer already holds
each partition's finished IPC bytes, and those bytes are byte-identical to the
range the on-disk reader addresses through the index, so storing them is a
matter of not writing.
Admission, not eviction. A task that skipped its write has nowhere else to
serve from, so evicting an entry would lose data a downstream stage still
needs. The budget is checked when the entry is offered; a task that does not
fit writes to disk exactly as before. The store is a fast path, never a new
failure mode.
All three read paths consult the store before building a path: the
co-located local read, the Flight
do_get, and theIO_BLOCK_TRANSPORTdo_actionthat serves remote reads by default. The block path prepends theschema-header stream the way the on-disk path does, so the receiver cannot
tell them apart — which is why the store carries the header bytes.
Release happens when a job's data is reclaimed. Intermediate stages
already get that immediately on job completion via
clean_up_intermediate_job_data, rather than waiting for the delayedwhole-job cleanup.
ballista.shuffle.memory_store_limit_bytesbounds the executor, 1 GiB bydefault;
0disables it.Results
TPC-H SF10, 2 executors x 4 vcores,
max_partitions_per_task=0, store on vsoff within alternating rounds:
Neutral at the default, ~5% when sized to hold everything. That is well short
of the 40% the RAM disk shows, and the reason is coverage: only
sort-shuffle goes through this path. The passthrough
ShuffleWriterExecstillwrites ~855 files per run — including every query's final result, which makes
a full disk round trip even for a one-row answer — and those bytes were in RAM
in the RAM-disk experiment.
Extending the store to that writer is the follow-up that would realise the
rest, and is the reason this is a draft rather than a merge candidate on its
own numbers.
Testing
cargo test -p ballista-core --lib— 287 pass, including 7 new store tests(admission, budget exhaustion, per-job and per-stage release, disabled via
0, shared global).small_shuffle_is_kept_in_memory_and_readabledrives a two-input task andasserts the job directory is empty — the write was skipped, not moved —
then reads every row back out of the store.
memory_store_limit_bytes = 0, sincethey assert on files the memory path deliberately does not create.
DataFusion with
--verify, with the store on.