common_function/
function_factory.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_expr::ScalarUDF;
18
19use crate::function::{FunctionContext, FunctionRef};
20use crate::scalars::udf::create_udf;
21
22/// A factory for creating `ScalarUDF` that require a function context.
23#[derive(Clone)]
24pub struct ScalarFunctionFactory {
25    name: String,
26    factory: Arc<dyn Fn(FunctionContext) -> ScalarUDF + Send + Sync>,
27}
28
29impl ScalarFunctionFactory {
30    /// Returns the name of the function.
31    pub fn name(&self) -> &str {
32        &self.name
33    }
34
35    /// Returns a `ScalarUDF` when given a function context.
36    pub fn provide(&self, ctx: FunctionContext) -> ScalarUDF {
37        (self.factory)(ctx)
38    }
39}
40
41impl From<ScalarUDF> for ScalarFunctionFactory {
42    fn from(df_udf: ScalarUDF) -> Self {
43        let name = df_udf.name().to_string();
44        let func = Arc::new(move |_ctx| df_udf.clone());
45        Self {
46            name,
47            factory: func,
48        }
49    }
50}
51
52impl From<FunctionRef> for ScalarFunctionFactory {
53    fn from(func: FunctionRef) -> Self {
54        let name = func.name().to_string();
55        let func = Arc::new(move |ctx: FunctionContext| {
56            create_udf(func.clone(), ctx.query_ctx, ctx.state)
57        });
58        Self {
59            name,
60            factory: func,
61        }
62    }
63}