Trait datafusion_expr::ScalarUDFImpl

pub trait ScalarUDFImpl: Debug + Send + Sync {
    // Required methods
    fn as_any(&self) -> &dyn Any;
    fn name(&self) -> &str;
    fn signature(&self) -> &Signature;
    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>;
    fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue>;

    // Provided methods
    fn display_name(&self, args: &[Expr]) -> Result<String> { ... }
    fn return_type_from_exprs(
        &self,
        _args: &[Expr],
        _schema: &dyn ExprSchema,
        arg_types: &[DataType],
    ) -> Result<DataType> { ... }
    fn invoke_no_args(&self, _number_rows: usize) -> Result<ColumnarValue> { ... }
    fn aliases(&self) -> &[String] { ... }
    fn simplify(
        &self,
        args: Vec<Expr>,
        _info: &dyn SimplifyInfo,
    ) -> Result<ExprSimplifyResult> { ... }
    fn short_circuits(&self) -> bool { ... }
    fn evaluate_bounds(&self, _input: &[&Interval]) -> Result<Interval> { ... }
    fn propagate_constraints(
        &self,
        _interval: &Interval,
        _inputs: &[&Interval],
    ) -> Result<Option<Vec<Interval>>> { ... }
    fn output_ordering(
        &self,
        _inputs: &[ExprProperties],
    ) -> Result<SortProperties> { ... }
    fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> { ... }
}

Trait for implementing ScalarUDF.

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.

Basic Example

#[derive(Debug)]
struct AddOne {
  signature: Signature
};

impl AddOne {
  fn new() -> Self {
    Self {
      signature: Signature::uniform(1, vec![DataType::Int32], Volatility::Immutable)
     }
  }
}

/// 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(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> { unimplemented!() }
}

// 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")]);

Required Methods

fn as_any(&self) -> &dyn Any

Returns this object as an Any trait object

fn name(&self) -> &str

Returns this function’s name

fn signature(&self) -> &Signature

Returns the function’s Signature for information about what input types are accepted and the function’s Volatility.

fn return_type(&self, arg_types: &[DataType]) -> Result<DataType>

What DataType will be returned by this function, given the types of the arguments.

Notes

If you provide an implementation for Self::return_type_from_exprs, DataFusion will not call return_type (this function). In this case it is recommended to return DataFusionError::Internal.

fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue>

Invoke the function on args, returning the appropriate result

The function will be invoked passed with the slice of ColumnarValue (either scalar or array).

If the function does not take any arguments, please use invoke_no_args instead and return not_impl_err for this function.

Performance

For the best performance, the implementations of invoke 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.

Provided Methods

fn display_name(&self, args: &[Expr]) -> Result<String>

Returns the user-defined display name of the UDF given the arguments

fn return_type_from_exprs( &self, _args: &[Expr], _schema: &dyn ExprSchema, arg_types: &[DataType], ) -> Result<DataType>

What DataType will be returned by this function, given the arguments?

Note most UDFs should implement Self::return_type and not this function. The output type for most functions only depends on the types of their inputs (e.g. sqrt(f32) is always f32).

By default, this function calls Self::return_type with the types of each argument.

This method can be overridden for functions that return different types based on the values of their arguments.

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

Notes:

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 invoke_no_args(&self, _number_rows: usize) -> Result<ColumnarValue>

Invoke the function without args, instead the number of rows are provided, returning the appropriate result.

fn aliases(&self) -> &[String]

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 simplify( &self, args: Vec<Expr>, _info: &dyn SimplifyInfo, ) -> Result<ExprSimplifyResult>

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

fn short_circuits(&self) -> bool

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

fn evaluate_bounds(&self, _input: &[&Interval]) -> Result<Interval>

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 propagate_constraints( &self, _interval: &Interval, _inputs: &[&Interval], ) -> Result<Option<Vec<Interval>>>

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].

fn output_ordering(&self, _inputs: &[ExprProperties]) -> Result<SortProperties>

Calculates the SortProperties of this function based on its children’s properties.

fn coerce_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>>

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 documentation for more details on type coercion

For example, if your function requires a floating point arguments, but the user calls it like my_func(1::int) (aka with 1 as an integer), coerce_types could return [DataType::Float64] to ensure the argument was cast 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.

Implementors

impl ScalarUDFImpl for SimpleScalarUDF