query/
query_engine.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
15mod context;
16mod default_serializer;
17pub mod options;
18pub mod runtime;
19mod state;
20use std::any::Any;
21use std::sync::Arc;
22
23use async_trait::async_trait;
24use catalog::CatalogManagerRef;
25use common_base::Plugins;
26use common_function::function_factory::ScalarFunctionFactory;
27use common_function::function_registry::FUNCTION_REGISTRY;
28use common_function::handlers::{
29    FlowServiceHandlerRef, ProcedureServiceHandlerRef, TableMutationHandlerRef,
30};
31use common_query::Output;
32use datafusion::catalog::TableFunction;
33use datafusion::dataframe::DataFrame;
34use datafusion_expr::{AggregateUDF, LogicalPlan, WindowUDF};
35pub use default_serializer::{DefaultPlanDecoder, DefaultSerializer};
36use partition::manager::PartitionRuleManagerRef;
37use session::context::QueryContextRef;
38use table::TableRef;
39
40use crate::datafusion::DatafusionQueryEngine;
41use crate::error::Result;
42use crate::options::QueryOptions;
43use crate::planner::LogicalPlanner;
44pub use crate::query_engine::context::QueryEngineContext;
45pub use crate::query_engine::runtime::{
46    DefaultQueryRuntimeProvider, QueryRuntimeContext, QueryRuntimeProvider, QueryRuntimeProviderRef,
47};
48pub use crate::query_engine::state::QueryEngineState;
49use crate::region_query::RegionQueryHandlerRef;
50
51/// Describe statement result
52#[derive(Debug)]
53pub struct DescribeResult {
54    /// The logical plan for statement
55    pub logical_plan: LogicalPlan,
56}
57
58#[async_trait]
59pub trait QueryEngine: Send + Sync {
60    /// Returns the query engine as Any
61    /// so that it can be downcast to a specific implementation.
62    fn as_any(&self) -> &dyn Any;
63
64    /// Returns the logical planner
65    fn planner(&self) -> Arc<dyn LogicalPlanner>;
66
67    /// Returns the query engine name.
68    fn name(&self) -> &str;
69
70    /// Describe the given [`LogicalPlan`].
71    async fn describe(
72        &self,
73        plan: LogicalPlan,
74        query_ctx: QueryContextRef,
75    ) -> Result<DescribeResult>;
76
77    /// Execute the given [`LogicalPlan`].
78    async fn execute(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<Output>;
79
80    /// Register an aggregate function.
81    ///
82    /// # Panics
83    /// Will panic if the function with same name is already registered.
84    fn register_aggregate_function(&self, func: AggregateUDF);
85
86    /// Register a scalar function.
87    /// Will override if the function with same name is already registered.
88    fn register_scalar_function(&self, func: ScalarFunctionFactory);
89
90    /// Register table function
91    fn register_table_function(&self, func: Arc<TableFunction>);
92
93    /// Register a window function (UDWF).
94    fn register_window_function(&self, func: WindowUDF);
95
96    /// Create a DataFrame from a table.
97    fn read_table(&self, table: TableRef) -> Result<DataFrame>;
98
99    /// Create a [`QueryEngineContext`].
100    fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext;
101
102    /// Retrieve the query engine state [`QueryEngineState`]
103    fn engine_state(&self) -> &QueryEngineState;
104}
105
106pub struct QueryEngineFactory {
107    query_engine: Arc<dyn QueryEngine>,
108}
109
110impl QueryEngineFactory {
111    pub fn new(
112        catalog_manager: CatalogManagerRef,
113        region_query_handler: Option<RegionQueryHandlerRef>,
114        table_mutation_handler: Option<TableMutationHandlerRef>,
115        procedure_service_handler: Option<ProcedureServiceHandlerRef>,
116        flow_service_handler: Option<FlowServiceHandlerRef>,
117        with_dist_planner: bool,
118        options: QueryOptions,
119    ) -> Self {
120        Self::try_new(
121            catalog_manager,
122            region_query_handler,
123            table_mutation_handler,
124            procedure_service_handler,
125            flow_service_handler,
126            with_dist_planner,
127            options,
128        )
129        .expect("Failed to build query engine factory")
130    }
131
132    pub fn try_new(
133        catalog_manager: CatalogManagerRef,
134        region_query_handler: Option<RegionQueryHandlerRef>,
135        table_mutation_handler: Option<TableMutationHandlerRef>,
136        procedure_service_handler: Option<ProcedureServiceHandlerRef>,
137        flow_service_handler: Option<FlowServiceHandlerRef>,
138        with_dist_planner: bool,
139        options: QueryOptions,
140    ) -> datafusion::error::Result<Self> {
141        Self::try_new_with_plugins(
142            catalog_manager,
143            None,
144            region_query_handler,
145            table_mutation_handler,
146            procedure_service_handler,
147            flow_service_handler,
148            with_dist_planner,
149            Default::default(),
150            options,
151        )
152    }
153
154    #[allow(clippy::too_many_arguments)]
155    pub fn new_with_plugins(
156        catalog_manager: CatalogManagerRef,
157        partition_rule_manager: Option<PartitionRuleManagerRef>,
158        region_query_handler: Option<RegionQueryHandlerRef>,
159        table_mutation_handler: Option<TableMutationHandlerRef>,
160        procedure_service_handler: Option<ProcedureServiceHandlerRef>,
161        flow_service_handler: Option<FlowServiceHandlerRef>,
162        with_dist_planner: bool,
163        plugins: Plugins,
164        options: QueryOptions,
165    ) -> Self {
166        Self::try_new_with_plugins(
167            catalog_manager,
168            partition_rule_manager,
169            region_query_handler,
170            table_mutation_handler,
171            procedure_service_handler,
172            flow_service_handler,
173            with_dist_planner,
174            plugins,
175            options,
176        )
177        .expect("Failed to build query engine factory")
178    }
179
180    #[allow(clippy::too_many_arguments)]
181    pub fn try_new_with_plugins(
182        catalog_manager: CatalogManagerRef,
183        partition_rule_manager: Option<PartitionRuleManagerRef>,
184        region_query_handler: Option<RegionQueryHandlerRef>,
185        table_mutation_handler: Option<TableMutationHandlerRef>,
186        procedure_service_handler: Option<ProcedureServiceHandlerRef>,
187        flow_service_handler: Option<FlowServiceHandlerRef>,
188        with_dist_planner: bool,
189        plugins: Plugins,
190        options: QueryOptions,
191    ) -> datafusion::error::Result<Self> {
192        let state = Arc::new(QueryEngineState::try_new(
193            catalog_manager,
194            partition_rule_manager,
195            region_query_handler,
196            table_mutation_handler,
197            procedure_service_handler,
198            flow_service_handler,
199            with_dist_planner,
200            plugins.clone(),
201            options,
202        )?);
203        let query_engine = Arc::new(DatafusionQueryEngine::new(state, plugins));
204        register_functions(&query_engine);
205        Ok(Self { query_engine })
206    }
207
208    pub fn query_engine(&self) -> QueryEngineRef {
209        self.query_engine.clone()
210    }
211}
212
213/// Register all functions implemented by GreptimeDB
214fn register_functions(query_engine: &Arc<DatafusionQueryEngine>) {
215    for func in FUNCTION_REGISTRY.scalar_functions() {
216        query_engine.register_scalar_function(func);
217    }
218
219    for accumulator in FUNCTION_REGISTRY.aggregate_functions() {
220        query_engine.register_aggregate_function(accumulator);
221    }
222
223    for table_function in FUNCTION_REGISTRY.table_functions() {
224        query_engine.register_table_function(table_function);
225    }
226
227    for window_function in FUNCTION_REGISTRY.window_functions() {
228        query_engine.register_window_function(window_function);
229    }
230}
231
232pub type QueryEngineRef = Arc<dyn QueryEngine>;
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237
238    #[test]
239    fn test_query_engine_factory() {
240        let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap();
241        let factory = QueryEngineFactory::new(
242            catalog_list,
243            None,
244            None,
245            None,
246            None,
247            false,
248            QueryOptions::default(),
249        );
250
251        let engine = factory.query_engine();
252
253        assert_eq!("datafusion", engine.name());
254    }
255}
query/query_engine.rs

query/
query_engine.rs
