Return an appropriately configured `RuntimeEnvBuilder`
(&self)
| 98 | |
| 99 | /// Return an appropriately configured `RuntimeEnvBuilder` |
| 100 | pub fn runtime_env_builder(&self) -> Result<RuntimeEnvBuilder> { |
| 101 | let mut rt_builder = RuntimeEnvBuilder::new(); |
| 102 | const NUM_TRACKED_CONSUMERS: usize = 5; |
| 103 | // Use CLI --memory-limit if provided, otherwise fall back to |
| 104 | // DATAFUSION_RUNTIME_MEMORY_LIMIT env var |
| 105 | let memory_limit = self.memory_limit.or_else(|| { |
| 106 | std::env::var("DATAFUSION_RUNTIME_MEMORY_LIMIT") |
| 107 | .ok() |
| 108 | .and_then(|val| parse_capacity_limit(&val).ok()) |
| 109 | }); |
| 110 | |
| 111 | if let Some(memory_limit) = memory_limit { |
| 112 | let pool: Arc<dyn MemoryPool> = match self.mem_pool_type.as_str() { |
| 113 | "fair" => Arc::new(TrackConsumersPool::new( |
| 114 | FairSpillPool::new(memory_limit), |
| 115 | NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(), |
| 116 | )), |
| 117 | "greedy" => Arc::new(TrackConsumersPool::new( |
| 118 | GreedyMemoryPool::new(memory_limit), |
| 119 | NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(), |
| 120 | )), |
| 121 | _ => { |
| 122 | return Err(DataFusionError::Configuration(format!( |
| 123 | "Invalid memory pool type: {}", |
| 124 | self.mem_pool_type |
| 125 | ))); |
| 126 | } |
| 127 | }; |
| 128 | rt_builder = rt_builder |
| 129 | .with_memory_pool(pool) |
| 130 | .with_disk_manager_builder(DiskManagerBuilder::default()); |
| 131 | } |
| 132 | Ok(rt_builder) |
| 133 | } |
| 134 | |
| 135 | /// Build the runtime environment, optionally wrapping the local filesystem |
| 136 | /// with a throttled object store to simulate remote storage latency. |