micromegas_analytics/lakehouse/
runtime.rs

1use anyhow::Result;
2use datafusion::execution::{
3    memory_pool::{GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool},
4    runtime_env::{RuntimeEnv, RuntimeEnvBuilder},
5};
6use std::{num::NonZeroUsize, sync::Arc};
7
8/// Creates a new DataFusion `RuntimeEnv` with a configurable memory pool.
9pub fn make_runtime_env() -> Result<RuntimeEnv> {
10    let nb_top_consumers = NonZeroUsize::new(5).unwrap();
11    let pool: Arc<dyn MemoryPool> = match std::env::var("MICROMEGAS_DATAFUSION_MEMORY_BUDGET_MB") {
12        Ok(mb_str) => {
13            let bytes = mb_str.parse::<usize>()? * 1024 * 1024;
14            Arc::new(TrackConsumersPool::new(
15                GreedyMemoryPool::new(bytes),
16                nb_top_consumers,
17            ))
18        }
19        Err(_) => Arc::new(TrackConsumersPool::new(
20            UnboundedMemoryPool::default(),
21            nb_top_consumers,
22        )),
23    };
24    Ok(RuntimeEnvBuilder::new().with_memory_pool(pool).build()?)
25}

micromegas_analytics/lakehouse/runtime.rs

micromegas_analytics/lakehouse/
runtime.rs