## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/execution/src/runtime_env.rs

    Line
    Count
    Source
    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
    21
    #[allow(deprecated)]
    22
    use crate::disk_manager::DiskManagerConfig;
    23
    use crate::{
    24
        disk_manager::{DiskManager, DiskManagerBuilder, DiskManagerMode},
    25
        memory_pool::{
    26
            GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool,
    27
        },
    28
        object_store::{DefaultObjectStoreRegistry, ObjectStoreRegistry},
    29
    };
    30
    31
    use crate::cache::cache_manager::{CacheManager, CacheManagerConfig};
    32
    #[cfg(feature = "parquet_encryption")]
    33
    use crate::parquet_encryption::{EncryptionFactory, EncryptionFactoryRegistry};
    34
    use datafusion_common::{config::ConfigEntry, Result};
    35
    use object_store::ObjectStore;
    36
    use std::path::PathBuf;
    37
    use std::sync::Arc;
    38
    use std::{
    39
        fmt::{Debug, Formatter},
    40
        num::NonZeroUsize,
    41
    };
    42
    use url::Url;
    43
    44
    #[derive(Clone)]
    45
    /// Execution runtime environment that manages system resources such
    46
    /// as memory, disk, cache and storage.
    47
    ///
    48
    /// A [`RuntimeEnv`] can be created using [`RuntimeEnvBuilder`] and has the
    49
    /// following resource management functionality:
    50
    ///
    51
    /// * [`MemoryPool`]: Manage memory
    52
    /// * [`DiskManager`]: Manage temporary files on local disk
    53
    /// * [`CacheManager`]: Manage temporary cache data during the session lifetime
    54
    /// * [`ObjectStoreRegistry`]: Manage mapping URLs to object store instances
    55
    ///
    56
    /// # Example: Create default `RuntimeEnv`
    57
    /// ```
    58
    /// # use datafusion_execution::runtime_env::RuntimeEnv;
    59
    /// let runtime_env = RuntimeEnv::default();
    60
    /// ```
    61
    ///
    62
    /// # Example: Create a `RuntimeEnv` from [`RuntimeEnvBuilder`] with a new memory pool
    63
    /// ```
    64
    /// # use std::sync::Arc;
    65
    /// # use datafusion_execution::memory_pool::GreedyMemoryPool;
    66
    /// # use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
    67
    /// // restrict to using at most 100MB of memory
    68
    /// let pool_size = 100 * 1024 * 1024;
    69
    /// let runtime_env = RuntimeEnvBuilder::new()
    70
    ///     .with_memory_pool(Arc::new(GreedyMemoryPool::new(pool_size)))
    71
    ///     .build()
    72
    ///     .unwrap();
    73
    /// ```
    74
    pub struct RuntimeEnv {
    75
        /// Runtime memory management
    76
        pub memory_pool: Arc<dyn MemoryPool>,
    77
        /// Manage temporary files during query execution
    78
        pub disk_manager: Arc<DiskManager>,
    79
        /// Manage temporary cache during query execution
    80
        pub cache_manager: Arc<CacheManager>,
    81
        /// Object Store Registry
    82
        pub object_store_registry: Arc<dyn ObjectStoreRegistry>,
    83
        /// Parquet encryption factory registry
    84
        #[cfg(feature = "parquet_encryption")]
    85
        pub parquet_encryption_factory_registry: Arc<EncryptionFactoryRegistry>,
    86
    }
    87
    88
    impl Debug for RuntimeEnv {
    89
        fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
    90
            write!(f, "RuntimeEnv")
    91
        }
    92
    }
    93
    94
    impl RuntimeEnv {
    95
        /// Registers a custom `ObjectStore` to be used with a specific url.
    96
        /// This allows DataFusion to create external tables from urls that do not have
    97
        /// built in support such as `hdfs://namenode:port/...`.
    98
        ///
    99
        /// Returns the [`ObjectStore`] previously registered for this
    100
        /// scheme, if any.
    101
        ///
    102
        /// See [`ObjectStoreRegistry`] for more details
    103
        ///
    104
        /// # Example: Register local file system object store
    105
        /// ```
    106
        /// # use std::sync::Arc;
    107
        /// # use url::Url;
    108
        /// # use datafusion_execution::runtime_env::RuntimeEnv;
    109
        /// # let runtime_env = RuntimeEnv::default();
    110
        /// let url = Url::try_from("file://").unwrap();
    111
        /// let object_store = object_store::local::LocalFileSystem::new();
    112
        /// // register the object store with the runtime environment
    113
        /// runtime_env.register_object_store(&url, Arc::new(object_store));
    114
        /// ```
    115
        ///
    116
        /// # Example: Register remote URL object store like [Github](https://github.com)
    117
        /// ```
    118
        /// # use std::sync::Arc;
    119
        /// # use url::Url;
    120
        /// # use datafusion_execution::runtime_env::RuntimeEnv;
    121
        /// # let runtime_env = RuntimeEnv::default();
    122
        /// # // use local store for example as http feature is not enabled
    123
        /// # let http_store = object_store::local::LocalFileSystem::new();
    124
        /// // create a new object store via object_store::http::HttpBuilder;
    125
        /// let base_url = Url::parse("https://github.com").unwrap();
    126
        /// // (note this example can't depend on the http feature)
    127
        /// // let http_store = HttpBuilder::new()
    128
        /// //    .with_url(base_url.clone())
    129
        /// //    .build()
    130
        /// //    .unwrap();
    131
        /// // register the object store with the runtime environment
    132
        /// runtime_env.register_object_store(&base_url, Arc::new(http_store));
    133
        /// ```
    134
        pub fn register_object_store(
    135
            &self,
    136
            url: &Url,
    137
            object_store: Arc<dyn ObjectStore>,
    138
        ) -> Option<Arc<dyn ObjectStore>> {
    139
            self.object_store_registry.register_store(url, object_store)
    140
        }
    141
    142
        /// Deregisters a custom `ObjectStore` previously registered for a specific url.
    143
        /// See [`ObjectStoreRegistry::deregister_store`] for more details.
    144
        pub fn deregister_object_store(&self, url: &Url) -> Result<Arc<dyn ObjectStore>> {
    145
            self.object_store_registry.deregister_store(url)
    146
        }
    147
    148
        /// Retrieves a `ObjectStore` instance for a url by consulting the
    149
        /// registry. See [`ObjectStoreRegistry::get_store`] for more
    150
        /// details.
    151
    0
        pub fn object_store(&self, url: impl AsRef<Url>) -> Result<Arc<dyn ObjectStore>> {
    152
    0
            self.object_store_registry.get_store(url.as_ref())
    153
    0
        }
    154
    155
        /// Register an [`EncryptionFactory`] with an associated identifier that can be later
    156
        /// used to configure encryption when reading or writing Parquet.
    157
        /// If an encryption factory with the same identifier was already registered, it is replaced and returned.
    158
        #[cfg(feature = "parquet_encryption")]
    159
        pub fn register_parquet_encryption_factory(
    160
            &self,
    161
            id: &str,
    162
            encryption_factory: Arc<dyn EncryptionFactory>,
    163
        ) -> Option<Arc<dyn EncryptionFactory>> {
    164
            self.parquet_encryption_factory_registry
    165
                .register_factory(id, encryption_factory)
    166
        }
    167
    168
        /// Retrieve an [`EncryptionFactory`] by its identifier
    169
        #[cfg(feature = "parquet_encryption")]
    170
        pub fn parquet_encryption_factory(
    171
            &self,
    172
            id: &str,
    173
        ) -> Result<Arc<dyn EncryptionFactory>> {
    174
            self.parquet_encryption_factory_registry.get_factory(id)
    175
        }
    176
    }
    177
    178
    impl Default for RuntimeEnv {
    179
        fn default() -> Self {
    180
            RuntimeEnvBuilder::new().build().unwrap()
    181
        }
    182
    }
    183
    184
    /// Execution runtime configuration builder.
    185
    ///
    186
    /// See example on [`RuntimeEnv`]
    187
    #[derive(Clone)]
    188
    pub struct RuntimeEnvBuilder {
    189
        #[allow(deprecated)]
    190
        /// DiskManager to manage temporary disk file usage
    191
        pub disk_manager: DiskManagerConfig,
    192
        /// DiskManager builder to manager temporary disk file usage
    193
        pub disk_manager_builder: Option<DiskManagerBuilder>,
    194
        /// [`MemoryPool`] from which to allocate memory
    195
        ///
    196
        /// Defaults to using an [`UnboundedMemoryPool`] if `None`
    197
        pub memory_pool: Option<Arc<dyn MemoryPool>>,
    198
        /// CacheManager to manage cache data
    199
        pub cache_manager: CacheManagerConfig,
    200
        /// ObjectStoreRegistry to get object store based on url
    201
        pub object_store_registry: Arc<dyn ObjectStoreRegistry>,
    202
        /// Parquet encryption factory registry
    203
        #[cfg(feature = "parquet_encryption")]
    204
        pub parquet_encryption_factory_registry: Arc<EncryptionFactoryRegistry>,
    205
    }
    206
    207
    impl Default for RuntimeEnvBuilder {
    208
        fn default() -> Self {
    209
            Self::new()
    210
        }
    211
    }
    212
    213
    impl RuntimeEnvBuilder {
    214
        /// New with default values
    215
        pub fn new() -> Self {
    216
            Self {
    217
                disk_manager: Default::default(),
    218
                disk_manager_builder: Default::default(),
    219
                memory_pool: Default::default(),
    220
                cache_manager: Default::default(),
    221
                object_store_registry: Arc::new(DefaultObjectStoreRegistry::default()),
    222
                #[cfg(feature = "parquet_encryption")]
    223
                parquet_encryption_factory_registry: Default::default(),
    224
            }
    225
        }
    226
    227
        #[allow(deprecated)]
    228
        #[deprecated(since = "48.0.0", note = "Use with_disk_manager_builder instead")]
    229
        /// Customize disk manager
    230
        pub fn with_disk_manager(mut self, disk_manager: DiskManagerConfig) -> Self {
    231
            self.disk_manager = disk_manager;
    232
            self
    233
        }
    234
    235
        /// Customize the disk manager builder
    236
        pub fn with_disk_manager_builder(mut self, disk_manager: DiskManagerBuilder) -> Self {
    237
            self.disk_manager_builder = Some(disk_manager);
    238
            self
    239
        }
    240
    241
        /// Customize memory policy
    242
        pub fn with_memory_pool(mut self, memory_pool: Arc<dyn MemoryPool>) -> Self {
    243
            self.memory_pool = Some(memory_pool);
    244
            self
    245
        }
    246
    247
        /// Customize cache policy
    248
        pub fn with_cache_manager(mut self, cache_manager: CacheManagerConfig) -> Self {
    249
            self.cache_manager = cache_manager;
    250
            self
    251
        }
    252
    253
        /// Customize object store registry
    254
        pub fn with_object_store_registry(
    255
            mut self,
    256
            object_store_registry: Arc<dyn ObjectStoreRegistry>,
    257
        ) -> Self {
    258
            self.object_store_registry = object_store_registry;
    259
            self
    260
        }
    261
    262
        /// Specify the total memory to use while running the DataFusion
    263
        /// plan to `max_memory * memory_fraction` in bytes.
    264
        ///
    265
        /// This defaults to using [`GreedyMemoryPool`] wrapped in the
    266
        /// [`TrackConsumersPool`] with a maximum of 5 consumers.
    267
        ///
    268
        /// Note DataFusion does not yet respect this limit in all cases.
    269
        pub fn with_memory_limit(self, max_memory: usize, memory_fraction: f64) -> Self {
    270
            let pool_size = (max_memory as f64 * memory_fraction) as usize;
    271
            self.with_memory_pool(Arc::new(TrackConsumersPool::new(
    272
                GreedyMemoryPool::new(pool_size),
    273
                NonZeroUsize::new(5).unwrap(),
    274
            )))
    275
        }
    276
    277
        /// Use the specified path to create any needed temporary files
    278
    0
        pub fn with_temp_file_path(mut self, path: impl Into<PathBuf>) -> Self {
    279
    0
            let builder = self.disk_manager_builder.take().unwrap_or_default();
    280
    0
            self.with_disk_manager_builder(
    281
    0
                builder.with_mode(DiskManagerMode::Directories(vec![path.into()])),
    282
            )
    283
    0
        }
    284
    285
        /// Specify a limit on the size of the temporary file directory in bytes
    286
        pub fn with_max_temp_directory_size(mut self, size: u64) -> Self {
    287
            let builder = self.disk_manager_builder.take().unwrap_or_default();
    288
            self.with_disk_manager_builder(builder.with_max_temp_directory_size(size))
    289
        }
    290
    291
        /// Specify the limit of the file-embedded metadata cache, in bytes.
    292
        pub fn with_metadata_cache_limit(mut self, limit: usize) -> Self {
    293
            self.cache_manager = self.cache_manager.with_metadata_cache_limit(limit);
    294
            self
    295
        }
    296
    297
        /// Build a RuntimeEnv
    298
        pub fn build(self) -> Result<RuntimeEnv> {
    299
            let Self {
    300
                disk_manager,
    301
                disk_manager_builder,
    302
                memory_pool,
    303
                cache_manager,
    304
                object_store_registry,
    305
                #[cfg(feature = "parquet_encryption")]
    306
                parquet_encryption_factory_registry,
    307
            } = self;
    308
            let memory_pool =
    309
    0
                memory_pool.unwrap_or_else(|| Arc::new(UnboundedMemoryPool::default()));
    310
    311
            Ok(RuntimeEnv {
    312
                memory_pool,
    313
                disk_manager: if let Some(builder) = disk_manager_builder {
    314
                    Arc::new(builder.build()?)
    315
                } else {
    316
                    #[allow(deprecated)]
    317
                    DiskManager::try_new(disk_manager)?
    318
                },
    319
                cache_manager: CacheManager::try_new(&cache_manager)?,
    320
                object_store_registry,
    321
                #[cfg(feature = "parquet_encryption")]
    322
                parquet_encryption_factory_registry,
    323
            })
    324
        }
    325
    326
        /// Convenience method to create a new `Arc<RuntimeEnv>`
    327
        pub fn build_arc(self) -> Result<Arc<RuntimeEnv>> {
    328
            self.build().map(Arc::new)
    329
        }
    330
    331
        /// Create a new RuntimeEnvBuilder from an existing RuntimeEnv
    332
        pub fn from_runtime_env(runtime_env: &RuntimeEnv) -> Self {
    333
            let cache_config = CacheManagerConfig {
    334
                table_files_statistics_cache: runtime_env
    335
                    .cache_manager
    336
                    .get_file_statistic_cache(),
    337
                list_files_cache: runtime_env.cache_manager.get_list_files_cache(),
    338
                file_metadata_cache: Some(
    339
                    runtime_env.cache_manager.get_file_metadata_cache(),
    340
                ),
    341
                metadata_cache_limit: runtime_env.cache_manager.get_metadata_cache_limit(),
    342
            };
    343
    344
            Self {
    345
                #[allow(deprecated)]
    346
                disk_manager: DiskManagerConfig::Existing(Arc::clone(
    347
                    &runtime_env.disk_manager,
    348
                )),
    349
                disk_manager_builder: None,
    350
                memory_pool: Some(Arc::clone(&runtime_env.memory_pool)),
    351
                cache_manager: cache_config,
    352
                object_store_registry: Arc::clone(&runtime_env.object_store_registry),
    353
                #[cfg(feature = "parquet_encryption")]
    354
                parquet_encryption_factory_registry: Arc::clone(
    355
                    &runtime_env.parquet_encryption_factory_registry,
    356
                ),
    357
            }
    358
        }
    359
    360
        /// Returns a list of all available runtime configurations with their current values and descriptions
    361
        pub fn entries(&self) -> Vec<ConfigEntry> {
    362
            vec![
    363
                ConfigEntry {
    364
                    key: "datafusion.runtime.memory_limit".to_string(),
    365
                    value: None, // Default is system-dependent
    366
                    description: "Maximum memory limit for query execution. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes). Example: '2G' for 2 gigabytes.",
    367
                },
    368
                ConfigEntry {
    369
                    key: "datafusion.runtime.max_temp_directory_size".to_string(),
    370
                    value: Some("100G".to_string()),
    371
                    description: "Maximum temporary file directory size. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes). Example: '2G' for 2 gigabytes.",
    372
                },
    373
                ConfigEntry {
    374
                    key: "datafusion.runtime.temp_directory".to_string(),
    375
                    value: None, // Default is system-dependent
    376
                    description: "The path to the temporary file directory.",
    377
                },
    378
                ConfigEntry {
    379
                    key: "datafusion.runtime.metadata_cache_limit".to_string(),
    380
                    value: Some("50M".to_owned()),
    381
                    description: "Maximum memory to use for file metadata cache such as Parquet metadata. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes). Example: '2G' for 2 gigabytes.",
    382
                }
    383
            ]
    384
        }
    385
    386
        /// Generate documentation that can be included in the user guide
    387
        pub fn generate_config_markdown() -> String {
    388
            use std::fmt::Write as _;
    389
    390
            let s = Self::default();
    391
    392
            let mut docs = "| key | default | description |\n".to_string();
    393
            docs += "|-----|---------|-------------|\n";
    394
            let mut entries = s.entries();
    395
    0
            entries.sort_unstable_by(|a, b| a.key.cmp(&b.key));
    396
    397
            for entry in &entries {
    398
                let _ = writeln!(
    399
                    &mut docs,
    400
                    "| {} | {} | {} |",
    401
                    entry.key,
    402
                    entry.value.as_deref().unwrap_or("NULL"),
    403
                    entry.description
    404
                );
    405
            }
    406
            docs
    407
        }
    408
    }