1 // Licensed to the Apache Software Foundation (ASF) under one
2 // or more contributor license agreements. See the NOTICE file
3 // distributed with this work for additional information
4 // regarding copyright ownership. The ASF licenses this file
5 // to you under the Apache License, Version 2.0 (the
6 // "License"); you may not use this file except in compliance
7 // with the License. You may obtain a copy of the License at
8 //
9 // http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing,
12 // software distributed under the License is distributed on an
13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 // KIND, either express or implied. See the License for the
15 // specific language governing permissions and limitations
16 // under the License.
17
18 //! UDF support
19
20 use fmt::{Debug, Formatter};
21 use std::fmt;
22
23 use arrow::datatypes::Schema;
24
25 use crate::error::Result;
26 use crate::{logical_plan::Expr, physical_plan::PhysicalExpr};
27
28 use super::{
29 functions::{
30 ReturnTypeFunction, ScalarFunctionExpr, ScalarFunctionImplementation, Signature,
31 },
32 type_coercion::coerce,
33 };
34 use std::sync::Arc;
35
36 /// Logical representation of a UDF.
37 #[derive(Clone)]
38 pub struct ScalarUDF {
39 /// name
40 pub name: String,
41 /// signature
42 pub signature: Signature,
43 /// Return type
44 pub return_type: ReturnTypeFunction,
45 /// actual implementation
46 pub fun: ScalarFunctionImplementation,
47 }
48
49 impl Debug for ScalarUDF {
fmt(&self, f: &mut Formatter<'_>) -> fmt::Result50 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
51 f.debug_struct("ScalarUDF")
52 .field("name", &self.name)
53 .field("signature", &self.signature)
54 .field("fun", &"<FUNC>")
55 .finish()
56 }
57 }
58
59 impl PartialEq for ScalarUDF {
eq(&self, other: &Self) -> bool60 fn eq(&self, other: &Self) -> bool {
61 self.name == other.name && self.signature == other.signature
62 }
63 }
64
65 impl ScalarUDF {
66 /// Create a new ScalarUDF
new( name: &str, signature: &Signature, return_type: &ReturnTypeFunction, fun: &ScalarFunctionImplementation, ) -> Self67 pub fn new(
68 name: &str,
69 signature: &Signature,
70 return_type: &ReturnTypeFunction,
71 fun: &ScalarFunctionImplementation,
72 ) -> Self {
73 Self {
74 name: name.to_owned(),
75 signature: signature.clone(),
76 return_type: return_type.clone(),
77 fun: fun.clone(),
78 }
79 }
80
81 /// creates a logical expression with a call of the UDF
82 /// This utility allows using the UDF without requiring access to the registry.
call(&self, args: Vec<Expr>) -> Expr83 pub fn call(&self, args: Vec<Expr>) -> Expr {
84 Expr::ScalarUDF {
85 fun: Arc::new(self.clone()),
86 args,
87 }
88 }
89 }
90
91 /// Create a physical expression of the UDF.
92 /// This function errors when `args`' can't be coerced to a valid argument type of the UDF.
create_physical_expr( fun: &ScalarUDF, args: &[Arc<dyn PhysicalExpr>], input_schema: &Schema, ) -> Result<Arc<dyn PhysicalExpr>>93 pub fn create_physical_expr(
94 fun: &ScalarUDF,
95 args: &[Arc<dyn PhysicalExpr>],
96 input_schema: &Schema,
97 ) -> Result<Arc<dyn PhysicalExpr>> {
98 // coerce
99 let args = coerce(args, input_schema, &fun.signature)?;
100
101 let arg_types = args
102 .iter()
103 .map(|e| e.data_type(input_schema))
104 .collect::<Result<Vec<_>>>()?;
105
106 Ok(Arc::new(ScalarFunctionExpr::new(
107 &fun.name,
108 fun.fun.clone(),
109 args,
110 (fun.return_type)(&arg_types)?.as_ref(),
111 )))
112 }
113