Minor: Add `RuntimeEnvBuilder::build_arc() (#12213)

Author: alamb

datafusion/core/src/execution/context/mod.rs

@@ -219,13 +219,13 @@ where
 /// // configure a memory limit of 1GB with 20%  slop
 ///  let runtime_env = RuntimeEnvBuilder::new()
 ///     .with_memory_limit(1024 * 1024 * 1024, 0.80)
-///     .build()
+///     .build_arc()
 ///     .unwrap();
 ///
 /// // Create a SessionState using the config and runtime_env
 /// let state = SessionStateBuilder::new()
 ///   .with_config(config)
-///   .with_runtime_env(Arc::new(runtime_env))
+///   .with_runtime_env(runtime_env)
 ///   // include support for built in functions and configurations
 ///   .with_default_features()
 ///   .build();
@@ -1758,7 +1758,7 @@ mod tests {
         let path = path.join("tests/tpch-csv");
         let url = format!("file://{}", path.display());
 
-        let runtime = Arc::new(RuntimeEnvBuilder::new().build()?);
+        let runtime = RuntimeEnvBuilder::new().build_arc()?;
         let cfg = SessionConfig::new()
             .set_str("datafusion.catalog.location", url.as_str())
             .set_str("datafusion.catalog.format", "CSV")

datafusion/core/tests/fuzz_cases/sort_fuzz.rs

@@ -136,12 +136,10 @@ impl SortTest {
                     .sort_spill_reservation_bytes,
             );
 
-            let runtime = Arc::new(
-                RuntimeEnvBuilder::new()
-                    .with_memory_pool(Arc::new(GreedyMemoryPool::new(pool_size)))
-                    .build()
-                    .unwrap(),
-            );
+            let runtime = RuntimeEnvBuilder::new()
+                .with_memory_pool(Arc::new(GreedyMemoryPool::new(pool_size)))
+                .build_arc()
+                .unwrap();
             SessionContext::new_with_config_rt(session_config, runtime)
         } else {
             SessionContext::new_with_config(session_config)

datafusion/core/tests/memory_limit/mod.rs

@@ -509,21 +509,20 @@ impl TestCase {
 
         let table = scenario.table();
 
-        let rt_config = RuntimeEnvBuilder::new()
+        let mut builder = RuntimeEnvBuilder::new()
             // disk manager setting controls the spilling
             .with_disk_manager(disk_manager_config)
             .with_memory_limit(memory_limit, MEMORY_FRACTION);
 
-        let runtime = if let Some(pool) = memory_pool {
-            rt_config.with_memory_pool(pool).build().unwrap()
-        } else {
-            rt_config.build().unwrap()
+        if let Some(pool) = memory_pool {
+            builder = builder.with_memory_pool(pool);
         };
+        let runtime = builder.build_arc().unwrap();
 
         // Configure execution
         let builder = SessionStateBuilder::new()
             .with_config(config)
-            .with_runtime_env(Arc::new(runtime))
+            .with_runtime_env(runtime)
             .with_default_features();
         let builder = match scenario.rules() {
             Some(rules) => builder.with_physical_optimizer_rules(rules),

datafusion/core/tests/parquet/file_statistics.rs

@@ -197,12 +197,11 @@ fn get_cache_runtime_state() -> (
         .with_files_statistics_cache(Some(file_static_cache.clone()))
         .with_list_files_cache(Some(list_file_cache.clone()));
 
-    let rt = Arc::new(
-        RuntimeEnvBuilder::new()
-            .with_cache_manager(cache_config)
-            .build()
-            .expect("could not build runtime environment"),
-    );
+    let rt = RuntimeEnvBuilder::new()
+        .with_cache_manager(cache_config)
+        .build_arc()
+        .expect("could not build runtime environment");
+
     let state = SessionContext::new_with_config_rt(SessionConfig::default(), rt).state();
 
     (file_static_cache, list_file_cache, state)

datafusion/execution/src/runtime_env.rs

@@ -246,4 +246,9 @@ impl RuntimeEnvBuilder {
             object_store_registry: self.object_store_registry,
         })
     }
+
+    /// Convenience method to create a new `Arc<RuntimeEnv>`
+    pub fn build_arc(self) -> Result<Arc<RuntimeEnv>> {
+        self.build().map(Arc::new)
+    }
 }

datafusion/execution/src/task.rs

@@ -58,7 +58,7 @@ pub struct TaskContext {
 impl Default for TaskContext {
     fn default() -> Self {
         let runtime = RuntimeEnvBuilder::new()
-            .build()
+            .build_arc()
             .expect("default runtime created successfully");
 
         // Create a default task context, mostly useful for testing
@@ -69,7 +69,7 @@ impl Default for TaskContext {
             scalar_functions: HashMap::new(),
             aggregate_functions: HashMap::new(),
             window_functions: HashMap::new(),
-            runtime: Arc::new(runtime),
+            runtime,
         }
     }
 }

datafusion/physical-plan/src/aggregates/mod.rs

@@ -1331,12 +1331,10 @@ mod tests {
 
     fn new_spill_ctx(batch_size: usize, max_memory: usize) -> Arc<TaskContext> {
         let session_config = SessionConfig::new().with_batch_size(batch_size);
-        let runtime = Arc::new(
-            RuntimeEnvBuilder::default()
-                .with_memory_pool(Arc::new(FairSpillPool::new(max_memory)))
-                .build()
-                .unwrap(),
-        );
+        let runtime = RuntimeEnvBuilder::default()
+            .with_memory_pool(Arc::new(FairSpillPool::new(max_memory)))
+            .build_arc()
+            .unwrap();
         let task_ctx = TaskContext::default()
             .with_session_config(session_config)
             .with_runtime(runtime);
@@ -1815,11 +1813,9 @@ mod tests {
         let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(true));
         let input_schema = input.schema();
 
-        let runtime = Arc::new(
-            RuntimeEnvBuilder::default()
-                .with_memory_limit(1, 1.0)
-                .build()?,
-        );
+        let runtime = RuntimeEnvBuilder::default()
+            .with_memory_limit(1, 1.0)
+            .build_arc()?;
         let task_ctx = TaskContext::default().with_runtime(runtime);
         let task_ctx = Arc::new(task_ctx);
 

datafusion/physical-plan/src/joins/cross_join.rs

@@ -673,11 +673,9 @@ mod tests {
 
     #[tokio::test]
     async fn test_overallocation() -> Result<()> {
-        let runtime = Arc::new(
-            RuntimeEnvBuilder::new()
-                .with_memory_limit(100, 1.0)
-                .build()?,
-        );
+        let runtime = RuntimeEnvBuilder::new()
+            .with_memory_limit(100, 1.0)
+            .build_arc()?;
         let task_ctx = TaskContext::default().with_runtime(runtime);
         let task_ctx = Arc::new(task_ctx);
 

datafusion/physical-plan/src/joins/hash_join.rs

@@ -3798,11 +3798,9 @@ mod tests {
         ];
 
         for join_type in join_types {
-            let runtime = Arc::new(
-                RuntimeEnvBuilder::new()
-                    .with_memory_limit(100, 1.0)
-                    .build()?,
-            );
+            let runtime = RuntimeEnvBuilder::new()
+                .with_memory_limit(100, 1.0)
+                .build_arc()?;
             let task_ctx = TaskContext::default().with_runtime(runtime);
             let task_ctx = Arc::new(task_ctx);
 
@@ -3874,11 +3872,9 @@ mod tests {
         ];
 
         for join_type in join_types {
-            let runtime = Arc::new(
-                RuntimeEnvBuilder::new()
-                    .with_memory_limit(100, 1.0)
-                    .build()?,
-            );
+            let runtime = RuntimeEnvBuilder::new()
+                .with_memory_limit(100, 1.0)
+                .build_arc()?;
             let session_config = SessionConfig::default().with_batch_size(50);
             let task_ctx = TaskContext::default()
                 .with_session_config(session_config)

datafusion/physical-plan/src/joins/nested_loop_join.rs

@@ -1019,11 +1019,9 @@ mod tests {
         ];
 
         for join_type in join_types {
-            let runtime = Arc::new(
-                RuntimeEnvBuilder::new()
-                    .with_memory_limit(100, 1.0)
-                    .build()?,
-            );
+            let runtime = RuntimeEnvBuilder::new()
+                .with_memory_limit(100, 1.0)
+                .build_arc()?;
             let task_ctx = TaskContext::default().with_runtime(runtime);
             let task_ctx = Arc::new(task_ctx);