apache / apache/datafusion-ballista

Buffer shuffle output in memory and spill to disk only when a configured budget is exceeded

Open
#2,318 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.**

Ballista always materializes shuffle output to the filesystem. `ShuffleWriterExec` and `SortShuffleWriterExec` stream each output partition to an Arrow IPC file under `{work_dir}/{job_id}/{stage_id}/{partition_id}/`, and the downstream `ShuffleReaderExec` reads those files back, either locally or over Arrow Flight. That happens regardless of how large the intermediate result is, so a query with many small stages pays a full write-then-read round trip per stage boundary.

Measurements shared on #2308 suggest this is the single largest avoidable cost in the current model. Pointing `work_dir` at a RAM disk large enough to hold the whole shuffle (4 GiB in that experiment) cut TPC-H runtime by roughly 40% and recovered most of the remaining gap against a pipelined engine, once AQE and the related join configs were enabled. In other words, a large share of what looks like "the cost of the blocking shuffle" is really the cost of going through the filesystem, not the cost of the barrier.

A RAM disk is not a general answer. It has to be sized for the worst-case shuffle up front, it is a deployment-level decision rather than a per-query one, and it turns an overflow into a hard failure instead of a slowdown.

**Describe the solution you'd like**

Keep shuffle output in a bounded in-memory buffer and spill to the existing file path only once that buffer is exceeded.

Sketch:

- A configurable per-executor (or per-task) budget, e.g. `ballista.shuffle.writer.memory_buffer_bytes`, ideally accounted through the executor memory pool rather than a standalone limit so it composes with `--memory-pool-size`.
- The writer buffers finished partitions in memory and records them as in-memory locations. Once the budget is exhausted, subsequent partitions (or the oldest buffered ones) are flushed to files exactly as today.
- `ShuffleReaderExec` and the Flight service serve from memory when the block is resident and fall back to the file path otherwise. The `IO_BLOCK_TRANSPORT` path should be able to serve buffered blocks without a round trip through the filesystem.
- Fault tolerance: buffered blocks are lost when an executor dies. The existing `FetchPartitionError` handling already covers this by re-running the map tasks that produced the missing partitions, so the recovery story does not change, but deployments that want durability should be able to force the file path (or a future remote shuffle service, #1539).

**Describe alternatives you've considered**

- Point `work_dir` at tmpfs. Works today and is what the experiment did, but it has the sizing and failure-mode problems above.
- A remote shuffle service (#1539). Complementary rather than an alternative: it changes where durable shuffle data lives, not whether small intermediates touch a filesystem at all.
- Full streaming/pipelined exchange (#1151, #2003). Much larger change, and gives up the properties documented in the shuffle design page.

**Additional context**

- Design rationale for the current model: `docs/source/contributors-guide/shuffle.md` (added in #2308), section "Directions that do not require abandoning the model".
- Related: #320 (memory management in the shuffle writer), #1952 (bounding decoded reader memory), #660 (cheaper on-disk shuffle format).
- Measurement reported by @Dandandan in #2308.

Contributor guide

Open the contributing guide

Research direction

Start by reading ShuffleWriterExec and SortShuffleWriterExec, then inspect ShuffleReaderExec, the IO_BLOCK_TRANSPORT path, and the shuffle design in docs/source/contributors-guide/shuffle.md. Trace existing FetchPartitionError recovery and related shuffle memory work in #320 and #1952. Done means a configured memory budget keeps eligible blocks in memory, spills excess through the existing file path, and lets readers and Flight serve resident blocks without filesystem round trips.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
distributed-systems, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.