Minor: Remove code duplication in memory_limit derivation for dataf…

…usion-cli (#10814) * Fix: Sort Merge Join crashes on TPCH Q21 * Fix LeftAnti SMJ join when the join filter is set * rm dbg * Minor: Remove code duplication in `memory_limit` derivation for datafusion-cli
apache · Jun 7, 2024 · 8fcb3e4 · 8fcb3e4
1 parent 586241f
commit 8fcb3e4
Showing 1 changed file with 6 additions and 10 deletions.
diff --git 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