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_error::ext::{BoxedError, PlainError};
19use common_error::status_code::StatusCode;
20use common_query::error::{ExecuteSnafu, Result};
21use datafusion::arrow::datatypes::DataType;
22use datafusion::logical_expr::ColumnarValue;
23use datafusion_expr::{ScalarFunctionArgs, Signature};
24use datatypes::vectors::VectorRef;
25use session::context::{QueryContextBuilder, QueryContextRef};
26use snafu::ResultExt;
27
28use crate::state::FunctionState;
29
30/// The function execution context
31#[derive(Clone)]
32pub struct FunctionContext {
33    pub query_ctx: QueryContextRef,
34    pub state: Arc<FunctionState>,
35}
36
37impl FunctionContext {
38    /// Create a mock [`FunctionContext`] for test.
39    #[cfg(any(test, feature = "testing"))]
40    pub fn mock() -> Self {
41        Self {
42            query_ctx: QueryContextBuilder::default().build().into(),
43            state: Arc::new(FunctionState::mock()),
44        }
45    }
46}
47
48impl std::fmt::Display for FunctionContext {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        write!(f, "FunctionContext {{ query_ctx: {} }}", self.query_ctx)
51    }
52}
53
54impl Default for FunctionContext {
55    fn default() -> Self {
56        Self {
57            query_ctx: QueryContextBuilder::default().build().into(),
58            state: Arc::new(FunctionState::default()),
59        }
60    }
61}
62
63/// Scalar function trait, modified from databend to adapt datafusion
64/// TODO(dennis): optimize function by it's features such as monotonicity etc.
65pub trait Function: fmt::Display + Sync + Send {
66    /// Returns the name of the function, should be unique.
67    fn name(&self) -> &str;
68
69    /// The returned data type of function execution.
70    fn return_type(&self, input_types: &[DataType]) -> Result<DataType>;
71
72    /// The signature of function.
73    fn signature(&self) -> Signature;
74
75    fn invoke_with_args(
76        &self,
77        args: ScalarFunctionArgs,
78    ) -> datafusion_common::Result<ColumnarValue> {
79        // TODO(LFC): Remove default implementation once all UDFs have implemented this function.
80        let _ = args;
81        Err(datafusion_common::DataFusionError::NotImplemented(
82            "invoke_with_args".to_string(),
83        ))
84    }
85
86    /// Evaluate the function, e.g. run/execute the function.
87    /// TODO(LFC): Remove `eval` when all UDFs are rewritten to `invoke_with_args`
88    fn eval(&self, _: &FunctionContext, _: &[VectorRef]) -> Result<VectorRef> {
89        Err(BoxedError::new(PlainError::new(
90            "unsupported".to_string(),
91            StatusCode::Unsupported,
92        )))
93        .context(ExecuteSnafu)
94    }
95
96    fn aliases(&self) -> &[String] {
97        &[]
98    }
99}
100
101pub type FunctionRef = Arc<dyn Function>;
common_function/function.rs

common_function/
function.rs
