You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/04/29 19:07:38 UTC
(datafusion-ballista) branch main updated: feat(executor): bound executor memory via --memory-pool-size (#1624)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new e5f221641 feat(executor): bound executor memory via --memory-pool-size (#1624)
e5f221641 is described below
commit e5f2216411692a4bb34e281263035ca9ce521ada
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Apr 29 13:07:34 2026 -0600
feat(executor): bound executor memory via --memory-pool-size (#1624)
---
Cargo.lock | 7 ++
ballista/executor/Cargo.toml | 1 +
ballista/executor/src/config.rs | 49 +++++++++++++
ballista/executor/src/executor_process.rs | 112 ++++++++++++++++++++++++++++++
docs/source/user-guide/tuning-guide.md | 29 +++++++-
5 files changed, 196 insertions(+), 2 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 2b9b6014a..08c6b0292 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1102,6 +1102,7 @@ dependencies = [
"arrow-flight",
"async-trait",
"ballista-core",
+ "bytesize",
"clap 4.6.1",
"dashmap",
"datafusion",
@@ -1421,6 +1422,12 @@ dependencies = [
"either",
]
+[[package]]
+name = "bytesize"
+version = "2.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6bd91ee7b2422bcb158d90ef4d14f75ef67f340943fc4149891dcce8f8b972a3"
+
[[package]]
name = "bzip2"
version = "0.6.1"
diff --git a/ballista/executor/Cargo.toml b/ballista/executor/Cargo.toml
index 6ba1352d9..8a0043ddf 100644
--- a/ballista/executor/Cargo.toml
+++ b/ballista/executor/Cargo.toml
@@ -43,6 +43,7 @@ arrow = { workspace = true }
arrow-flight = { workspace = true }
async-trait = { workspace = true }
ballista-core = { path = "../core", version = "53.0.0" }
+bytesize = "2"
clap = { workspace = true, optional = true }
dashmap = { workspace = true }
datafusion = { workspace = true }
diff --git a/ballista/executor/src/config.rs b/ballista/executor/src/config.rs
index 0bf8c97b9..860617647 100644
--- a/ballista/executor/src/config.rs
+++ b/ballista/executor/src/config.rs
@@ -31,6 +31,17 @@ use ballista_core::error::BallistaError;
use crate::executor_process::ExecutorProcessConfig;
use crate::metrics::ExecutorMetricCollectionPolicy;
+/// Parse a human-readable size string into a byte count.
+///
+/// Accepts decimal SI suffixes (`KB`, `MB`, `GB`) where 1KB = 1000 bytes,
+/// binary IEC suffixes (`KiB`, `MiB`, `GiB`) where 1KiB = 1024 bytes, and
+/// plain integers (interpreted as bytes).
+fn parse_memory_pool_size(s: &str) -> Result<u64, String> {
+ s.parse::<bytesize::ByteSize>()
+ .map(|b| b.as_u64())
+ .map_err(|e| format!("invalid byte size '{s}': {e}"))
+}
+
/// Command-line arguments for configuring a Ballista executor.
///
/// This struct is parsed from command-line arguments using clap and contains
@@ -151,6 +162,15 @@ pub struct Config {
help = "Metric collection policy of this executor instance"
)]
pub metric_collection_policy: ExecutorMetricCollectionPolicy,
+ /// Optional total memory budget for the executor. Accepts human-readable
+ /// values like "8GB", "512MiB", or a plain byte count. When set, every
+ /// task gets a FairSpillPool of size `memory_pool_size / concurrent_tasks`.
+ #[arg(
+ long,
+ value_parser = parse_memory_pool_size,
+ help = "Optional total executor memory budget (e.g. \"8GB\", \"512MiB\"). Each concurrent task receives an equal share."
+ )]
+ pub memory_pool_size: Option<u64>,
}
impl TryFrom<Config> for ExecutorProcessConfig {
@@ -179,6 +199,7 @@ impl TryFrom<Config> for ExecutorProcessConfig {
grpc_server_config: ballista_core::utils::GrpcServerConfig::default(),
executor_heartbeat_interval_seconds: opt.executor_heartbeat_interval_seconds,
metric_collection_policy: opt.metric_collection_policy,
+ memory_pool_size: opt.memory_pool_size,
override_execution_engine: None,
override_function_registry: None,
override_config_producer: None,
@@ -190,3 +211,31 @@ impl TryFrom<Config> for ExecutorProcessConfig {
})
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::parse_memory_pool_size;
+
+ #[test]
+ fn parse_decimal_suffix() {
+ assert_eq!(parse_memory_pool_size("8GB").unwrap(), 8_000_000_000);
+ assert_eq!(parse_memory_pool_size("1KB").unwrap(), 1_000);
+ }
+
+ #[test]
+ fn parse_binary_suffix() {
+ assert_eq!(parse_memory_pool_size("512MiB").unwrap(), 512 * 1024 * 1024);
+ assert_eq!(parse_memory_pool_size("1KiB").unwrap(), 1024);
+ }
+
+ #[test]
+ fn parse_plain_integer_is_bytes() {
+ assert_eq!(parse_memory_pool_size("1024").unwrap(), 1024);
+ }
+
+ #[test]
+ fn parse_rejects_invalid() {
+ assert!(parse_memory_pool_size("banana").is_err());
+ assert!(parse_memory_pool_size("").is_err());
+ }
+}
diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs
index 40a8545ff..6447b8106 100644
--- a/ballista/executor/src/executor_process.rs
+++ b/ballista/executor/src/executor_process.rs
@@ -41,7 +41,9 @@ use tokio::task::JoinHandle;
use tokio::{fs, time};
use uuid::Uuid;
+use datafusion::execution::memory_pool::{FairSpillPool, MemoryPool};
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+use datafusion::prelude::SessionConfig;
use ballista_core::config::{LogRotationPolicy, TaskSchedulingPolicy};
use ballista_core::error::BallistaError;
@@ -72,6 +74,33 @@ use crate::shutdown::ShutdownNotifier;
use crate::{ArrowFlightServerProvider, terminate};
use crate::{execution_loop, executor_server};
+/// Wrap a [`RuntimeProducer`] so that every produced
+/// [`RuntimeEnv`](datafusion::execution::runtime_env::RuntimeEnv) carries a
+/// fresh [`FairSpillPool`] of size `total_bytes / concurrent_tasks`.
+///
+/// Returns an error if the per-task share would be zero (i.e. `total_bytes <
+/// concurrent_tasks`). The inner env's disk manager, cache manager, and
+/// object store registry are preserved via [`RuntimeEnvBuilder::from_runtime_env`].
+fn wrap_runtime_producer_with_memory_pool(
+ inner: RuntimeProducer,
+ total_bytes: u64,
+ concurrent_tasks: usize,
+) -> Result<RuntimeProducer, BallistaError> {
+ let per_task = (total_bytes / concurrent_tasks as u64) as usize;
+ if per_task == 0 {
+ return Err(BallistaError::Configuration(format!(
+ "memory_pool_size ({total_bytes} bytes) is smaller than concurrent_tasks ({concurrent_tasks})"
+ )));
+ }
+ Ok(Arc::new(move |session_config: &SessionConfig| {
+ let inner_env = inner(session_config)?;
+ let pool: Arc<dyn MemoryPool> = Arc::new(FairSpillPool::new(per_task));
+ RuntimeEnvBuilder::from_runtime_env(&inner_env)
+ .with_memory_pool(pool)
+ .build_arc()
+ }))
+}
+
/// Configuration for the executor process.
///
/// This struct contains all settings needed to run an executor, including
@@ -120,6 +149,11 @@ pub struct ExecutorProcessConfig {
pub executor_heartbeat_interval_seconds: u64,
/// Metric collection policy of this executor instance
pub metric_collection_policy: ExecutorMetricCollectionPolicy,
+ /// Optional total memory pool size in bytes. When set, every task's
+ /// runtime env receives a FairSpillPool of size
+ /// `memory_pool_size / concurrent_tasks`. When `None`, no pool is
+ /// installed and DataFusion falls back to its unbounded default.
+ pub memory_pool_size: Option<u64>,
/// Optional execution engine to use to execute physical plans, will default to
/// DataFusion if none is provided.
pub override_execution_engine: Option<Arc<dyn ExecutionEngine>>,
@@ -176,6 +210,7 @@ impl Default for ExecutorProcessConfig {
grpc_server_config: Default::default(),
executor_heartbeat_interval_seconds: 60,
metric_collection_policy: ExecutorMetricCollectionPolicy::default(),
+ memory_pool_size: None,
override_execution_engine: None,
override_function_registry: None,
override_runtime_producer: None,
@@ -261,6 +296,21 @@ pub async fn start_executor_process(
})
});
+ let runtime_producer = if let Some(total) = opt.memory_pool_size {
+ let producer = wrap_runtime_producer_with_memory_pool(
+ runtime_producer,
+ total,
+ concurrent_tasks,
+ )?;
+ let per_task = total / concurrent_tasks as u64;
+ info!(
+ "Memory pool: total {total} bytes split into {concurrent_tasks} tasks ({per_task} bytes each)"
+ );
+ producer
+ } else {
+ runtime_producer
+ };
+
let logical = opt
.override_logical_codec
.clone()
@@ -928,3 +978,65 @@ mod tests {
path
}
}
+
+#[cfg(test)]
+mod memory_pool_tests {
+ use super::*;
+ use datafusion::execution::memory_pool::MemoryLimit;
+ use datafusion::execution::object_store::DefaultObjectStoreRegistry;
+ use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
+ use std::sync::Arc;
+
+ fn baseline_producer() -> RuntimeProducer {
+ Arc::new(|_| Ok(Arc::new(RuntimeEnv::default())))
+ }
+
+ #[test]
+ fn returns_error_when_total_smaller_than_concurrent_tasks() {
+ let inner = baseline_producer();
+ let result = wrap_runtime_producer_with_memory_pool(inner, 4, 8);
+ assert!(result.is_err());
+ let msg = result.err().unwrap().to_string();
+ assert!(msg.contains("memory_pool_size"));
+ assert!(msg.contains("concurrent_tasks"));
+ }
+
+ #[test]
+ fn produces_runtime_with_fair_spill_pool_of_per_task_size() {
+ let total = 8u64 * 1024 * 1024 * 1024;
+ let concurrent = 8usize;
+ let expected_per_task = (total / concurrent as u64) as usize;
+
+ let wrapped = wrap_runtime_producer_with_memory_pool(
+ baseline_producer(),
+ total,
+ concurrent,
+ )
+ .unwrap();
+ let env = wrapped(&SessionConfig::new()).unwrap();
+
+ match env.memory_pool.memory_limit() {
+ MemoryLimit::Finite(n) => assert_eq!(n, expected_per_task),
+ MemoryLimit::Infinite => panic!("expected Finite limit, got Infinite"),
+ MemoryLimit::Unknown => panic!("expected Finite limit, got Unknown"),
+ }
+ }
+
+ #[test]
+ fn preserves_inner_object_store_registry() {
+ let registry: Arc<dyn datafusion::execution::object_store::ObjectStoreRegistry> =
+ Arc::new(DefaultObjectStoreRegistry::new());
+ let registry_for_inner = registry.clone();
+ let inner: RuntimeProducer = Arc::new(move |_| {
+ let env = RuntimeEnvBuilder::new()
+ .with_object_store_registry(registry_for_inner.clone())
+ .build()?;
+ Ok(Arc::new(env))
+ });
+
+ let wrapped = wrap_runtime_producer_with_memory_pool(inner, 1024, 1).unwrap();
+ let env = wrapped(&SessionConfig::new()).unwrap();
+
+ assert!(Arc::ptr_eq(&env.object_store_registry, ®istry));
+ }
+}
diff --git a/docs/source/user-guide/tuning-guide.md b/docs/source/user-guide/tuning-guide.md
index dfe3c2937..9a6334d22 100644
--- a/docs/source/user-guide/tuning-guide.md
+++ b/docs/source/user-guide/tuning-guide.md
@@ -63,8 +63,33 @@ Increasing this configuration setting will increase the number of tasks that eac
this will also mean that the executor will use more memory. If executors are failing due to out-of-memory errors then
decreasing the number of concurrent tasks may help.
-In the future, Ballista will have better support for tracking memory usage and allocating tasks based on available
-memory, as well as supporting spill-to-disk to reduce memory pressure.
+## Configuring Executor Memory Pool
+
+By default the executor uses DataFusion's unbounded memory pool, so spillable
+operators (sort, hash join, hash aggregate) grow until the host runs out of
+memory. To bound executor memory and let those operators spill to disk under
+pressure, pass `--memory-pool-size` when starting the executor:
+
+```sh
+ballista-executor --memory-pool-size 8GB --concurrent-tasks 8
+```
+
+The argument accepts human-readable sizes (`8GB`, `512MiB`) or a plain byte
+count. SI suffixes (`KB`/`MB`/`GB`) are powers of 10; IEC suffixes
+(`KiB`/`MiB`/`GiB`) are powers of 2.
+
+The total budget is divided equally across concurrent task slots: each task
+receives its own `FairSpillPool` of size `memory_pool_size / concurrent_tasks`.
+With `--memory-pool-size 8GB --concurrent-tasks 8`, every task sees a 1 GB
+pool, fully isolated from other tasks. Idle slots do not lend their share to
+busy ones, which keeps task memory predictable at the cost of some unused
+capacity when the executor is under-utilized.
+
+The executor refuses to start if the per-task share would round to zero (i.e.
+`memory_pool_size < concurrent_tasks`).
+
+When `--memory-pool-size` is not set, the executor behaves as before with no
+memory pool installed.
## Shuffle Implementation
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org