| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you under the Apache License, Version 2.0 (the |
| // "License"); you may not use this file except in compliance |
| // with the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, |
| // software distributed under the License is distributed on an |
| // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| // KIND, either express or implied. See the License for the |
| // specific language governing permissions and limitations |
| // under the License. |
| |
| //! UDF support |
| |
| use fmt::{Debug, Formatter}; |
| use std::fmt; |
| |
| use arrow::datatypes::Schema; |
| |
| use crate::error::Result; |
| use crate::{logical_plan::Expr, physical_plan::PhysicalExpr}; |
| |
| use super::{ |
| functions::{ |
| ReturnTypeFunction, ScalarFunctionExpr, ScalarFunctionImplementation, Signature, |
| }, |
| type_coercion::coerce, |
| }; |
| use std::sync::Arc; |
| |
| /// Logical representation of a UDF. |
| #[derive(Clone)] |
| pub struct ScalarUDF { |
| /// name |
| pub name: String, |
| /// signature |
| pub signature: Signature, |
| /// Return type |
| pub return_type: ReturnTypeFunction, |
| /// actual implementation |
| pub fun: ScalarFunctionImplementation, |
| } |
| |
| impl Debug for ScalarUDF { |
| fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { |
| f.debug_struct("ScalarUDF") |
| .field("name", &self.name) |
| .field("signature", &self.signature) |
| .field("fun", &"<FUNC>") |
| .finish() |
| } |
| } |
| |
| impl PartialEq for ScalarUDF { |
| fn eq(&self, other: &Self) -> bool { |
| self.name == other.name && self.signature == other.signature |
| } |
| } |
| |
| impl ScalarUDF { |
| /// Create a new ScalarUDF |
| pub fn new( |
| name: &str, |
| signature: &Signature, |
| return_type: &ReturnTypeFunction, |
| fun: &ScalarFunctionImplementation, |
| ) -> Self { |
| Self { |
| name: name.to_owned(), |
| signature: signature.clone(), |
| return_type: return_type.clone(), |
| fun: fun.clone(), |
| } |
| } |
| |
| /// creates a logical expression with a call of the UDF |
| /// This utility allows using the UDF without requiring access to the registry. |
| pub fn call(&self, args: Vec<Expr>) -> Expr { |
| Expr::ScalarUDF { |
| fun: Arc::new(self.clone()), |
| args, |
| } |
| } |
| } |
| |
| /// Create a physical expression of the UDF. |
| /// This function errors when `args`' can't be coerced to a valid argument type of the UDF. |
| pub fn create_physical_expr( |
| fun: &ScalarUDF, |
| args: &[Arc<dyn PhysicalExpr>], |
| input_schema: &Schema, |
| ) -> Result<Arc<dyn PhysicalExpr>> { |
| // coerce |
| let args = coerce(args, input_schema, &fun.signature)?; |
| |
| let arg_types = args |
| .iter() |
| .map(|e| e.data_type(input_schema)) |
| .collect::<Result<Vec<_>>>()?; |
| |
| Ok(Arc::new(ScalarFunctionExpr::new( |
| &fun.name, |
| fun.fun.clone(), |
| args, |
| (fun.return_type)(&arg_types)?.as_ref(), |
| ))) |
| } |