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 f67969cf6 feat: Share read-side runtime state across a session's tasks
on the executor (#1995)
f67969cf6 is described below
commit f67969cf6b642d49114674f10d30c520839d155e
Author: Andy Grove <[email protected]>
AuthorDate: Sun Jul 12 07:45:01 2026 -0600
feat: Share read-side runtime state across a session's tasks on the
executor (#1995)
---
Cargo.lock | 1 +
ballista/executor/Cargo.toml | 1 +
ballista/executor/src/config.rs | 10 ++
ballista/executor/src/execution_loop.rs | 3 +-
ballista/executor/src/executor.rs | 105 ++++++++++++++
ballista/executor/src/executor_process.rs | 171 +++++++++++++----------
ballista/executor/src/executor_server.rs | 4 +-
ballista/executor/src/lib.rs | 2 +
ballista/executor/src/runtime_cache.rs | 218 ++++++++++++++++++++++++++++++
python/Cargo.lock | 24 +++-
10 files changed, 462 insertions(+), 77 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index fa3e259c2..b918eb711 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1159,6 +1159,7 @@ dependencies = [
"futures",
"libc",
"log",
+ "lru 0.18.0",
"memory-stats",
"mimalloc",
"parking_lot",
diff --git a/ballista/executor/Cargo.toml b/ballista/executor/Cargo.toml
index c80e68b9c..aa60cb2fd 100644
--- a/ballista/executor/Cargo.toml
+++ b/ballista/executor/Cargo.toml
@@ -50,6 +50,7 @@ datafusion = { workspace = true }
datafusion-proto = { workspace = true }
futures = { workspace = true }
log = { workspace = true }
+lru = "0.18"
memory-stats = { version = "1.2.0", features = ["always_use_statm"] }
mimalloc = { workspace = true, optional = true }
parking_lot = { workspace = true }
diff --git a/ballista/executor/src/config.rs b/ballista/executor/src/config.rs
index 8ddeb4ab3..48b5a28c9 100644
--- a/ballista/executor/src/config.rs
+++ b/ballista/executor/src/config.rs
@@ -171,6 +171,15 @@ pub struct Config {
help = "Optional total executor memory budget (e.g. \"8GB\",
\"512MiB\"). Each concurrent task receives an equal share."
)]
pub memory_pool_size: Option<u64>,
+ /// Maximum number of sessions whose shared runtime state (object-store
+ /// clients, Parquet footer cache) is retained on the executor (LRU). `0`
+ /// disables caching.
+ #[arg(
+ long,
+ default_value_t = 16,
+ help = "Max number of sessions whose shared runtime state
(object-store clients, Parquet footer cache) is retained on the executor (LRU).
0 disables caching."
+ )]
+ pub session_runtime_cache_capacity: usize,
/// Number of seconds established client connection should be cached if
not used (0 means no cache)
#[arg(
long,
@@ -207,6 +216,7 @@ impl TryFrom<Config> for ExecutorProcessConfig {
executor_heartbeat_interval_seconds:
opt.executor_heartbeat_interval_seconds,
metric_collection_policy: opt.metric_collection_policy,
memory_pool_size: opt.memory_pool_size,
+ session_runtime_cache_capacity: opt.session_runtime_cache_capacity,
override_execution_engine: None,
override_function_registry: None,
override_config_producer: None,
diff --git a/ballista/executor/src/execution_loop.rs
b/ballista/executor/src/execution_loop.rs
index 8a759f501..ff0838a7e 100644
--- a/ballista/executor/src/execution_loop.rs
+++ b/ballista/executor/src/execution_loop.rs
@@ -262,7 +262,8 @@ async fn run_received_task<T: 'static + AsLogicalPlan, U:
'static + AsExecutionP
let task_higher_order_functions =
executor.function_registry.higher_order_functions.clone();
- let runtime = executor.produce_runtime(&session_config)?;
+ let runtime =
+ executor.produce_runtime_for_session(&task.session_id,
&session_config)?;
let session_id = task.session_id.clone();
let task_context = Arc::new(TaskContext::new(
diff --git a/ballista/executor/src/executor.rs
b/ballista/executor/src/executor.rs
index 09c0af31c..d5f90b40a 100644
--- a/ballista/executor/src/executor.rs
+++ b/ballista/executor/src/executor.rs
@@ -23,6 +23,7 @@ use crate::execution_engine::QueryStageExecutor;
use crate::execution_loop::any_to_string;
use crate::metrics::ExecutorMetricsCollector;
use crate::metrics::LoggingMetricsCollector;
+use crate::runtime_cache::SessionRuntimeCache;
use ballista_core::ConfigProducer;
use ballista_core::JobId;
use ballista_core::RuntimeProducer;
@@ -97,6 +98,11 @@ pub struct Executor {
/// Execution engine that the executor will delegate to
/// for executing query stages
pub(crate) execution_engine: Arc<dyn ExecutionEngine>,
+
+ /// Optional session-keyed cache of shared base runtime envs. When set,
+ /// `produce_runtime_for_session` reuses read-side state across a session's
+ /// tasks; when `None`, each task builds a runtime from `runtime_producer`.
+ session_runtime_cache: Option<Arc<dyn SessionRuntimeCache>>,
}
impl Executor {
@@ -144,6 +150,7 @@ impl Executor {
concurrent_tasks,
abort_handles: Default::default(),
execution_engine,
+ session_runtime_cache: None,
}
}
/// Creates new Executor with default `ExecutionEngine`.
@@ -167,6 +174,7 @@ impl Executor {
concurrent_tasks,
abort_handles: Default::default(),
execution_engine: Arc::new(DefaultExecutionEngine::new()),
+ session_runtime_cache: None,
}
}
}
@@ -180,6 +188,30 @@ impl Executor {
(self.runtime_producer)(config)
}
+ /// Attaches (or clears) the session-keyed runtime cache. Cloning the
+ /// `Executor` shares the same cache via `Arc`.
+ pub fn with_session_runtime_cache(
+ mut self,
+ cache: Option<Arc<dyn SessionRuntimeCache>>,
+ ) -> Self {
+ self.session_runtime_cache = cache;
+ self
+ }
+
+ /// Produces the runtime for a task, reusing the session's shared read-side
+ /// state when a session cache is attached; otherwise builds a runtime per
+ /// task via the `runtime_producer`.
+ pub fn produce_runtime_for_session(
+ &self,
+ session_id: &str,
+ config: &SessionConfig,
+ ) -> datafusion::error::Result<Arc<RuntimeEnv>> {
+ match &self.session_runtime_cache {
+ Some(cache) => cache.produce_runtime(session_id, config),
+ None => (self.runtime_producer)(config),
+ }
+ }
+
/// Creates a default [`SessionConfig`] using the configured config
producer.
pub fn produce_config(&self) -> SessionConfig {
(self.config_producer)()
@@ -269,6 +301,9 @@ impl Executor {
mod test {
use crate::execution_engine::{DefaultQueryStageExec, ShuffleWriterVariant};
use crate::executor::Executor;
+ use crate::runtime_cache::{
+ DefaultSessionRuntimeCache, MemoryPoolPolicy, SessionRuntimeCache,
+ };
use ballista_core::RuntimeProducer;
use ballista_core::execution_plans::ShuffleWriterExec;
use ballista_core::serde::protobuf::ExecutorRegistration;
@@ -278,11 +313,13 @@ mod test {
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::error::{DataFusionError, Result};
use datafusion::execution::context::TaskContext;
+ use datafusion::execution::runtime_env::RuntimeEnv;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning,
PlanProperties,
RecordBatchStream, SendableRecordBatchStream,
};
+ use datafusion::prelude::SessionConfig;
use datafusion::prelude::SessionContext;
use futures::Stream;
use std::pin::Pin;
@@ -459,4 +496,72 @@ mod test {
let inner_result = result.unwrap().unwrap();
assert!(inner_result.is_err());
}
+
+ #[test]
+ fn produce_runtime_for_session_shares_read_side_state() {
+ let executor_registration = ExecutorRegistration {
+ id: "executor".to_string(),
+ port: 0,
+ grpc_port: 0,
+ specification: None,
+ host: None,
+ os_info: None,
+ };
+ let config_producer = Arc::new(default_config_producer);
+
+ // A base producer that builds a fresh env each call, plus an identity
+ // pool policy, so shared read-side state is observable via ptr
equality.
+ let base_producer: RuntimeProducer =
+ Arc::new(|_| Ok(Arc::new(RuntimeEnv::default())));
+ let identity: MemoryPoolPolicy = Arc::new(|base, _| Ok(base));
+ let cache: Arc<dyn SessionRuntimeCache> = Arc::new(
+ DefaultSessionRuntimeCache::new(base_producer.clone(), identity,
4),
+ );
+
+ let executor = Executor::new_basic(
+ executor_registration,
+ "/tmp",
+ base_producer,
+ config_producer,
+ 2,
+ )
+ .with_session_runtime_cache(Some(cache));
+
+ let cfg = SessionConfig::new();
+ let e1 = executor.produce_runtime_for_session("s1", &cfg).unwrap();
+ let e2 = executor.produce_runtime_for_session("s1", &cfg).unwrap();
+ let e3 = executor.produce_runtime_for_session("s2", &cfg).unwrap();
+
+ assert!(Arc::ptr_eq(&e1.cache_manager, &e2.cache_manager));
+ assert!(!Arc::ptr_eq(&e1.cache_manager, &e3.cache_manager));
+ }
+
+ #[test]
+ fn produce_runtime_for_session_falls_back_without_cache() {
+ let executor_registration = ExecutorRegistration {
+ id: "executor".to_string(),
+ port: 0,
+ grpc_port: 0,
+ specification: None,
+ host: None,
+ os_info: None,
+ };
+ let config_producer = Arc::new(default_config_producer);
+ let base_producer: RuntimeProducer =
+ Arc::new(|_| Ok(Arc::new(RuntimeEnv::default())));
+
+ // No cache attached: each call builds a fresh env.
+ let executor = Executor::new_basic(
+ executor_registration,
+ "/tmp",
+ base_producer,
+ config_producer,
+ 2,
+ );
+
+ let cfg = SessionConfig::new();
+ let e1 = executor.produce_runtime_for_session("s1", &cfg).unwrap();
+ let e2 = executor.produce_runtime_for_session("s1", &cfg).unwrap();
+ assert!(!Arc::ptr_eq(&e1.cache_manager, &e2.cache_manager));
+ }
}
diff --git a/ballista/executor/src/executor_process.rs
b/ballista/executor/src/executor_process.rs
index 4e37687e8..360a89491 100644
--- a/ballista/executor/src/executor_process.rs
+++ b/ballista/executor/src/executor_process.rs
@@ -42,7 +42,7 @@ use tokio::{fs, time};
use uuid::Uuid;
use datafusion::execution::memory_pool::{FairSpillPool, MemoryPool};
-use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
use datafusion::prelude::SessionConfig;
use ballista_core::config::{LogRotationPolicy, TaskSchedulingPolicy};
@@ -70,36 +70,46 @@ use crate::executor_server::TERMINATING;
use crate::flight_service::BallistaFlightService;
use crate::metrics::ExecutorMetricCollectionPolicy;
use crate::metrics::LoggingMetricsCollector;
+use crate::runtime_cache::{
+ DefaultSessionRuntimeCache, MemoryPoolPolicy, SessionRuntimeCache,
+};
use crate::shutdown::Shutdown;
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`.
+/// Builds a per-task memory-pool policy: each task's runtime is rebuilt from
the
+/// shared base env with a fresh [`FairSpillPool`] of size
+/// `total_bytes / concurrent_tasks`. The base env's disk manager, cache
manager,
+/// and object-store registry are preserved via
+/// [`RuntimeEnvBuilder::from_runtime_env`].
///
/// 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,
+/// concurrent_tasks`).
+fn memory_pool_policy(
total_bytes: u64,
concurrent_tasks: usize,
-) -> Result<RuntimeProducer, BallistaError> {
+) -> Result<MemoryPoolPolicy, 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()
- }))
+ Ok(Arc::new(
+ move |base: Arc<RuntimeEnv>, _config: &SessionConfig| {
+ let pool: Arc<dyn MemoryPool> =
Arc::new(FairSpillPool::new(per_task));
+ RuntimeEnvBuilder::from_runtime_env(&base)
+ .with_memory_pool(pool)
+ .build_arc()
+ },
+ ))
+}
+
+/// A no-op policy: the task uses the shared base env unchanged (DataFusion's
+/// default unbounded pool). Used when `--memory-pool-size` is unset.
+fn identity_pool_policy() -> MemoryPoolPolicy {
+ Arc::new(|base, _config| Ok(base))
}
/// Configuration for the executor process.
@@ -155,6 +165,11 @@ pub struct ExecutorProcessConfig {
/// `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>,
+ /// Maximum number of sessions whose shared base runtime env is retained on
+ /// the executor (LRU). Sharing reuses object-store clients and the Parquet
+ /// footer cache across a session's tasks and queries. `0` disables caching
+ /// and builds a fresh runtime per task.
+ pub session_runtime_cache_capacity: usize,
/// 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>>,
@@ -214,6 +229,7 @@ impl Default for ExecutorProcessConfig {
executor_heartbeat_interval_seconds: 60,
metric_collection_policy:
ExecutorMetricCollectionPolicy::default(),
memory_pool_size: None,
+ session_runtime_cache_capacity: 16,
override_execution_engine: None,
override_function_registry: None,
override_runtime_producer: None,
@@ -290,7 +306,7 @@ pub async fn start_executor_process(
.unwrap_or_else(|| Arc::new(default_config_producer));
let wd = work_dir.clone();
- let runtime_producer: RuntimeProducer =
+ let base_runtime_producer: RuntimeProducer =
opt.override_runtime_producer.clone().unwrap_or_else(|| {
Arc::new(move |_| {
let runtime_env = RuntimeEnvBuilder::new()
@@ -300,19 +316,24 @@ 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 pool_policy: MemoryPoolPolicy = if let Some(total) =
opt.memory_pool_size {
+ let policy = memory_pool_policy(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
+ policy
} else {
- runtime_producer
+ identity_pool_policy()
+ };
+
+ // Combined producer preserving the current per-task behavior: build a
fresh
+ // base env, then apply the pool policy. Used by
`Executor::produce_runtime`
+ // and as the fallback when session caching is disabled.
+ let runtime_producer: RuntimeProducer = {
+ let base = base_runtime_producer.clone();
+ let policy = pool_policy.clone();
+ Arc::new(move |config: &SessionConfig| policy(base(config)?, config))
};
let logical = opt
@@ -330,26 +351,43 @@ pub async fn start_executor_process(
datafusion_proto::protobuf::PhysicalPlanNode,
> = BallistaCodec::new(logical, physical);
- let executor = Arc::new(Executor::new(
- executor_meta.clone(),
- &work_dir,
- runtime_producer,
- config_producer,
- opt.override_function_registry.clone().unwrap_or_default(),
- metrics_collector,
- concurrent_tasks,
- opt.override_execution_engine.clone().unwrap_or_else(|| {
- if opt.client_ttl > 0 {
- let client_pool =
- Arc::new(DefaultBallistaClientPool::with_eviction_thread(
- Duration::from_secs(opt.client_ttl),
- ));
- Arc::new(DefaultExecutionEngine::with_client_pool(client_pool))
- } else {
- Arc::new(DefaultExecutionEngine::new())
- }
- }),
- ));
+ // Session caching applies only to the default runtime producer. With an
+ // override producer we cannot split base + pool, so we serve per-task as
+ // before by leaving the cache unset.
+ let session_runtime_cache: Option<Arc<dyn SessionRuntimeCache>> =
+ if opt.override_runtime_producer.is_none() {
+ Some(Arc::new(DefaultSessionRuntimeCache::new(
+ base_runtime_producer.clone(),
+ pool_policy.clone(),
+ opt.session_runtime_cache_capacity,
+ )))
+ } else {
+ None
+ };
+
+ let executor = Arc::new(
+ Executor::new(
+ executor_meta.clone(),
+ &work_dir,
+ runtime_producer,
+ config_producer,
+ opt.override_function_registry.clone().unwrap_or_default(),
+ metrics_collector,
+ concurrent_tasks,
+ opt.override_execution_engine.clone().unwrap_or_else(|| {
+ if opt.client_ttl > 0 {
+ let client_pool =
+
Arc::new(DefaultBallistaClientPool::with_eviction_thread(
+ Duration::from_secs(opt.client_ttl),
+ ));
+
Arc::new(DefaultExecutionEngine::with_client_pool(client_pool))
+ } else {
+ Arc::new(DefaultExecutionEngine::new())
+ }
+ }),
+ )
+ .with_session_runtime_cache(session_runtime_cache),
+ );
let connect_timeout = opt.scheduler_connect_timeout_seconds as u64;
let session_config = (executor.config_producer)();
@@ -1064,14 +1102,9 @@ mod memory_pool_tests {
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);
+ let result = memory_pool_policy(4, 8);
assert!(result.is_err());
let msg = result.err().unwrap().to_string();
assert!(msg.contains("memory_pool_size"));
@@ -1084,13 +1117,9 @@ mod memory_pool_tests {
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();
+ let policy = memory_pool_policy(total, concurrent).unwrap();
+ let base = Arc::new(RuntimeEnv::default());
+ let env = policy(base, &SessionConfig::new()).unwrap();
match env.memory_pool.memory_limit() {
MemoryLimit::Finite(n) => assert_eq!(n, expected_per_task),
@@ -1100,20 +1129,26 @@ mod memory_pool_tests {
}
#[test]
- fn preserves_inner_object_store_registry() {
+ fn preserves_base_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 base = Arc::new(
+ RuntimeEnvBuilder::new()
+ .with_object_store_registry(registry.clone())
+ .build()
+ .unwrap(),
+ );
- let wrapped = wrap_runtime_producer_with_memory_pool(inner, 1024,
1).unwrap();
- let env = wrapped(&SessionConfig::new()).unwrap();
+ let policy = memory_pool_policy(1024, 1).unwrap();
+ let env = policy(base, &SessionConfig::new()).unwrap();
assert!(Arc::ptr_eq(&env.object_store_registry, ®istry));
}
+
+ #[test]
+ fn identity_policy_returns_base_unchanged() {
+ let base = Arc::new(RuntimeEnv::default());
+ let env = identity_pool_policy()(base.clone(),
&SessionConfig::new()).unwrap();
+ assert!(Arc::ptr_eq(&env, &base));
+ }
}
diff --git a/ballista/executor/src/executor_server.rs
b/ballista/executor/src/executor_server.rs
index e65dbb0e0..a78b006e9 100644
--- a/ballista/executor/src/executor_server.rs
+++ b/ballista/executor/src/executor_server.rs
@@ -386,7 +386,9 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> ExecutorServer<T,
&task.session_config,
);
- let runtime = self.executor.produce_runtime(&task.session_config);
+ let runtime = self
+ .executor
+ .produce_runtime_for_session(&task.session_id,
&task.session_config);
match (exec, runtime) {
(Ok(exec), Ok(runtime)) => {
diff --git a/ballista/executor/src/lib.rs b/ballista/executor/src/lib.rs
index f0d0f995e..70234ceaa 100644
--- a/ballista/executor/src/lib.rs
+++ b/ballista/executor/src/lib.rs
@@ -39,6 +39,8 @@ pub mod executor_server;
pub mod flight_service;
/// Metrics collection for executor runtime statistics.
pub mod metrics;
+/// Session-scoped cache of shared executor runtime environments.
+pub mod runtime_cache;
/// Graceful shutdown coordination for executor components.
pub mod shutdown;
/// Signal handling for process termination.
diff --git a/ballista/executor/src/runtime_cache.rs
b/ballista/executor/src/runtime_cache.rs
new file mode 100644
index 000000000..2cc0e05fb
--- /dev/null
+++ b/ballista/executor/src/runtime_cache.rs
@@ -0,0 +1,218 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Session-scoped reuse of executor runtime state.
+
+use std::num::NonZeroUsize;
+use std::sync::Arc;
+
+use ballista_core::RuntimeProducer;
+use datafusion::execution::runtime_env::RuntimeEnv;
+use datafusion::prelude::SessionConfig;
+use lru::LruCache;
+use parking_lot::Mutex;
+
+/// Derives a task's runtime from a shared base [`RuntimeEnv`] by installing a
+/// fresh per-task memory pool. The base's object-store registry and the cache
+/// manager's underlying file-metadata (footer) cache are preserved via
+///
[`RuntimeEnvBuilder::from_runtime_env`](datafusion::execution::runtime_env::RuntimeEnvBuilder::from_runtime_env)
+/// — the outer `CacheManager` is rebuilt around that same inner cache — so
+/// read-side state is shared while the memory pool stays per task.
+pub type MemoryPoolPolicy = Arc<
+ dyn Fn(Arc<RuntimeEnv>, &SessionConfig) ->
datafusion::error::Result<Arc<RuntimeEnv>>
+ + Send
+ + Sync,
+>;
+
+/// Produces a task's [`RuntimeEnv`], optionally reusing read-side state across
+/// the tasks of a session.
+///
+/// The executor calls [`produce_runtime`](Self::produce_runtime) for every
task;
+/// implementors decide whether and how to share base runtime state
(object-store
+/// clients, file-metadata cache) between a session's tasks.
+/// [`DefaultSessionRuntimeCache`] is the built-in implementation — provide a
+/// custom one to change the caching/sharing strategy.
+pub trait SessionRuntimeCache: Send + Sync {
+ /// Returns the per-task runtime for `session_id`.
+ fn produce_runtime(
+ &self,
+ session_id: &str,
+ config: &SessionConfig,
+ ) -> datafusion::error::Result<Arc<RuntimeEnv>>;
+}
+
+/// A bounded, session-keyed cache of shared *base* [`RuntimeEnv`]s.
+///
+/// A base env carries the read-side state safe to share across all tasks of a
+/// session: the object-store registry and the cache manager (whose Parquet
+/// footer cache is thereby reused across the session's tasks and queries).
+/// `RuntimeEnvBuilder::from_runtime_env` also carries over the disk manager
+/// (rooted at the executor's `work_dir`), so a session's tasks share one
+/// `DiskManager` too; this is safe because spill temp files are uniquely
+/// named, matching the standard one-`RuntimeEnv`-per-`SessionContext` model.
+/// Each task's real runtime is produced by applying [`MemoryPoolPolicy`] to
+/// the shared base, which installs a fresh per-task memory pool — so memory
+/// isolation is unchanged.
+///
+/// The cache is bounded by an LRU of `capacity` sessions. A capacity of `0`
+/// disables caching entirely: every call builds a fresh base env, matching the
+/// behavior of building a runtime per task.
+pub struct DefaultSessionRuntimeCache {
+ base_producer: RuntimeProducer,
+ pool_policy: MemoryPoolPolicy,
+ cache: Option<Mutex<LruCache<String, Arc<RuntimeEnv>>>>,
+}
+
+impl DefaultSessionRuntimeCache {
+ /// Creates a new cache that produces per-task runtimes from
`base_producer`
+ /// (invoked at most once per cached session) and `pool_policy` (invoked on
+ /// every call). `capacity` bounds the number of distinct sessions kept in
+ /// the LRU; `0` disables caching.
+ pub fn new(
+ base_producer: RuntimeProducer,
+ pool_policy: MemoryPoolPolicy,
+ capacity: usize,
+ ) -> Self {
+ let cache = NonZeroUsize::new(capacity).map(|cap|
Mutex::new(LruCache::new(cap)));
+ Self {
+ base_producer,
+ pool_policy,
+ cache,
+ }
+ }
+}
+
+impl SessionRuntimeCache for DefaultSessionRuntimeCache {
+ /// Returns the per-task runtime for `session_id`, reusing a cached base
env
+ /// when present and building + caching one on miss.
+ fn produce_runtime(
+ &self,
+ session_id: &str,
+ config: &SessionConfig,
+ ) -> datafusion::error::Result<Arc<RuntimeEnv>> {
+ let base = match &self.cache {
+ None => (self.base_producer)(config)?,
+ Some(cache) => {
+ if let Some(base) = cache.lock().get(session_id) {
+ base.clone()
+ } else {
+ // Build the base env without holding the lock, so a miss
+ // never stalls other sessions' lookups. A rare concurrent
+ // first-miss for the same session may build twice; that is
+ // harmless (idempotent, cheap) — the last writer wins and
+ // the extra env is dropped.
+ let base = (self.base_producer)(config)?;
+ cache.lock().put(session_id.to_string(), base.clone());
+ base
+ }
+ }
+ };
+ (self.pool_policy)(base, config)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use datafusion::execution::memory_pool::GreedyMemoryPool;
+ use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+
+ fn base_producer() -> RuntimeProducer {
+ Arc::new(|_| Ok(Arc::new(RuntimeEnv::default())))
+ }
+
+ /// Identity policy: the task shares the base env unchanged (no pool swap).
+ fn identity_policy() -> MemoryPoolPolicy {
+ Arc::new(|base, _| Ok(base))
+ }
+
+ /// Rebuilding policy: mimics production — a fresh per-task pool layered
onto
+ /// the shared base, preserving the base's read-side state.
+ fn per_task_pool_policy() -> MemoryPoolPolicy {
+ Arc::new(|base, _| {
+ RuntimeEnvBuilder::from_runtime_env(&base)
+ .with_memory_pool(Arc::new(GreedyMemoryPool::new(1024)))
+ .build_arc()
+ })
+ }
+
+ #[test]
+ fn same_session_shares_cache_manager() {
+ let cache =
+ DefaultSessionRuntimeCache::new(base_producer(),
identity_policy(), 4);
+ let cfg = SessionConfig::new();
+ let e1 = cache.produce_runtime("s1", &cfg).unwrap();
+ let e2 = cache.produce_runtime("s1", &cfg).unwrap();
+ assert!(Arc::ptr_eq(&e1.cache_manager, &e2.cache_manager));
+ }
+
+ #[test]
+ fn different_sessions_get_different_base() {
+ let cache =
+ DefaultSessionRuntimeCache::new(base_producer(),
identity_policy(), 4);
+ let cfg = SessionConfig::new();
+ let e1 = cache.produce_runtime("s1", &cfg).unwrap();
+ let e2 = cache.produce_runtime("s2", &cfg).unwrap();
+ assert!(!Arc::ptr_eq(&e1.cache_manager, &e2.cache_manager));
+ }
+
+ #[test]
+ fn per_task_pool_shares_footer_cache_but_not_env() {
+ let cache =
+ DefaultSessionRuntimeCache::new(base_producer(),
per_task_pool_policy(), 4);
+ let cfg = SessionConfig::new();
+ let e1 = cache.produce_runtime("s1", &cfg).unwrap();
+ let e2 = cache.produce_runtime("s1", &cfg).unwrap();
+ // Different runtime envs (fresh per-task pool)...
+ assert!(!Arc::ptr_eq(&e1, &e2));
+ // ...but the shared read-side state is reused: object-store registry
+ // passes through unchanged, and the inner footer (file-metadata) cache
+ // is the same instance even though the outer CacheManager wrapper is
+ // rebuilt. Do NOT compare the outer Arc<CacheManager>.
+ assert!(Arc::ptr_eq(
+ &e1.object_store_registry,
+ &e2.object_store_registry
+ ));
+ assert!(Arc::ptr_eq(
+ &e1.cache_manager.get_file_metadata_cache(),
+ &e2.cache_manager.get_file_metadata_cache(),
+ ));
+ }
+
+ #[test]
+ fn capacity_zero_disables_cache() {
+ let cache =
+ DefaultSessionRuntimeCache::new(base_producer(),
identity_policy(), 0);
+ let cfg = SessionConfig::new();
+ let e1 = cache.produce_runtime("s1", &cfg).unwrap();
+ let e2 = cache.produce_runtime("s1", &cfg).unwrap();
+ assert!(!Arc::ptr_eq(&e1.cache_manager, &e2.cache_manager));
+ }
+
+ #[test]
+ fn evicts_least_recently_used() {
+ let cache =
+ DefaultSessionRuntimeCache::new(base_producer(),
identity_policy(), 2);
+ let cfg = SessionConfig::new();
+ let s1_a = cache.produce_runtime("s1", &cfg).unwrap();
+ cache.produce_runtime("s2", &cfg).unwrap();
+ // Inserting s3 (capacity 2) evicts the least-recently-used, s1.
+ cache.produce_runtime("s3", &cfg).unwrap();
+ let s1_b = cache.produce_runtime("s1", &cfg).unwrap();
+ assert!(!Arc::ptr_eq(&s1_a.cache_manager, &s1_b.cache_manager));
+ }
+}
diff --git a/python/Cargo.lock b/python/Cargo.lock
index ca63ef385..e7d0ef3c9 100644
--- a/python/Cargo.lock
+++ b/python/Cargo.lock
@@ -548,6 +548,7 @@ dependencies = [
"futures",
"libc",
"log",
+ "lru",
"memory-stats",
"parking_lot",
"serde",
@@ -1847,7 +1848,7 @@ source =
"registry+https://github.com/rust-lang/crates.io-index";
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
- "windows-sys 0.52.0",
+ "windows-sys 0.61.2",
]
[[package]]
@@ -2617,6 +2618,15 @@ version = "0.4.33"
source = "registry+https://github.com/rust-lang/crates.io-index";
checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad"
+[[package]]
+name = "lru"
+version = "0.18.1"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "0b6180140927ee907000b0aa540091f6ea512ead4447c92b8fc35bc72788a5a6"
+dependencies = [
+ "hashbrown 0.17.1",
+]
+
[[package]]
name = "lru-slab"
version = "0.1.2"
@@ -3504,7 +3514,7 @@ dependencies = [
"errno",
"libc",
"linux-raw-sys",
- "windows-sys 0.52.0",
+ "windows-sys 0.61.2",
]
[[package]]
@@ -3784,7 +3794,7 @@ source =
"registry+https://github.com/rust-lang/crates.io-index";
checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e"
dependencies = [
"libc",
- "windows-sys 0.60.2",
+ "windows-sys 0.61.2",
]
[[package]]
@@ -3859,7 +3869,7 @@ dependencies = [
"cfg-if",
"libc",
"psm",
- "windows-sys 0.60.2",
+ "windows-sys 0.61.2",
]
[[package]]
@@ -3938,10 +3948,10 @@ source =
"registry+https://github.com/rust-lang/crates.io-index";
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
- "getrandom 0.3.4",
+ "getrandom 0.4.2",
"once_cell",
"rustix",
- "windows-sys 0.52.0",
+ "windows-sys 0.61.2",
]
[[package]]
@@ -4545,7 +4555,7 @@ version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index";
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
- "windows-sys 0.52.0",
+ "windows-sys 0.61.2",
]
[[package]]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
