Skip to content

feat(executor): bound executor memory via --memory-pool-size#1624

Merged
andygrove merged 8 commits intoapache:mainfrom
andygrove:feat/executor-memory-pool
Apr 29, 2026
Merged

feat(executor): bound executor memory via --memory-pool-size#1624
andygrove merged 8 commits intoapache:mainfrom
andygrove:feat/executor-memory-pool

Conversation

@andygrove
Copy link
Copy Markdown
Member

@andygrove andygrove commented Apr 29, 2026

Which issue does this PR close?

Closes #1563

Rationale for this change

The executor today builds a RuntimeEnv with no MemoryPool, so DataFusion uses its unbounded default and spillable operators (sort, hash join, hash agg) grow until the host OOMs. There is no way to bound executor memory or trigger spilling. The tuning-guide.md already flagged this as future work.

What changes are included in this PR?

Adds an opt-in --memory-pool-size <SIZE> flag (e.g. 8GB, 512MiB, plain bytes) on the executor binary. When set, every task receives an isolated FairSpillPool of size total / concurrent_tasks. Wrapping is applied after the base RuntimeProducer is resolved, so it composes with embedder-supplied producers (including the existing S3 helper) and preserves their DiskManager, CacheManager, and ObjectStoreRegistry. Hard error at startup if the per-task share would round to zero.

Adds bytesize as a dep on ballista-executor for size parsing. No public API changes outside the executor crate.

Are there any user-facing changes?

Yes. New optional CLI flag and a new section in docs/source/user-guide/tuning-guide.md. Default behavior is unchanged when the flag is omitted.

Adds wrap_runtime_producer_with_memory_pool, which takes a RuntimeProducer
and a total byte budget, and returns a new producer that installs a
FairSpillPool of size total/concurrent_tasks on every produced RuntimeEnv.
Also moves TryFrom<Config> impl above the test module in config.rs to fix
a pre-existing clippy::items_after_test_module lint.
Wire wrap_runtime_producer_with_memory_pool into start_executor_process
so that when --memory-pool-size is set, each task's RuntimeEnv receives
a FairSpillPool sized to total / concurrent_tasks. Remove the now-unused
#[allow(dead_code)] attribute and pub(crate) visibility from the helper.
@andygrove andygrove marked this pull request as ready for review April 29, 2026 15:19
Resolves rustdoc broken-intra-doc-link error in executor_process.rs.
RuntimeEnv was not imported, so the bare [RuntimeEnv] link could not
resolve. Use a fully qualified path link instead of adding an otherwise
unused import.
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Apr 29, 2026
total,
concurrent_tasks,
)?;
let per_task = total / concurrent_tasks as u64;
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i wonder should we divide usable memory across tasks or use shared memory pool across all of them? is spark doing this ?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i guess this is consistent with spark executor, is it ?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the approach that we use in Comet - each task gets its own fair pool and within a task, the operators share that pool. This keeps things predictable and stable.

We can certainly explore adding other config options in the future.

@andygrove andygrove merged commit e5f2216 into apache:main Apr 29, 2026
18 checks passed
@andygrove andygrove deleted the feat/executor-memory-pool branch April 29, 2026 19:07
@andygrove
Copy link
Copy Markdown
Member Author

Thanks for the review @milenkovicm

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

How to limit the maximum memory usage of the ballista executor?

2 participants