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

query/
query_engine.rs
