common_function/
function.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::fmt;
16use std::sync::Arc;
17
18use common_query::error::Result;
19use common_query::prelude::Signature;
20use datatypes::data_type::ConcreteDataType;
21use datatypes::vectors::VectorRef;
22use session::context::{QueryContextBuilder, QueryContextRef};
23
24use crate::state::FunctionState;
25
26/// The function execution context
27#[derive(Clone)]
28pub struct FunctionContext {
29    pub query_ctx: QueryContextRef,
30    pub state: Arc<FunctionState>,
31}
32
33impl FunctionContext {
34    /// Create a mock [`FunctionContext`] for test.
35    #[cfg(any(test, feature = "testing"))]
36    pub fn mock() -> Self {
37        Self {
38            query_ctx: QueryContextBuilder::default().build().into(),
39            state: Arc::new(FunctionState::mock()),
40        }
41    }
42}
43
44impl Default for FunctionContext {
45    fn default() -> Self {
46        Self {
47            query_ctx: QueryContextBuilder::default().build().into(),
48            state: Arc::new(FunctionState::default()),
49        }
50    }
51}
52
53/// Scalar function trait, modified from databend to adapt datafusion
54/// TODO(dennis): optimize function by it's features such as monotonicity etc.
55pub trait Function: fmt::Display + Sync + Send {
56    /// Returns the name of the function, should be unique.
57    fn name(&self) -> &str;
58
59    /// The returned data type of function execution.
60    fn return_type(&self, input_types: &[ConcreteDataType]) -> Result<ConcreteDataType>;
61
62    /// The signature of function.
63    fn signature(&self) -> Signature;
64
65    /// Evaluate the function, e.g. run/execute the function.
66    fn eval(&self, ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef>;
67}
68
69pub type FunctionRef = Arc<dyn Function>;
70
71/// Async Scalar function trait
72#[async_trait::async_trait]
73pub trait AsyncFunction: fmt::Display + Sync + Send {
74    /// Returns the name of the function, should be unique.
75    fn name(&self) -> &str;
76
77    /// The returned data type of function execution.
78    fn return_type(&self, input_types: &[ConcreteDataType]) -> Result<ConcreteDataType>;
79
80    /// The signature of function.
81    fn signature(&self) -> Signature;
82
83    /// Evaluate the function, e.g. run/execute the function.
84    /// TODO(dennis): simplify the signature and refactor all the admin functions.
85    async fn eval(&self, _func_ctx: FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef>;
86}
87
88pub type AsyncFunctionRef = Arc<dyn AsyncFunction>;