datafusion_execution/
runtime_env.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Execution [`RuntimeEnv`] environment that manages access to object
19//! store, memory manager, disk manager.
20
21use crate::disk_manager::SpillingProgress;
22use crate::{
23    disk_manager::{DiskManager, DiskManagerBuilder, DiskManagerMode},
24    memory_pool::{
25        GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool,
26    },
27    object_store::{DefaultObjectStoreRegistry, ObjectStoreRegistry},
28};
29
30use crate::cache::cache_manager::{CacheManager, CacheManagerConfig};
31#[cfg(feature = "parquet_encryption")]
32use crate::parquet_encryption::{EncryptionFactory, EncryptionFactoryRegistry};
33use datafusion_common::{Result, config::ConfigEntry};
34use object_store::ObjectStore;
35use std::sync::Arc;
36use std::{
37    fmt::{Debug, Formatter},
38    num::NonZeroUsize,
39};
40use std::{path::PathBuf, time::Duration};
41use url::Url;
42
43#[derive(Clone)]
44/// Execution runtime environment that manages system resources such
45/// as memory, disk, cache and storage.
46///
47/// A [`RuntimeEnv`] can be created using [`RuntimeEnvBuilder`] and has the
48/// following resource management functionality:
49///
50/// * [`MemoryPool`]: Manage memory
51/// * [`DiskManager`]: Manage temporary files on local disk
52/// * [`CacheManager`]: Manage temporary cache data during the session lifetime
53/// * [`ObjectStoreRegistry`]: Manage mapping URLs to object store instances
54///
55/// # Example: Create default `RuntimeEnv`
56/// ```
57/// # use datafusion_execution::runtime_env::RuntimeEnv;
58/// let runtime_env = RuntimeEnv::default();
59/// ```
60///
61/// # Example: Create a `RuntimeEnv` from [`RuntimeEnvBuilder`] with a new memory pool
62/// ```
63/// # use std::sync::Arc;
64/// # use datafusion_execution::memory_pool::GreedyMemoryPool;
65/// # use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
66/// // restrict to using at most 100MB of memory
67/// let pool_size = 100 * 1024 * 1024;
68/// let runtime_env = RuntimeEnvBuilder::new()
69///     .with_memory_pool(Arc::new(GreedyMemoryPool::new(pool_size)))
70///     .build()
71///     .unwrap();
72/// ```
73pub struct RuntimeEnv {
74    /// Runtime memory management
75    pub memory_pool: Arc<dyn MemoryPool>,
76    /// Manage temporary files during query execution
77    pub disk_manager: Arc<DiskManager>,
78    /// Manage temporary cache during query execution
79    pub cache_manager: Arc<CacheManager>,
80    /// Object Store Registry
81    pub object_store_registry: Arc<dyn ObjectStoreRegistry>,
82    /// Parquet encryption factory registry
83    #[cfg(feature = "parquet_encryption")]
84    pub parquet_encryption_factory_registry: Arc<EncryptionFactoryRegistry>,
85}
86
87impl Debug for RuntimeEnv {
88    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
89        write!(f, "RuntimeEnv")
90    }
91}
92
93struct RuntimeConfigValues {
94    memory_limit: Option<String>,
95    max_temp_directory_size: Option<String>,
96    max_spill_merge_fan_in: Option<String>,
97    temp_directory: Option<String>,
98    metadata_cache_limit: Option<String>,
99    list_files_cache_limit: Option<String>,
100    list_files_cache_ttl: Option<String>,
101    file_statistics_cache_limit: Option<String>,
102}
103
104impl RuntimeConfigValues {
105    /// Creates runtime configuration entries with the provided values.
106    ///
107    /// This defines the structure and metadata for all runtime configuration
108    /// entries to avoid duplication between `RuntimeEnv::config_entries()` and
109    /// `RuntimeEnvBuilder::entries()`.
110    fn into_config_entries(self) -> Vec<ConfigEntry> {
111        let Self {
112            memory_limit,
113            max_temp_directory_size,
114            max_spill_merge_fan_in,
115            temp_directory,
116            metadata_cache_limit,
117            list_files_cache_limit,
118            list_files_cache_ttl,
119            file_statistics_cache_limit,
120        } = self;
121        vec![
122            ConfigEntry {
123                key: "datafusion.runtime.memory_limit".to_string(),
124                value: memory_limit,
125                description: "Maximum memory limit for query execution. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.",
126            },
127            ConfigEntry {
128                key: "datafusion.runtime.max_temp_directory_size".to_string(),
129                value: max_temp_directory_size,
130                description: "Maximum temporary file directory size. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.",
131            },
132            ConfigEntry {
133                key: "datafusion.runtime.max_spill_merge_fan_in".to_string(),
134                value: max_spill_merge_fan_in,
135                description: "Maximum number of spill files opened by one external merge pass. Use 0 for unlimited. Values below 2 still use 2 so a merge can make progress.",
136            },
137            ConfigEntry {
138                key: "datafusion.runtime.temp_directory".to_string(),
139                value: temp_directory,
140                description: "The path to the temporary file directory.",
141            },
142            ConfigEntry {
143                key: "datafusion.runtime.metadata_cache_limit".to_string(),
144                value: metadata_cache_limit,
145                description: "Maximum memory to use for file metadata cache such as Parquet metadata. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.",
146            },
147            ConfigEntry {
148                key: "datafusion.runtime.list_files_cache_limit".to_string(),
149                value: list_files_cache_limit,
150                description: "Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.",
151            },
152            ConfigEntry {
153                key: "datafusion.runtime.list_files_cache_ttl".to_string(),
154                value: list_files_cache_ttl,
155                description: "TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes.",
156            },
157            ConfigEntry {
158                key: "datafusion.runtime.file_statistics_cache_limit".to_string(),
159                value: file_statistics_cache_limit,
160                description: "Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.",
161            },
162        ]
163    }
164}
165
166impl RuntimeEnv {
167    /// Registers a custom `ObjectStore` to be used with a specific url.
168    /// This allows DataFusion to create external tables from urls that do not have
169    /// built in support such as `hdfs://namenode:port/...`.
170    ///
171    /// Returns the [`ObjectStore`] previously registered for this
172    /// scheme, if any.
173    ///
174    /// See [`ObjectStoreRegistry`] for more details
175    ///
176    /// # Example: Register local file system object store
177    /// ```
178    /// # use std::sync::Arc;
179    /// # use url::Url;
180    /// # use datafusion_execution::runtime_env::RuntimeEnv;
181    /// # let runtime_env = RuntimeEnv::default();
182    /// let url = Url::try_from("file://").unwrap();
183    /// let object_store = object_store::local::LocalFileSystem::new();
184    /// // register the object store with the runtime environment
185    /// runtime_env.register_object_store(&url, Arc::new(object_store));
186    /// ```
187    ///
188    /// # Example: Register remote URL object store like [Github](https://github.com)
189    /// ```
190    /// # use std::sync::Arc;
191    /// # use url::Url;
192    /// # use datafusion_execution::runtime_env::RuntimeEnv;
193    /// # let runtime_env = RuntimeEnv::default();
194    /// # // use local store for example as http feature is not enabled
195    /// # let http_store = object_store::local::LocalFileSystem::new();
196    /// // create a new object store via object_store::http::HttpBuilder;
197    /// let base_url = Url::parse("https://github.com").unwrap();
198    /// // (note this example can't depend on the http feature)
199    /// // let http_store = HttpBuilder::new()
200    /// //    .with_url(base_url.clone())
201    /// //    .build()
202    /// //    .unwrap();
203    /// // register the object store with the runtime environment
204    /// runtime_env.register_object_store(&base_url, Arc::new(http_store));
205    /// ```
206    pub fn register_object_store(
207        &self,
208        url: &Url,
209        object_store: Arc<dyn ObjectStore>,
210    ) -> Option<Arc<dyn ObjectStore>> {
211        self.object_store_registry.register_store(url, object_store)
212    }
213
214    /// Deregisters a custom `ObjectStore` previously registered for a specific url.
215    /// See [`ObjectStoreRegistry::deregister_store`] for more details.
216    pub fn deregister_object_store(&self, url: &Url) -> Result<Arc<dyn ObjectStore>> {
217        self.object_store_registry.deregister_store(url)
218    }
219
220    /// Retrieves a `ObjectStore` instance for a url by consulting the
221    /// registry. See [`ObjectStoreRegistry::get_store`] for more
222    /// details.
223    pub fn object_store(&self, url: impl AsRef<Url>) -> Result<Arc<dyn ObjectStore>> {
224        self.object_store_registry.get_store(url.as_ref())
225    }
226
227    /// Returns the current spilling progress
228    pub fn spilling_progress(&self) -> SpillingProgress {
229        self.disk_manager.spilling_progress()
230    }
231
232    /// Register an [`EncryptionFactory`] with an associated identifier that can be later
233    /// used to configure encryption when reading or writing Parquet.
234    /// If an encryption factory with the same identifier was already registered, it is replaced and returned.
235    #[cfg(feature = "parquet_encryption")]
236    pub fn register_parquet_encryption_factory(
237        &self,
238        id: &str,
239        encryption_factory: Arc<dyn EncryptionFactory>,
240    ) -> Option<Arc<dyn EncryptionFactory>> {
241        self.parquet_encryption_factory_registry
242            .register_factory(id, encryption_factory)
243    }
244
245    /// Retrieve an [`EncryptionFactory`] by its identifier
246    #[cfg(feature = "parquet_encryption")]
247    pub fn parquet_encryption_factory(
248        &self,
249        id: &str,
250    ) -> Result<Arc<dyn EncryptionFactory>> {
251        self.parquet_encryption_factory_registry.get_factory(id)
252    }
253
254    /// Returns the current runtime configuration entries
255    pub fn config_entries(&self) -> Vec<ConfigEntry> {
256        use crate::memory_pool::MemoryLimit;
257
258        /// Convert bytes to a human-readable format
259        fn format_byte_size(size: u64) -> String {
260            const GB: u64 = 1024 * 1024 * 1024;
261            const MB: u64 = 1024 * 1024;
262            const KB: u64 = 1024;
263
264            match size {
265                s if s >= GB => format!("{}G", s / GB),
266                s if s >= MB => format!("{}M", s / MB),
267                s if s >= KB => format!("{}K", s / KB),
268                s => format!("{s}"),
269            }
270        }
271
272        fn format_duration(duration: Duration) -> String {
273            let total = duration.as_secs();
274            let mins = total / 60;
275            let secs = total % 60;
276
277            format!("{mins}m{secs}s")
278        }
279
280        let memory_limit_value = match self.memory_pool.memory_limit() {
281            MemoryLimit::Finite(size) => Some(format_byte_size(
282                size.try_into()
283                    .expect("Memory limit size conversion failed"),
284            )),
285            MemoryLimit::Infinite => Some("unlimited".to_string()),
286            MemoryLimit::Unknown => None,
287        };
288
289        let max_temp_dir_size = self.disk_manager.max_temp_directory_size();
290        let max_temp_dir_value = format_byte_size(max_temp_dir_size);
291        let max_spill_merge_fan_in =
292            self.disk_manager.max_spill_merge_fan_in().to_string();
293
294        let temp_paths = self.disk_manager.temp_dir_paths();
295        let temp_dir_value = if temp_paths.is_empty() {
296            None
297        } else {
298            Some(
299                temp_paths
300                    .iter()
301                    .map(|p| p.display().to_string())
302                    .collect::<Vec<_>>()
303                    .join(","),
304            )
305        };
306
307        let metadata_cache_limit = self.cache_manager.get_metadata_cache_limit();
308        let metadata_cache_value = format_byte_size(
309            metadata_cache_limit
310                .try_into()
311                .expect("Metadata cache size conversion failed"),
312        );
313
314        let list_files_cache_limit = self.cache_manager.get_list_files_cache_limit();
315        let list_files_cache_value = format_byte_size(
316            list_files_cache_limit
317                .try_into()
318                .expect("List files cache size conversion failed"),
319        );
320
321        let list_files_cache_ttl = self
322            .cache_manager
323            .get_list_files_cache_ttl()
324            .map(format_duration);
325
326        let file_statistics_cache_limit =
327            self.cache_manager.get_file_statistic_cache_limit();
328        let file_statistics_cache_value = format_byte_size(
329            file_statistics_cache_limit
330                .try_into()
331                .expect("File statistics cache size conversion failed"),
332        );
333
334        RuntimeConfigValues {
335            memory_limit: memory_limit_value,
336            max_temp_directory_size: Some(max_temp_dir_value),
337            max_spill_merge_fan_in: Some(max_spill_merge_fan_in),
338            temp_directory: temp_dir_value,
339            metadata_cache_limit: Some(metadata_cache_value),
340            list_files_cache_limit: Some(list_files_cache_value),
341            list_files_cache_ttl,
342            file_statistics_cache_limit: Some(file_statistics_cache_value),
343        }
344        .into_config_entries()
345    }
346}
347
348impl Default for RuntimeEnv {
349    fn default() -> Self {
350        RuntimeEnvBuilder::new().build().unwrap()
351    }
352}
353
354/// Execution runtime configuration builder.
355///
356/// See example on [`RuntimeEnv`]
357#[derive(Clone)]
358pub struct RuntimeEnvBuilder {
359    /// DiskManager to manage temporary disk file usage
360    pub disk_manager: Option<Arc<DiskManager>>,
361    /// DiskManager builder to manager temporary disk file usage
362    pub disk_manager_builder: Option<DiskManagerBuilder>,
363    /// [`MemoryPool`] from which to allocate memory
364    ///
365    /// Defaults to using an [`UnboundedMemoryPool`] if `None`
366    pub memory_pool: Option<Arc<dyn MemoryPool>>,
367    /// CacheManager to manage cache data
368    pub cache_manager: CacheManagerConfig,
369    /// ObjectStoreRegistry to get object store based on url
370    pub object_store_registry: Arc<dyn ObjectStoreRegistry>,
371    /// Parquet encryption factory registry
372    #[cfg(feature = "parquet_encryption")]
373    pub parquet_encryption_factory_registry: Arc<EncryptionFactoryRegistry>,
374}
375
376impl Default for RuntimeEnvBuilder {
377    fn default() -> Self {
378        Self::new()
379    }
380}
381
382impl RuntimeEnvBuilder {
383    /// New with default values
384    pub fn new() -> Self {
385        Self {
386            disk_manager: Default::default(),
387            disk_manager_builder: Default::default(),
388            memory_pool: Default::default(),
389            cache_manager: Default::default(),
390            object_store_registry: Arc::new(DefaultObjectStoreRegistry::default()),
391            #[cfg(feature = "parquet_encryption")]
392            parquet_encryption_factory_registry: Default::default(),
393        }
394    }
395
396    /// Customize the disk manager builder
397    pub fn with_disk_manager_builder(mut self, disk_manager: DiskManagerBuilder) -> Self {
398        self.disk_manager_builder = Some(disk_manager);
399        self
400    }
401
402    /// Customize memory policy
403    pub fn with_memory_pool(mut self, memory_pool: Arc<dyn MemoryPool>) -> Self {
404        self.memory_pool = Some(memory_pool);
405        self
406    }
407
408    /// Customize cache policy
409    pub fn with_cache_manager(mut self, cache_manager: CacheManagerConfig) -> Self {
410        self.cache_manager = cache_manager;
411        self
412    }
413
414    /// Customize object store registry
415    pub fn with_object_store_registry(
416        mut self,
417        object_store_registry: Arc<dyn ObjectStoreRegistry>,
418    ) -> Self {
419        self.object_store_registry = object_store_registry;
420        self
421    }
422
423    /// Specify the total memory to use while running the DataFusion
424    /// plan to `max_memory * memory_fraction` in bytes.
425    ///
426    /// This defaults to using [`GreedyMemoryPool`] wrapped in the
427    /// [`TrackConsumersPool`] with a maximum of 5 consumers.
428    ///
429    /// Note DataFusion does not yet respect this limit in all cases.
430    pub fn with_memory_limit(self, max_memory: usize, memory_fraction: f64) -> Self {
431        let pool_size = (max_memory as f64 * memory_fraction) as usize;
432        self.with_memory_pool(Arc::new(TrackConsumersPool::new(
433            GreedyMemoryPool::new(pool_size),
434            NonZeroUsize::new(5).unwrap(),
435        )))
436    }
437
438    /// Use the specified path to create any needed temporary files
439    pub fn with_temp_file_path(mut self, path: impl Into<PathBuf>) -> Self {
440        let builder = self.disk_manager_builder.take().unwrap_or_default();
441        self.with_disk_manager_builder(
442            builder.with_mode(DiskManagerMode::Directories(vec![path.into()])),
443        )
444    }
445
446    /// Specify a limit on the size of the temporary file directory in bytes
447    pub fn with_max_temp_directory_size(mut self, size: u64) -> Self {
448        let builder = self.disk_manager_builder.take().unwrap_or_default();
449        self.with_disk_manager_builder(builder.with_max_temp_directory_size(size))
450    }
451
452    /// Limit the number of spill files opened by one external merge pass.
453    ///
454    /// A value of 0 means unlimited.
455    pub fn with_max_spill_merge_fan_in(mut self, fan_in: usize) -> Self {
456        let builder = self.disk_manager_builder.take().unwrap_or_default();
457        self.with_disk_manager_builder(builder.with_max_spill_merge_fan_in(fan_in))
458    }
459
460    /// Specify the limit of the file-embedded metadata cache, in bytes.
461    pub fn with_metadata_cache_limit(mut self, limit: usize) -> Self {
462        self.cache_manager = self.cache_manager.with_metadata_cache_limit(limit);
463        self
464    }
465
466    /// Specifies the memory limit for the object list cache, in bytes.
467    pub fn with_object_list_cache_limit(mut self, limit: usize) -> Self {
468        self.cache_manager = self.cache_manager.with_list_files_cache_limit(limit);
469        self
470    }
471
472    /// Specifies the duration entries in the object list cache will be considered valid.
473    pub fn with_object_list_cache_ttl(mut self, ttl: Option<Duration>) -> Self {
474        self.cache_manager = self.cache_manager.with_list_files_cache_ttl(ttl);
475        self
476    }
477
478    pub fn with_file_statistics_cache_limit(mut self, limit: usize) -> Self {
479        self.cache_manager = self.cache_manager.with_file_statistics_cache_limit(limit);
480        self
481    }
482
483    /// Build a RuntimeEnv
484    pub fn build(self) -> Result<RuntimeEnv> {
485        let Self {
486            disk_manager,
487            disk_manager_builder,
488            memory_pool,
489            cache_manager,
490            object_store_registry,
491            #[cfg(feature = "parquet_encryption")]
492            parquet_encryption_factory_registry,
493        } = self;
494        let memory_pool =
495            memory_pool.unwrap_or_else(|| Arc::new(UnboundedMemoryPool::default()));
496
497        let disk_manager: Arc<DiskManager> = match (disk_manager, disk_manager_builder) {
498            (_, Some(builder)) => Arc::new(builder.build()?),
499            (Some(manager), None) => manager,
500            (None, None) => Arc::new(DiskManagerBuilder::default().build()?),
501        };
502
503        Ok(RuntimeEnv {
504            memory_pool,
505            disk_manager,
506            cache_manager: CacheManager::try_new(&cache_manager)?,
507            object_store_registry,
508            #[cfg(feature = "parquet_encryption")]
509            parquet_encryption_factory_registry,
510        })
511    }
512
513    /// Convenience method to create a new `Arc<RuntimeEnv>`
514    pub fn build_arc(self) -> Result<Arc<RuntimeEnv>> {
515        self.build().map(Arc::new)
516    }
517
518    /// Create a new RuntimeEnvBuilder from an existing RuntimeEnv
519    pub fn from_runtime_env(runtime_env: &RuntimeEnv) -> Self {
520        let cache_config = CacheManagerConfig {
521            file_statistics_cache: runtime_env.cache_manager.get_file_statistic_cache(),
522            file_statistics_cache_limit: runtime_env
523                .cache_manager
524                .get_file_statistic_cache_limit(),
525            list_files_cache: runtime_env.cache_manager.get_list_files_cache(),
526            list_files_cache_limit: runtime_env
527                .cache_manager
528                .get_list_files_cache_limit(),
529            list_files_cache_ttl: runtime_env.cache_manager.get_list_files_cache_ttl(),
530            file_metadata_cache: Some(
531                runtime_env.cache_manager.get_file_metadata_cache(),
532            ),
533            metadata_cache_limit: runtime_env.cache_manager.get_metadata_cache_limit(),
534        };
535
536        Self {
537            disk_manager: Some(Arc::clone(&runtime_env.disk_manager)),
538            disk_manager_builder: None,
539            memory_pool: Some(Arc::clone(&runtime_env.memory_pool)),
540            cache_manager: cache_config,
541            object_store_registry: Arc::clone(&runtime_env.object_store_registry),
542            #[cfg(feature = "parquet_encryption")]
543            parquet_encryption_factory_registry: Arc::clone(
544                &runtime_env.parquet_encryption_factory_registry,
545            ),
546        }
547    }
548
549    /// Returns a list of all available runtime configurations with their current values and descriptions
550    pub fn entries(&self) -> Vec<ConfigEntry> {
551        RuntimeConfigValues {
552            memory_limit: None,
553            max_temp_directory_size: Some("100G".to_string()),
554            max_spill_merge_fan_in: Some("0".to_string()),
555            temp_directory: None,
556            metadata_cache_limit: Some("50M".to_owned()),
557            list_files_cache_limit: Some("1M".to_owned()),
558            list_files_cache_ttl: None,
559            file_statistics_cache_limit: Some("20M".to_owned()),
560        }
561        .into_config_entries()
562    }
563
564    /// Generate documentation that can be included in the user guide
565    pub fn generate_config_markdown() -> String {
566        use std::fmt::Write as _;
567
568        let s = Self::default();
569
570        let mut docs = "| key | default | description |\n".to_string();
571        docs += "|-----|---------|-------------|\n";
572        let mut entries = s.entries();
573        entries.sort_unstable_by(|a, b| a.key.cmp(&b.key));
574
575        for entry in &entries {
576            let _ = writeln!(
577                &mut docs,
578                "| {} | {} | {} |",
579                entry.key,
580                entry.value.as_deref().unwrap_or("NULL"),
581                entry.description
582            );
583        }
584        docs
585    }
586}
datafusion_execution/runtime_env.rs

datafusion_execution/
runtime_env.rs
