datafusion_python/expr/
create_function.rs

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
18use std::{
19    fmt::{self, Display, Formatter},
20    sync::Arc,
21};
22
23use datafusion::logical_expr::{
24    CreateFunction, CreateFunctionBody, OperateFunctionArg, Volatility,
25};
26use pyo3::{prelude::*, IntoPyObjectExt};
27
28use super::logical_node::LogicalNode;
29use super::PyExpr;
30use crate::common::{data_type::PyDataType, df_schema::PyDFSchema};
31use crate::sql::logical::PyLogicalPlan;
32
33#[pyclass(name = "CreateFunction", module = "datafusion.expr", subclass)]
34#[derive(Clone)]
35pub struct PyCreateFunction {
36    create: CreateFunction,
37}
38
39impl From<PyCreateFunction> for CreateFunction {
40    fn from(create: PyCreateFunction) -> Self {
41        create.create
42    }
43}
44
45impl From<CreateFunction> for PyCreateFunction {
46    fn from(create: CreateFunction) -> PyCreateFunction {
47        PyCreateFunction { create }
48    }
49}
50
51impl Display for PyCreateFunction {
52    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
53        write!(f, "CreateFunction: name {:?}", self.create.name)
54    }
55}
56
57#[pyclass(name = "OperateFunctionArg", module = "datafusion.expr", subclass)]
58#[derive(Clone)]
59pub struct PyOperateFunctionArg {
60    arg: OperateFunctionArg,
61}
62
63#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
64#[pyclass(eq, eq_int, name = "Volatility", module = "datafusion.expr")]
65pub enum PyVolatility {
66    Immutable,
67    Stable,
68    Volatile,
69}
70
71#[pyclass(name = "CreateFunctionBody", module = "datafusion.expr", subclass)]
72#[derive(Clone)]
73pub struct PyCreateFunctionBody {
74    body: CreateFunctionBody,
75}
76
77#[pymethods]
78impl PyCreateFunctionBody {
79    pub fn language(&self) -> Option<String> {
80        self.body
81            .language
82            .as_ref()
83            .map(|language| language.to_string())
84    }
85
86    pub fn behavior(&self) -> Option<PyVolatility> {
87        self.body.behavior.as_ref().map(|behavior| match behavior {
88            Volatility::Immutable => PyVolatility::Immutable,
89            Volatility::Stable => PyVolatility::Stable,
90            Volatility::Volatile => PyVolatility::Volatile,
91        })
92    }
93
94    pub fn function_body(&self) -> Option<PyExpr> {
95        self.body
96            .function_body
97            .as_ref()
98            .map(|function_body| function_body.clone().into())
99    }
100}
101
102#[pymethods]
103impl PyCreateFunction {
104    #[new]
105    #[pyo3(signature = (or_replace, temporary, name, params, schema, return_type=None, args=None))]
106    pub fn new(
107        or_replace: bool,
108        temporary: bool,
109        name: String,
110        params: PyCreateFunctionBody,
111        schema: PyDFSchema,
112        return_type: Option<PyDataType>,
113        args: Option<Vec<PyOperateFunctionArg>>,
114    ) -> Self {
115        PyCreateFunction {
116            create: CreateFunction {
117                or_replace,
118                temporary,
119                name,
120                args: args.map(|args| args.into_iter().map(|arg| arg.arg).collect()),
121                return_type: return_type.map(|return_type| return_type.data_type),
122                params: params.body,
123                schema: Arc::new(schema.into()),
124            },
125        }
126    }
127
128    pub fn or_replace(&self) -> bool {
129        self.create.or_replace
130    }
131
132    pub fn temporary(&self) -> bool {
133        self.create.temporary
134    }
135
136    pub fn name(&self) -> String {
137        self.create.name.clone()
138    }
139
140    pub fn params(&self) -> PyCreateFunctionBody {
141        PyCreateFunctionBody {
142            body: self.create.params.clone(),
143        }
144    }
145
146    pub fn schema(&self) -> PyDFSchema {
147        (*self.create.schema).clone().into()
148    }
149
150    pub fn return_type(&self) -> Option<PyDataType> {
151        self.create
152            .return_type
153            .as_ref()
154            .map(|return_type| return_type.clone().into())
155    }
156
157    pub fn args(&self) -> Option<Vec<PyOperateFunctionArg>> {
158        self.create.args.as_ref().map(|args| {
159            args.iter()
160                .map(|arg| PyOperateFunctionArg { arg: arg.clone() })
161                .collect()
162        })
163    }
164
165    fn __repr__(&self) -> PyResult<String> {
166        Ok(format!("CreateFunction({})", self))
167    }
168
169    fn __name__(&self) -> PyResult<String> {
170        Ok("CreateFunction".to_string())
171    }
172}
173
174impl LogicalNode for PyCreateFunction {
175    fn inputs(&self) -> Vec<PyLogicalPlan> {
176        vec![]
177    }
178
179    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
180        self.clone().into_bound_py_any(py)
181    }
182}
datafusion_python/expr/create_function.rs

datafusion_python/expr/
create_function.rs
