query/query_engine/
runtime.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use datafusion::error::Result as DfResult;
18use datafusion::execution::context::SessionConfig;
19use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
20
21use crate::options::QueryOptions;
22use crate::query_engine::state::MetricsMemoryPool;
23
24/// Reference-counted query runtime provider.
25pub type QueryRuntimeProviderRef = Arc<dyn QueryRuntimeProvider>;
26
27/// Context for building query runtime components.
28#[derive(Clone, Copy)]
29#[non_exhaustive]
30pub struct QueryRuntimeContext<'a> {
31    /// Query options used by the query engine.
32    pub query_options: &'a QueryOptions,
33    /// Resolved memory pool size in bytes.
34    pub resolved_memory_pool_size: usize,
35}
36
37impl<'a> QueryRuntimeContext<'a> {
38    /// Creates a new query runtime context.
39    pub fn new(query_options: &'a QueryOptions, resolved_memory_pool_size: usize) -> Self {
40        Self {
41            query_options,
42            resolved_memory_pool_size,
43        }
44    }
45}
46
47/// Provides DataFusion session and runtime setup for the query engine.
48pub trait QueryRuntimeProvider: Send + Sync + 'static {
49    /// Configures the DataFusion session config before building the session state.
50    fn configure_session_config(&self, _ctx: QueryRuntimeContext<'_>, _config: &mut SessionConfig) {
51    }
52
53    /// Builds the DataFusion runtime environment.
54    fn build_runtime_env(
55        &self,
56        _ctx: QueryRuntimeContext<'_>,
57        builder: RuntimeEnvBuilder,
58    ) -> DfResult<Arc<RuntimeEnv>> {
59        builder.build().map(Arc::new)
60    }
61}
62
63/// Default query runtime provider.
64#[derive(Debug, Default)]
65pub struct DefaultQueryRuntimeProvider;
66
67impl DefaultQueryRuntimeProvider {
68    /// Creates a default DataFusion runtime environment builder.
69    pub fn runtime_env_builder(ctx: QueryRuntimeContext<'_>) -> RuntimeEnvBuilder {
70        if ctx.resolved_memory_pool_size > 0 {
71            RuntimeEnvBuilder::new().with_memory_pool(Arc::new(MetricsMemoryPool::new(
72                ctx.resolved_memory_pool_size,
73            )))
74        } else {
75            RuntimeEnvBuilder::new()
76        }
77    }
78}
79
80impl QueryRuntimeProvider for DefaultQueryRuntimeProvider {}
query/query_engine/runtime.rs

query/query_engine/
runtime.rs
