From c0ef5330c85606fcb4fe8ee691e062df775b7354 Mon Sep 17 00:00:00 2001 From: comphead Date: Thu, 6 Jun 2024 15:13:30 -0700 Subject: [PATCH] Minor: Remove code duplication in `memory_limit` derivation for datafusion-cli --- datafusion-cli/src/main.rs | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs index d81a73df20a4..187f856894b2 100644 --- a/datafusion-cli/src/main.rs +++ b/datafusion-cli/src/main.rs @@ -177,18 +177,14 @@ async fn main_inner() -> Result<()> { let rt_config = // set memory pool size if let Some(memory_limit) = args.memory_limit { + // unwrap is safe here because is_valid_memory_pool_size already checked the value let memory_limit = extract_memory_pool_size(&memory_limit).unwrap(); // set memory pool type - if let Some(mem_pool_type) = args.mem_pool_type { - match mem_pool_type { - PoolType::Greedy => rt_config - .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit))), - PoolType::Fair => rt_config - .with_memory_pool(Arc::new(FairSpillPool::new(memory_limit))), - } - } else { - rt_config - .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit))) + match args.mem_pool_type { + Some(PoolType::Fair) => rt_config + .with_memory_pool(Arc::new(FairSpillPool::new(memory_limit))), + _ => rt_config + .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit))) } } else { rt_config