polars_plan/dsl/
udf.rs

1use polars_utils::pl_str::PlSmallStr;
2
3use super::{ColumnsUdf, Expr, GetOutput, OpaqueColumnUdf};
4use crate::prelude::{FunctionOptions, new_column_udf};
5
6/// Represents a user-defined function
7#[derive(Clone)]
8pub struct UserDefinedFunction {
9    /// name
10    pub name: PlSmallStr,
11    /// The function output type.
12    pub return_type: GetOutput,
13    /// The function implementation.
14    pub fun: OpaqueColumnUdf,
15    /// Options for the function.
16    pub options: FunctionOptions,
17}
18
19impl std::fmt::Debug for UserDefinedFunction {
20    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
21        f.debug_struct("UserDefinedFunction")
22            .field("name", &self.name)
23            .field("fun", &"<FUNC>")
24            .field("options", &self.options)
25            .finish()
26    }
27}
28
29impl UserDefinedFunction {
30    /// Create a new UserDefinedFunction
31    pub fn new(name: PlSmallStr, return_type: GetOutput, fun: impl ColumnsUdf + 'static) -> Self {
32        Self {
33            name,
34            return_type,
35            fun: new_column_udf(fun),
36            options: FunctionOptions::default(),
37        }
38    }
39
40    /// creates a logical expression with a call of the UDF
41    pub fn call(self, args: Vec<Expr>) -> Expr {
42        Expr::AnonymousFunction {
43            input: args,
44            function: self.fun,
45            output_type: self.return_type.clone(),
46            options: self.options,
47        }
48    }
49}