micromegas_analytics/lakehouse/
runtime.rs

1use super::scoped_memory_pool::ScopedMemoryPool;
2use anyhow::Result;
3use datafusion::execution::{
4    memory_pool::{GreedyMemoryPool, MemoryPool, TrackConsumersPool, UnboundedMemoryPool},
5    runtime_env::{RuntimeEnv, RuntimeEnvBuilder},
6};
7use std::{num::NonZeroUsize, sync::Arc};
8
9/// Applies `mb` (megabytes) as the disk manager's max temp directory size, in bytes.
10/// Factored out of `make_runtime_env` so it can be unit-tested without touching env vars.
11pub fn apply_max_temp_directory_mb(builder: RuntimeEnvBuilder, mb: u64) -> RuntimeEnvBuilder {
12    builder.with_max_temp_directory_size(mb * 1024 * 1024)
13}
14
15/// Parses the raw `MICROMEGAS_DATAFUSION_MAX_TEMP_DIRECTORY_MB` value (as returned by
16/// `std::env::var`) into an optional MB budget: `Ok(None)` when the variable is unset,
17/// `Err` when it is set to something that doesn't parse as `u64`. Factored out of
18/// `make_runtime_env` so the parse-failure path can be unit-tested without setting a real
19/// process env var, which would race with other tests running in parallel in the same
20/// process.
21pub fn parse_max_temp_directory_mb(var: Result<String, std::env::VarError>) -> Result<Option<u64>> {
22    match var {
23        Ok(mb_str) => Ok(Some(mb_str.parse::<u64>()?)),
24        Err(std::env::VarError::NotPresent) => Ok(None),
25        Err(err @ std::env::VarError::NotUnicode(_)) => Err(err.into()),
26    }
27}
28
29/// Creates a new DataFusion `RuntimeEnv` with a configurable memory pool.
30pub fn make_runtime_env() -> Result<RuntimeEnv> {
31    let nb_top_consumers = NonZeroUsize::new(5).unwrap();
32    let pool: Arc<dyn MemoryPool> = match std::env::var("MICROMEGAS_DATAFUSION_MEMORY_BUDGET_MB") {
33        Ok(mb_str) => {
34            let bytes = mb_str.parse::<usize>()? * 1024 * 1024;
35            Arc::new(TrackConsumersPool::new(
36                GreedyMemoryPool::new(bytes),
37                nb_top_consumers,
38            ))
39        }
40        Err(_) => Arc::new(TrackConsumersPool::new(
41            UnboundedMemoryPool::default(),
42            nb_top_consumers,
43        )),
44    };
45    let mut builder = RuntimeEnvBuilder::new().with_memory_pool(pool);
46    if let Some(mb) =
47        parse_max_temp_directory_mb(std::env::var("MICROMEGAS_DATAFUSION_MAX_TEMP_DIRECTORY_MB"))?
48    {
49        builder = apply_max_temp_directory_mb(builder, mb);
50    }
51    Ok(builder.build()?)
52}
53
54/// Builds a `RuntimeEnv` that reuses `shared`'s disk manager, caches and object-store
55/// registry but installs `scoped_pool` (already wrapping `shared`'s memory pool) as its
56/// memory pool. Takes the pool as a parameter, rather than constructing it internally, so
57/// callers can hand the (infallible) pool to `QueryAuditState` before this fallible step runs.
58pub fn scoped_runtime(
59    shared: &RuntimeEnv,
60    scoped_pool: Arc<ScopedMemoryPool>,
61) -> Result<Arc<RuntimeEnv>> {
62    Ok(Arc::new(
63        RuntimeEnvBuilder::from_runtime_env(shared)
64            .with_memory_pool(scoped_pool)
65            .build()?,
66    ))
67}
micromegas_analytics/lakehouse/runtime.rs

micromegas_analytics/lakehouse/
runtime.rs
