From 322967ea2624f7b52645555825f01b0d199cdf81 Mon Sep 17 00:00:00 2001 From: theirix Date: Sat, 24 Aug 2024 13:06:49 +0100 Subject: [PATCH 1/2] Allow to build RuntimeEnv from RuntimeConfig --- datafusion/core/tests/fuzz_cases/sort_fuzz.rs | 8 ++++---- datafusion/execution/src/runtime_env.rs | 5 +++++ 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/sort_fuzz.rs b/datafusion/core/tests/fuzz_cases/sort_fuzz.rs index f4b4f16aa160..418504e9e24c 100644 --- a/datafusion/core/tests/fuzz_cases/sort_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/sort_fuzz.rs @@ -22,7 +22,7 @@ use arrow::{ compute::SortOptions, record_batch::RecordBatch, }; -use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv}; +use datafusion::execution::runtime_env::RuntimeConfig; use datafusion::physical_plan::expressions::PhysicalSortExpr; use datafusion::physical_plan::memory::MemoryExec; use datafusion::physical_plan::sorts::sort::SortExec; @@ -136,9 +136,9 @@ impl SortTest { .sort_spill_reservation_bytes, ); - let runtime_config = RuntimeConfig::new() - .with_memory_pool(Arc::new(GreedyMemoryPool::new(pool_size))); - let runtime = Arc::new(RuntimeEnv::new(runtime_config).unwrap()); + let runtime_env = RuntimeConfig::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(pool_size))).build(); + let runtime = Arc::new(runtime_env.unwrap()); SessionContext::new_with_config_rt(session_config, runtime) } else { SessionContext::new_with_config(session_config) diff --git a/datafusion/execution/src/runtime_env.rs b/datafusion/execution/src/runtime_env.rs index 420246595558..bddce81f537b 100644 --- a/datafusion/execution/src/runtime_env.rs +++ b/datafusion/execution/src/runtime_env.rs @@ -228,4 +228,9 @@ impl RuntimeConfig { pub fn with_temp_file_path(self, path: impl Into) -> Self { self.with_disk_manager(DiskManagerConfig::new_specified(vec![path.into()])) } + + /// Build a `RuntimeEnv` object from the configuration + pub fn build(self) -> Result { + RuntimeEnv::new(self) + } } From 0ae5fbe356a602a16a2f42b9468cfc5a515cd4c5 Mon Sep 17 00:00:00 2001 From: theirix Date: Sun, 25 Aug 2024 09:52:15 +0100 Subject: [PATCH 2/2] Fix formatting --- datafusion/core/tests/fuzz_cases/sort_fuzz.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/fuzz_cases/sort_fuzz.rs b/datafusion/core/tests/fuzz_cases/sort_fuzz.rs index 418504e9e24c..eda306dd3de5 100644 --- a/datafusion/core/tests/fuzz_cases/sort_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/sort_fuzz.rs @@ -137,7 +137,8 @@ impl SortTest { ); let runtime_env = RuntimeConfig::new() - .with_memory_pool(Arc::new(GreedyMemoryPool::new(pool_size))).build(); + .with_memory_pool(Arc::new(GreedyMemoryPool::new(pool_size))) + .build(); let runtime = Arc::new(runtime_env.unwrap()); SessionContext::new_with_config_rt(session_config, runtime) } else {