| // 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. |
| |
| //! [`ScalarUDF`]: Scalar User Defined Functions |
| |
| use crate::async_udf::AsyncScalarUDF; |
| use crate::expr::schema_name_from_exprs_comma_separated_without_space; |
| use crate::simplify::{ExprSimplifyResult, SimplifyInfo}; |
| use crate::sort_properties::{ExprProperties, SortProperties}; |
| use crate::udf_eq::UdfEq; |
| use crate::{ColumnarValue, Documentation, Expr, Signature}; |
| use arrow::datatypes::{DataType, Field, FieldRef}; |
| use datafusion_common::config::ConfigOptions; |
| use datafusion_common::{not_impl_err, ExprSchema, Result, ScalarValue}; |
| use datafusion_expr_common::dyn_eq::{DynEq, DynHash}; |
| use datafusion_expr_common::interval_arithmetic::Interval; |
| use std::any::Any; |
| use std::cmp::Ordering; |
| use std::fmt::Debug; |
| use std::hash::{Hash, Hasher}; |
| use std::sync::Arc; |
| |
| /// Logical representation of a Scalar User Defined Function. |
| /// |
| /// A scalar function produces a single row output for each row of input. This |
| /// struct contains the information DataFusion needs to plan and invoke |
| /// functions you supply such as name, type signature, return type, and actual |
| /// implementation. |
| /// |
| /// 1. For simple use cases, use [`create_udf`] (examples in [`simple_udf.rs`]). |
| /// |
| /// 2. For advanced use cases, use [`ScalarUDFImpl`] which provides full API |
| /// access (examples in [`advanced_udf.rs`]). |
| /// |
| /// See [`Self::call`] to create an `Expr` which invokes a `ScalarUDF` with arguments. |
| /// |
| /// # API Note |
| /// |
| /// This is a separate struct from [`ScalarUDFImpl`] to maintain backwards |
| /// compatibility with the older API. |
| /// |
| /// [`create_udf`]: crate::expr_fn::create_udf |
| /// [`simple_udf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/simple_udf.rs |
| /// [`advanced_udf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/advanced_udf.rs |
| #[derive(Debug, Clone)] |
| pub struct ScalarUDF { |
| inner: Arc<dyn ScalarUDFImpl>, |
| } |
| |
| impl PartialEq for ScalarUDF { |
| fn eq(&self, other: &Self) -> bool { |
| self.inner.dyn_eq(other.inner.as_any()) |
| } |
| } |
| |
| impl PartialOrd for ScalarUDF { |
| fn partial_cmp(&self, other: &Self) -> Option<Ordering> { |
| let mut cmp = self.name().cmp(other.name()); |
| if cmp == Ordering::Equal { |
| cmp = self.signature().partial_cmp(other.signature())?; |
| } |
| if cmp == Ordering::Equal { |
| cmp = self.aliases().partial_cmp(other.aliases())?; |
| } |
| // Contract for PartialOrd and PartialEq consistency requires that |
| // a == b if and only if partial_cmp(a, b) == Some(Equal). |
| if cmp == Ordering::Equal && self != other { |
| // Functions may have other properties besides name and signature |
| // that differentiate two instances (e.g. type, or arbitrary parameters). |
| // We cannot return Some(Equal) in such case. |
| return None; |
| } |
| debug_assert!( |
| cmp == Ordering::Equal || self != other, |
| "Detected incorrect implementation of PartialEq when comparing functions: '{}' and '{}'. \ |
| The functions compare as equal, but they are not equal based on general properties that \ |
| the PartialOrd implementation observes,", |
| self.name(), other.name() |
| ); |
| Some(cmp) |
| } |
| } |
| |
| impl Eq for ScalarUDF {} |
| |
| impl Hash for ScalarUDF { |
| fn hash<H: Hasher>(&self, state: &mut H) { |
| self.inner.dyn_hash(state) |
| } |
| } |
| |
| impl ScalarUDF { |
| /// Create a new `ScalarUDF` from a `[ScalarUDFImpl]` trait object |
| /// |
| /// Note this is the same as using the `From` impl (`ScalarUDF::from`) |
| pub fn new_from_impl<F>(fun: F) -> ScalarUDF |
| where |
| F: ScalarUDFImpl + 'static, |
| { |
| Self::new_from_shared_impl(Arc::new(fun)) |
| } |
| |
| /// Create a new `ScalarUDF` from a `[ScalarUDFImpl]` trait object |
| pub fn new_from_shared_impl(fun: Arc<dyn ScalarUDFImpl>) -> ScalarUDF { |
| Self { inner: fun } |
| } |
| |
| /// Return the underlying [`ScalarUDFImpl`] trait object for this function |
| pub fn inner(&self) -> &Arc<dyn ScalarUDFImpl> { |
| &self.inner |
| } |
| |
| /// Adds additional names that can be used to invoke this function, in |
| /// addition to `name` |
| /// |
| /// If you implement [`ScalarUDFImpl`] directly you should return aliases directly. |
| pub fn with_aliases(self, aliases: impl IntoIterator<Item = &'static str>) -> Self { |
| Self::new_from_impl(AliasedScalarUDFImpl::new(Arc::clone(&self.inner), aliases)) |
| } |
| |
| /// Returns a [`Expr`] logical expression to call this UDF with specified |
| /// arguments. |
| /// |
| /// This utility allows easily calling UDFs |
| /// |
| /// # Example |
| /// ```no_run |
| /// use datafusion_expr::{col, lit, ScalarUDF}; |
| /// # fn my_udf() -> ScalarUDF { unimplemented!() } |
| /// let my_func: ScalarUDF = my_udf(); |
| /// // Create an expr for `my_func(a, 12.3)` |
| /// let expr = my_func.call(vec![col("a"), lit(12.3)]); |
| /// ``` |
| pub fn call(&self, args: Vec<Expr>) -> Expr { |
| Expr::ScalarFunction(crate::expr::ScalarFunction::new_udf( |
| Arc::new(self.clone()), |
| args, |
| )) |
| } |
| |
| /// Returns this function's name. |
| /// |
| /// See [`ScalarUDFImpl::name`] for more details. |
| pub fn name(&self) -> &str { |
| self.inner.name() |
| } |
| |
| /// Returns this function's display_name. |
| /// |
| /// See [`ScalarUDFImpl::display_name`] for more details |
| #[deprecated( |
| since = "50.0.0", |
| note = "This method is unused and will be removed in a future release" |
| )] |
| pub fn display_name(&self, args: &[Expr]) -> Result<String> { |
| #[expect(deprecated)] |
| self.inner.display_name(args) |
| } |
| |
| /// Returns this function's schema_name. |
| /// |
| /// See [`ScalarUDFImpl::schema_name`] for more details |
| pub fn schema_name(&self, args: &[Expr]) -> Result<String> { |
| self.inner.schema_name(args) |
| } |
| |
| /// Returns the aliases for this function. |
| /// |
| /// See [`ScalarUDF::with_aliases`] for more details |
| pub fn aliases(&self) -> &[String] { |
| self.inner.aliases() |
| } |
| |
| /// Returns this function's [`Signature`] (what input types are accepted). |
| /// |
| /// See [`ScalarUDFImpl::signature`] for more details. |
| pub fn signature(&self) -> &Signature { |
| self.inner.signature() |
| } |
| |
| /// The datatype this function returns given the input argument types. |
| /// This function is used when the input arguments are [`DataType`]s. |
| /// |
| /// # Notes |
| /// |
| /// If a function implement [`ScalarUDFImpl::return_field_from_args`], |
| /// its [`ScalarUDFImpl::return_type`] should raise an error. |
| /// |
| /// See [`ScalarUDFImpl::return_type`] for more details. |
| pub fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { |
| self.inner.return_type(arg_types) |
| } |
| |
| /// Return the datatype this function returns given the input argument types. |
| /// |
| /// See [`ScalarUDFImpl::return_field_from_args`] for more details. |
| pub fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { |
| self.inner.return_field_from_args(args) |
| } |
| |
| /// Do the function rewrite |
| /// |
| /// See [`ScalarUDFImpl::simplify`] for more details. |
| pub fn simplify( |
| &self, |
| args: Vec<Expr>, |
| info: &dyn SimplifyInfo, |
| ) -> Result<ExprSimplifyResult> { |
| self.inner.simplify(args, info) |
| } |
| |
| #[deprecated(since = "50.0.0", note = "Use `return_field_from_args` instead.")] |
| pub fn is_nullable(&self, args: &[Expr], schema: &dyn ExprSchema) -> bool { |
| #[allow(deprecated)] |
| self.inner.is_nullable(args, schema) |
| } |
| |
| /// Invoke the function on `args`, returning the appropriate result. |
| /// |
| /// See [`ScalarUDFImpl::invoke_with_args`] for details. |
| pub fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { |
| #[cfg(debug_assertions)] |
| let return_field = Arc::clone(&args.return_field); |
| let result = self.inner.invoke_with_args(args)?; |
| // Maybe this could be enabled always? |
| // This doesn't use debug_assert!, but it's meant to run anywhere except on production. It's same in spirit, thus conditioning on debug_assertions. |
| #[cfg(debug_assertions)] |
| { |
| if &result.data_type() != return_field.data_type() { |
| return datafusion_common::internal_err!("Function '{}' returned value of type '{:?}' while the following type was promised at planning time and expected: '{:?}'", |
| self.name(), |
| result.data_type(), |
| return_field.data_type() |
| ); |
| } |
| // TODO verify return data is non-null when it was promised to be? |
| } |
| Ok(result) |
| } |
| |
| /// Determines which of the arguments passed to this function are evaluated eagerly |
| /// and which may be evaluated lazily. |
| /// |
| /// See [ScalarUDFImpl::conditional_arguments] for more information. |
| pub fn conditional_arguments<'a>( |
| &self, |
| args: &'a [Expr], |
| ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> { |
| self.inner.conditional_arguments(args) |
| } |
| |
| /// Returns true if some of this `exprs` subexpressions may not be evaluated |
| /// and thus any side effects (like divide by zero) may not be encountered. |
| /// |
| /// See [ScalarUDFImpl::short_circuits] for more information. |
| pub fn short_circuits(&self) -> bool { |
| self.inner.short_circuits() |
| } |
| |
| /// Computes the output interval for a [`ScalarUDF`], given the input |
| /// intervals. |
| /// |
| /// # Parameters |
| /// |
| /// * `inputs` are the intervals for the inputs (children) of this function. |
| /// |
| /// # Example |
| /// |
| /// If the function is `ABS(a)`, and the input interval is `a: [-3, 2]`, |
| /// then the output interval would be `[0, 3]`. |
| pub fn evaluate_bounds(&self, inputs: &[&Interval]) -> Result<Interval> { |
| self.inner.evaluate_bounds(inputs) |
| } |
| |
| /// Updates bounds for child expressions, given a known interval for this |
| /// function. This is used to propagate constraints down through an expression |
| /// tree. |
| /// |
| /// # Parameters |
| /// |
| /// * `interval` is the currently known interval for this function. |
| /// * `inputs` are the current intervals for the inputs (children) of this function. |
| /// |
| /// # Returns |
| /// |
| /// A `Vec` of new intervals for the children, in order. |
| /// |
| /// If constraint propagation reveals an infeasibility for any child, returns |
| /// [`None`]. If none of the children intervals change as a result of |
| /// propagation, may return an empty vector instead of cloning `children`. |
| /// This is the default (and conservative) return value. |
| /// |
| /// # Example |
| /// |
| /// If the function is `ABS(a)`, the current `interval` is `[4, 5]` and the |
| /// input `a` is given as `[-7, 3]`, then propagation would return `[-5, 3]`. |
| pub fn propagate_constraints( |
| &self, |
| interval: &Interval, |
| inputs: &[&Interval], |
| ) -> Result<Option<Vec<Interval>>> { |
| self.inner.propagate_constraints(interval, inputs) |
| } |
| |
| /// Calculates the [`SortProperties`] of this function based on its |
| /// children's properties. |
| pub fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> { |
| self.inner.output_ordering(inputs) |
| } |
| |
| pub fn preserves_lex_ordering(&self, inputs: &[ExprProperties]) -> Result<bool> { |
| self.inner.preserves_lex_ordering(inputs) |
| } |
| |
| /// See [`ScalarUDFImpl::coerce_types`] for more details. |
| pub fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> { |
| self.inner.coerce_types(arg_types) |
| } |
| |
| /// Returns the documentation for this Scalar UDF. |
| /// |
| /// Documentation can be accessed programmatically as well as |
| /// generating publicly facing documentation. |
| pub fn documentation(&self) -> Option<&Documentation> { |
| self.inner.documentation() |
| } |
| |
| /// Return true if this function is an async function |
| pub fn as_async(&self) -> Option<&AsyncScalarUDF> { |
| self.inner().as_any().downcast_ref::<AsyncScalarUDF>() |
| } |
| } |
| |
| impl<F> From<F> for ScalarUDF |
| where |
| F: ScalarUDFImpl + 'static, |
| { |
| fn from(fun: F) -> Self { |
| Self::new_from_impl(fun) |
| } |
| } |
| |
| /// Arguments passed to [`ScalarUDFImpl::invoke_with_args`] when invoking a |
| /// scalar function. |
| #[derive(Debug, Clone)] |
| pub struct ScalarFunctionArgs { |
| /// The evaluated arguments to the function |
| pub args: Vec<ColumnarValue>, |
| /// Field associated with each arg, if it exists |
| pub arg_fields: Vec<FieldRef>, |
| /// The number of rows in record batch being evaluated |
| pub number_rows: usize, |
| /// The return field of the scalar function returned (from `return_type` |
| /// or `return_field_from_args`) when creating the physical expression |
| /// from the logical expression |
| pub return_field: FieldRef, |
| /// The config options at execution time |
| pub config_options: Arc<ConfigOptions>, |
| } |
| |
| impl ScalarFunctionArgs { |
| /// The return type of the function. See [`Self::return_field`] for more |
| /// details. |
| pub fn return_type(&self) -> &DataType { |
| self.return_field.data_type() |
| } |
| } |
| |
| /// Information about arguments passed to the function |
| /// |
| /// This structure contains metadata about how the function was called |
| /// such as the type of the arguments, any scalar arguments and if the |
| /// arguments can (ever) be null |
| /// |
| /// See [`ScalarUDFImpl::return_field_from_args`] for more information |
| #[derive(Debug)] |
| pub struct ReturnFieldArgs<'a> { |
| /// The data types of the arguments to the function |
| pub arg_fields: &'a [FieldRef], |
| /// Is argument `i` to the function a scalar (constant)? |
| /// |
| /// If the argument `i` is not a scalar, it will be None |
| /// |
| /// For example, if a function is called like `my_function(column_a, 5)` |
| /// this field will be `[None, Some(ScalarValue::Int32(Some(5)))]` |
| pub scalar_arguments: &'a [Option<&'a ScalarValue>], |
| } |
| |
| /// Trait for implementing user defined scalar functions. |
| /// |
| /// This trait exposes the full API for implementing user defined functions and |
| /// can be used to implement any function. |
| /// |
| /// See [`advanced_udf.rs`] for a full example with complete implementation and |
| /// [`ScalarUDF`] for other available options. |
| /// |
| /// [`advanced_udf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/advanced_udf.rs |
| /// |
| /// # Basic Example |
| /// ``` |
| /// # use std::any::Any; |
| /// # use std::sync::LazyLock; |
| /// # use arrow::datatypes::DataType; |
| /// # use datafusion_common::{DataFusionError, plan_err, Result}; |
| /// # use datafusion_expr::{col, ColumnarValue, Documentation, ScalarFunctionArgs, Signature, Volatility}; |
| /// # use datafusion_expr::{ScalarUDFImpl, ScalarUDF}; |
| /// # use datafusion_expr::scalar_doc_sections::DOC_SECTION_MATH; |
| /// /// This struct for a simple UDF that adds one to an int32 |
| /// #[derive(Debug, PartialEq, Eq, Hash)] |
| /// struct AddOne { |
| /// signature: Signature, |
| /// } |
| /// |
| /// impl AddOne { |
| /// fn new() -> Self { |
| /// Self { |
| /// signature: Signature::uniform(1, vec![DataType::Int32], Volatility::Immutable), |
| /// } |
| /// } |
| /// } |
| /// |
| /// static DOCUMENTATION: LazyLock<Documentation> = LazyLock::new(|| { |
| /// Documentation::builder(DOC_SECTION_MATH, "Add one to an int32", "add_one(2)") |
| /// .with_argument("arg1", "The int32 number to add one to") |
| /// .build() |
| /// }); |
| /// |
| /// fn get_doc() -> &'static Documentation { |
| /// &DOCUMENTATION |
| /// } |
| /// |
| /// /// Implement the ScalarUDFImpl trait for AddOne |
| /// impl ScalarUDFImpl for AddOne { |
| /// fn as_any(&self) -> &dyn Any { self } |
| /// fn name(&self) -> &str { "add_one" } |
| /// fn signature(&self) -> &Signature { &self.signature } |
| /// fn return_type(&self, args: &[DataType]) -> Result<DataType> { |
| /// if !matches!(args.get(0), Some(&DataType::Int32)) { |
| /// return plan_err!("add_one only accepts Int32 arguments"); |
| /// } |
| /// Ok(DataType::Int32) |
| /// } |
| /// // The actual implementation would add one to the argument |
| /// fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { |
| /// unimplemented!() |
| /// } |
| /// fn documentation(&self) -> Option<&Documentation> { |
| /// Some(get_doc()) |
| /// } |
| /// } |
| /// |
| /// // Create a new ScalarUDF from the implementation |
| /// let add_one = ScalarUDF::from(AddOne::new()); |
| /// |
| /// // Call the function `add_one(col)` |
| /// let expr = add_one.call(vec![col("a")]); |
| /// ``` |
| pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync { |
| /// Returns this object as an [`Any`] trait object |
| fn as_any(&self) -> &dyn Any; |
| |
| /// Returns this function's name |
| fn name(&self) -> &str; |
| |
| /// Returns any aliases (alternate names) for this function. |
| /// |
| /// Aliases can be used to invoke the same function using different names. |
| /// For example in some databases `now()` and `current_timestamp()` are |
| /// aliases for the same function. This behavior can be obtained by |
| /// returning `current_timestamp` as an alias for the `now` function. |
| /// |
| /// Note: `aliases` should only include names other than [`Self::name`]. |
| /// Defaults to `[]` (no aliases) |
| fn aliases(&self) -> &[String] { |
| &[] |
| } |
| |
| /// Returns the user-defined display name of function, given the arguments |
| /// |
| /// This can be used to customize the output column name generated by this |
| /// function. |
| /// |
| /// Defaults to `name(args[0], args[1], ...)` |
| #[deprecated( |
| since = "50.0.0", |
| note = "This method is unused and will be removed in a future release" |
| )] |
| fn display_name(&self, args: &[Expr]) -> Result<String> { |
| let names: Vec<String> = args.iter().map(ToString::to_string).collect(); |
| // TODO: join with ", " to standardize the formatting of Vec<Expr>, <https://github.com/apache/datafusion/issues/10364> |
| Ok(format!("{}({})", self.name(), names.join(","))) |
| } |
| |
| /// Returns the name of the column this expression would create |
| /// |
| /// See [`Expr::schema_name`] for details |
| fn schema_name(&self, args: &[Expr]) -> Result<String> { |
| Ok(format!( |
| "{}({})", |
| self.name(), |
| schema_name_from_exprs_comma_separated_without_space(args)? |
| )) |
| } |
| |
| /// Returns a [`Signature`] describing the argument types for which this |
| /// function has an implementation, and the function's [`Volatility`]. |
| /// |
| /// See [`Signature`] for more details on argument type handling |
| /// and [`Self::return_type`] for computing the return type. |
| /// |
| /// [`Volatility`]: datafusion_expr_common::signature::Volatility |
| fn signature(&self) -> &Signature; |
| |
| /// [`DataType`] returned by this function, given the types of the |
| /// arguments. |
| /// |
| /// # Arguments |
| /// |
| /// `arg_types` Data types of the arguments. The implementation of |
| /// `return_type` can assume that some other part of the code has coerced |
| /// the actual argument types to match [`Self::signature`]. |
| /// |
| /// # Notes |
| /// |
| /// If you provide an implementation for [`Self::return_field_from_args`], |
| /// DataFusion will not call `return_type` (this function). While it is |
| /// valid to to put [`unimplemented!()`] or [`unreachable!()`], it is |
| /// recommended to return [`DataFusionError::Internal`] instead, which |
| /// reduces the severity of symptoms if bugs occur (an error rather than a |
| /// panic). |
| /// |
| /// [`DataFusionError::Internal`]: datafusion_common::DataFusionError::Internal |
| fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>; |
| |
| /// Create a new instance of this function with updated configuration. |
| /// |
| /// This method is called when configuration options change at runtime |
| /// (e.g., via `SET` statements) to allow functions that depend on |
| /// configuration to update themselves accordingly. |
| /// |
| /// Note the current [`ConfigOptions`] are also passed to [`Self::invoke_with_args`] so |
| /// this API is not needed for functions where the values may |
| /// depend on the current options. |
| /// |
| /// This API is useful for functions where the return |
| /// **type** depends on the configuration options, such as the `now()` function |
| /// which depends on the current timezone. |
| /// |
| /// # Arguments |
| /// |
| /// * `config` - The updated configuration options |
| /// |
| /// # Returns |
| /// |
| /// * `Some(ScalarUDF)` - A new instance of this function configured with the new settings |
| /// * `None` - If this function does not change with new configuration settings (the default) |
| fn with_updated_config(&self, _config: &ConfigOptions) -> Option<ScalarUDF> { |
| None |
| } |
| |
| /// What type will be returned by this function, given the arguments? |
| /// |
| /// By default, this function calls [`Self::return_type`] with the |
| /// types of each argument. |
| /// |
| /// # Notes |
| /// |
| /// For the majority of UDFs, implementing [`Self::return_type`] is sufficient, |
| /// as the result type is typically a deterministic function of the input types |
| /// (e.g., `sqrt(f32)` consistently yields `f32`). Implementing this method directly |
| /// is generally unnecessary unless the return type depends on runtime values. |
| /// |
| /// This function can be used for more advanced cases such as: |
| /// |
| /// 1. specifying nullability |
| /// 2. return types based on the **values** of the arguments (rather than |
| /// their **types**. |
| /// |
| /// # Example creating `Field` |
| /// |
| /// Note the name of the [`Field`] is ignored, except for structured types such as |
| /// `DataType::Struct`. |
| /// |
| /// ```rust |
| /// # use std::sync::Arc; |
| /// # use arrow::datatypes::{DataType, Field, FieldRef}; |
| /// # use datafusion_common::Result; |
| /// # use datafusion_expr::ReturnFieldArgs; |
| /// # struct Example{} |
| /// # impl Example { |
| /// fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { |
| /// // report output is only nullable if any one of the arguments are nullable |
| /// let nullable = args.arg_fields.iter().any(|f| f.is_nullable()); |
| /// let field = Arc::new(Field::new("ignored_name", DataType::Int32, true)); |
| /// Ok(field) |
| /// } |
| /// # } |
| /// ``` |
| /// |
| /// # Output Type based on Values |
| /// |
| /// For example, the following two function calls get the same argument |
| /// types (something and a `Utf8` string) but return different types based |
| /// on the value of the second argument: |
| /// |
| /// * `arrow_cast(x, 'Int16')` --> `Int16` |
| /// * `arrow_cast(x, 'Float32')` --> `Float32` |
| /// |
| /// # Requirements |
| /// |
| /// This function **must** consistently return the same type for the same |
| /// logical input even if the input is simplified (e.g. it must return the same |
| /// value for `('foo' | 'bar')` as it does for ('foobar'). |
| fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { |
| let data_types = args |
| .arg_fields |
| .iter() |
| .map(|f| f.data_type()) |
| .cloned() |
| .collect::<Vec<_>>(); |
| let return_type = self.return_type(&data_types)?; |
| Ok(Arc::new(Field::new(self.name(), return_type, true))) |
| } |
| |
| #[deprecated( |
| since = "45.0.0", |
| note = "Use `return_field_from_args` instead. if you use `is_nullable` that returns non-nullable with `return_type`, you would need to switch to `return_field_from_args`, you might have error" |
| )] |
| fn is_nullable(&self, _args: &[Expr], _schema: &dyn ExprSchema) -> bool { |
| true |
| } |
| |
| /// Invoke the function returning the appropriate result. |
| /// |
| /// # Performance |
| /// |
| /// For the best performance, the implementations should handle the common case |
| /// when one or more of their arguments are constant values (aka |
| /// [`ColumnarValue::Scalar`]). |
| /// |
| /// [`ColumnarValue::values_to_arrays`] can be used to convert the arguments |
| /// to arrays, which will likely be simpler code, but be slower. |
| fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue>; |
| |
| /// Optionally apply per-UDF simplification / rewrite rules. |
| /// |
| /// This can be used to apply function specific simplification rules during |
| /// optimization (e.g. `arrow_cast` --> `Expr::Cast`). The default |
| /// implementation does nothing. |
| /// |
| /// Note that DataFusion handles simplifying arguments and "constant |
| /// folding" (replacing a function call with constant arguments such as |
| /// `my_add(1,2) --> 3` ). Thus, there is no need to implement such |
| /// optimizations manually for specific UDFs. |
| /// |
| /// # Arguments |
| /// * `args`: The arguments of the function |
| /// * `info`: The necessary information for simplification |
| /// |
| /// # Returns |
| /// [`ExprSimplifyResult`] indicating the result of the simplification NOTE |
| /// if the function cannot be simplified, the arguments *MUST* be returned |
| /// unmodified |
| /// |
| /// # Notes |
| /// |
| /// The returned expression must have the same schema as the original |
| /// expression, including both the data type and nullability. For example, |
| /// if the original expression is nullable, the returned expression must |
| /// also be nullable, otherwise it may lead to schema verification errors |
| /// later in query planning. |
| fn simplify( |
| &self, |
| args: Vec<Expr>, |
| _info: &dyn SimplifyInfo, |
| ) -> Result<ExprSimplifyResult> { |
| Ok(ExprSimplifyResult::Original(args)) |
| } |
| |
| /// Returns true if some of this `exprs` subexpressions may not be evaluated |
| /// and thus any side effects (like divide by zero) may not be encountered. |
| /// |
| /// Setting this to true prevents certain optimizations such as common |
| /// subexpression elimination |
| /// |
| /// When overriding this function to return `true`, [ScalarUDFImpl::conditional_arguments] can also be |
| /// overridden to report more accurately which arguments are eagerly evaluated and which ones |
| /// lazily. |
| fn short_circuits(&self) -> bool { |
| false |
| } |
| |
| /// Determines which of the arguments passed to this function are evaluated eagerly |
| /// and which may be evaluated lazily. |
| /// |
| /// If this function returns `None`, all arguments are eagerly evaluated. |
| /// Returning `None` is a micro optimization that saves a needless `Vec` |
| /// allocation. |
| /// |
| /// If the function returns `Some`, returns (`eager`, `lazy`) where `eager` |
| /// are the arguments that are always evaluated, and `lazy` are the |
| /// arguments that may be evaluated lazily (i.e. may not be evaluated at all |
| /// in some cases). |
| /// |
| /// Implementations must ensure that the two returned `Vec`s are disjunct, |
| /// and that each argument from `args` is present in one the two `Vec`s. |
| /// |
| /// When overriding this function, [ScalarUDFImpl::short_circuits] must |
| /// be overridden to return `true`. |
| fn conditional_arguments<'a>( |
| &self, |
| args: &'a [Expr], |
| ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> { |
| if self.short_circuits() { |
| Some((vec![], args.iter().collect())) |
| } else { |
| None |
| } |
| } |
| |
| /// Computes the output [`Interval`] for a [`ScalarUDFImpl`], given the input |
| /// intervals. |
| /// |
| /// # Parameters |
| /// |
| /// * `children` are the intervals for the children (inputs) of this function. |
| /// |
| /// # Example |
| /// |
| /// If the function is `ABS(a)`, and the input interval is `a: [-3, 2]`, |
| /// then the output interval would be `[0, 3]`. |
| fn evaluate_bounds(&self, _input: &[&Interval]) -> Result<Interval> { |
| // We cannot assume the input datatype is the same of output type. |
| Interval::make_unbounded(&DataType::Null) |
| } |
| |
| /// Updates bounds for child expressions, given a known [`Interval`]s for this |
| /// function. |
| /// |
| /// This function is used to propagate constraints down through an |
| /// expression tree. |
| /// |
| /// # Parameters |
| /// |
| /// * `interval` is the currently known interval for this function. |
| /// * `inputs` are the current intervals for the inputs (children) of this function. |
| /// |
| /// # Returns |
| /// |
| /// A `Vec` of new intervals for the children, in order. |
| /// |
| /// If constraint propagation reveals an infeasibility for any child, returns |
| /// [`None`]. If none of the children intervals change as a result of |
| /// propagation, may return an empty vector instead of cloning `children`. |
| /// This is the default (and conservative) return value. |
| /// |
| /// # Example |
| /// |
| /// If the function is `ABS(a)`, the current `interval` is `[4, 5]` and the |
| /// input `a` is given as `[-7, 3]`, then propagation would return `[-5, 3]`. |
| fn propagate_constraints( |
| &self, |
| _interval: &Interval, |
| _inputs: &[&Interval], |
| ) -> Result<Option<Vec<Interval>>> { |
| Ok(Some(vec![])) |
| } |
| |
| /// Calculates the [`SortProperties`] of this function based on its children's properties. |
| fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> { |
| if !self.preserves_lex_ordering(inputs)? { |
| return Ok(SortProperties::Unordered); |
| } |
| |
| let Some(first_order) = inputs.first().map(|p| &p.sort_properties) else { |
| return Ok(SortProperties::Singleton); |
| }; |
| |
| if inputs |
| .iter() |
| .skip(1) |
| .all(|input| &input.sort_properties == first_order) |
| { |
| Ok(*first_order) |
| } else { |
| Ok(SortProperties::Unordered) |
| } |
| } |
| |
| /// Returns true if the function preserves lexicographical ordering based on |
| /// the input ordering. |
| /// |
| /// For example, `concat(a || b)` preserves lexicographical ordering, but `abs(a)` does not. |
| fn preserves_lex_ordering(&self, _inputs: &[ExprProperties]) -> Result<bool> { |
| Ok(false) |
| } |
| |
| /// Coerce arguments of a function call to types that the function can evaluate. |
| /// |
| /// This function is only called if [`ScalarUDFImpl::signature`] returns |
| /// [`crate::TypeSignature::UserDefined`]. Most UDFs should return one of |
| /// the other variants of [`TypeSignature`] which handle common cases. |
| /// |
| /// See the [type coercion module](crate::type_coercion) |
| /// documentation for more details on type coercion |
| /// |
| /// [`TypeSignature`]: crate::TypeSignature |
| /// |
| /// For example, if your function requires a floating point arguments, but the user calls |
| /// it like `my_func(1::int)` (i.e. with `1` as an integer), coerce_types can return `[DataType::Float64]` |
| /// to ensure the argument is converted to `1::double` |
| /// |
| /// # Parameters |
| /// * `arg_types`: The argument types of the arguments this function with |
| /// |
| /// # Return value |
| /// A Vec the same length as `arg_types`. DataFusion will `CAST` the function call |
| /// arguments to these specific types. |
| fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> { |
| not_impl_err!("Function {} does not implement coerce_types", self.name()) |
| } |
| |
| /// Returns the documentation for this Scalar UDF. |
| /// |
| /// Documentation can be accessed programmatically as well as generating |
| /// publicly facing documentation. |
| fn documentation(&self) -> Option<&Documentation> { |
| None |
| } |
| } |
| |
| /// ScalarUDF that adds an alias to the underlying function. It is better to |
| /// implement [`ScalarUDFImpl`], which supports aliases, directly if possible. |
| #[derive(Debug, PartialEq, Eq, Hash)] |
| struct AliasedScalarUDFImpl { |
| inner: UdfEq<Arc<dyn ScalarUDFImpl>>, |
| aliases: Vec<String>, |
| } |
| |
| impl AliasedScalarUDFImpl { |
| pub fn new( |
| inner: Arc<dyn ScalarUDFImpl>, |
| new_aliases: impl IntoIterator<Item = &'static str>, |
| ) -> Self { |
| let mut aliases = inner.aliases().to_vec(); |
| aliases.extend(new_aliases.into_iter().map(|s| s.to_string())); |
| Self { |
| inner: inner.into(), |
| aliases, |
| } |
| } |
| } |
| |
| #[warn(clippy::missing_trait_methods)] // Delegates, so it should implement every single trait method |
| impl ScalarUDFImpl for AliasedScalarUDFImpl { |
| fn as_any(&self) -> &dyn Any { |
| self |
| } |
| |
| fn name(&self) -> &str { |
| self.inner.name() |
| } |
| |
| fn display_name(&self, args: &[Expr]) -> Result<String> { |
| #[expect(deprecated)] |
| self.inner.display_name(args) |
| } |
| |
| fn schema_name(&self, args: &[Expr]) -> Result<String> { |
| self.inner.schema_name(args) |
| } |
| |
| fn signature(&self) -> &Signature { |
| self.inner.signature() |
| } |
| |
| fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { |
| self.inner.return_type(arg_types) |
| } |
| |
| fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { |
| self.inner.return_field_from_args(args) |
| } |
| |
| fn is_nullable(&self, args: &[Expr], schema: &dyn ExprSchema) -> bool { |
| #[allow(deprecated)] |
| self.inner.is_nullable(args, schema) |
| } |
| |
| fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { |
| self.inner.invoke_with_args(args) |
| } |
| |
| fn with_updated_config(&self, _config: &ConfigOptions) -> Option<ScalarUDF> { |
| None |
| } |
| |
| fn aliases(&self) -> &[String] { |
| &self.aliases |
| } |
| |
| fn simplify( |
| &self, |
| args: Vec<Expr>, |
| info: &dyn SimplifyInfo, |
| ) -> Result<ExprSimplifyResult> { |
| self.inner.simplify(args, info) |
| } |
| |
| fn conditional_arguments<'a>( |
| &self, |
| args: &'a [Expr], |
| ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> { |
| self.inner.conditional_arguments(args) |
| } |
| |
| fn short_circuits(&self) -> bool { |
| self.inner.short_circuits() |
| } |
| |
| fn evaluate_bounds(&self, input: &[&Interval]) -> Result<Interval> { |
| self.inner.evaluate_bounds(input) |
| } |
| |
| fn propagate_constraints( |
| &self, |
| interval: &Interval, |
| inputs: &[&Interval], |
| ) -> Result<Option<Vec<Interval>>> { |
| self.inner.propagate_constraints(interval, inputs) |
| } |
| |
| fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> { |
| self.inner.output_ordering(inputs) |
| } |
| |
| fn preserves_lex_ordering(&self, inputs: &[ExprProperties]) -> Result<bool> { |
| self.inner.preserves_lex_ordering(inputs) |
| } |
| |
| fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> { |
| self.inner.coerce_types(arg_types) |
| } |
| |
| fn documentation(&self) -> Option<&Documentation> { |
| self.inner.documentation() |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| use datafusion_expr_common::signature::Volatility; |
| use std::hash::DefaultHasher; |
| |
| #[derive(Debug, PartialEq, Eq, Hash)] |
| struct TestScalarUDFImpl { |
| name: &'static str, |
| field: &'static str, |
| signature: Signature, |
| } |
| impl ScalarUDFImpl for TestScalarUDFImpl { |
| fn as_any(&self) -> &dyn Any { |
| self |
| } |
| |
| fn name(&self) -> &str { |
| self.name |
| } |
| |
| fn signature(&self) -> &Signature { |
| &self.signature |
| } |
| |
| fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { |
| unimplemented!() |
| } |
| |
| fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> { |
| unimplemented!() |
| } |
| } |
| |
| // PartialEq and Hash must be consistent, and also PartialEq and PartialOrd |
| // must be consistent, so they are tested together. |
| #[test] |
| fn test_partial_eq_hash_and_partial_ord() { |
| // A parameterized function |
| let f = test_func("foo", "a"); |
| |
| // Same like `f`, different instance |
| let f2 = test_func("foo", "a"); |
| assert_eq!(f, f2); |
| assert_eq!(hash(&f), hash(&f2)); |
| assert_eq!(f.partial_cmp(&f2), Some(Ordering::Equal)); |
| |
| // Different parameter |
| let b = test_func("foo", "b"); |
| assert_ne!(f, b); |
| assert_ne!(hash(&f), hash(&b)); // hash can collide for different values but does not collide in this test |
| assert_eq!(f.partial_cmp(&b), None); |
| |
| // Different name |
| let o = test_func("other", "a"); |
| assert_ne!(f, o); |
| assert_ne!(hash(&f), hash(&o)); // hash can collide for different values but does not collide in this test |
| assert_eq!(f.partial_cmp(&o), Some(Ordering::Less)); |
| |
| // Different name and parameter |
| assert_ne!(b, o); |
| assert_ne!(hash(&b), hash(&o)); // hash can collide for different values but does not collide in this test |
| assert_eq!(b.partial_cmp(&o), Some(Ordering::Less)); |
| } |
| |
| fn test_func(name: &'static str, parameter: &'static str) -> ScalarUDF { |
| ScalarUDF::from(TestScalarUDFImpl { |
| name, |
| field: parameter, |
| signature: Signature::any(1, Volatility::Immutable), |
| }) |
| } |
| |
| fn hash<T: Hash>(value: &T) -> u64 { |
| let hasher = &mut DefaultHasher::new(); |
| value.hash(hasher); |
| hasher.finish() |
| } |
| } |