From 3ca46b070b68a998bf540ad9cbb8342465bfd3c8 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Thu, 23 Feb 2023 16:36:33 -0500 Subject: [PATCH] Add bindings for case, cast, and trycast (#232) * Add bindings for case, cast, and trycast * cargo clippy * Python linters --- datafusion/__init__.py | 6 +++ datafusion/tests/test_imports.py | 6 +++ src/expr.rs | 5 ++ src/expr/case.rs | 57 +++++++++++++++++++++++ src/expr/cast.rs | 78 ++++++++++++++++++++++++++++++++ 5 files changed, 152 insertions(+) create mode 100644 src/expr/case.rs create mode 100644 src/expr/cast.rs diff --git a/datafusion/__init__.py b/datafusion/__init__.py index 2a81a4d68..3a6482d24 100644 --- a/datafusion/__init__.py +++ b/datafusion/__init__.py @@ -63,6 +63,9 @@ IsNotFalse, IsNotUnknown, Negative, + Case, + Cast, + TryCast, Between, ) @@ -102,6 +105,9 @@ "IsNotFalse", "IsNotUnknown", "Negative", + "Case", + "Cast", + "TryCast", "Between", ] diff --git a/datafusion/tests/test_imports.py b/datafusion/tests/test_imports.py index ccdfeff47..3d4a0d953 100644 --- a/datafusion/tests/test_imports.py +++ b/datafusion/tests/test_imports.py @@ -59,6 +59,9 @@ IsNotFalse, IsNotUnknown, Negative, + Case, + Cast, + TryCast, Between, ) @@ -108,6 +111,9 @@ def test_class_module_is_datafusion(): IsNotFalse, IsNotUnknown, Negative, + Case, + Cast, + TryCast, Between, ]: assert klass.__module__ == "datafusion.expr" diff --git a/src/expr.rs b/src/expr.rs index 3d184fa12..350eb6646 100644 --- a/src/expr.rs +++ b/src/expr.rs @@ -44,6 +44,8 @@ pub mod analyze; pub mod between; pub mod binary_expr; pub mod bool_expr; +pub mod case; +pub mod cast; pub mod column; pub mod empty_relation; pub mod filter; @@ -232,6 +234,9 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; // operators diff --git a/src/expr/case.rs b/src/expr/case.rs new file mode 100644 index 000000000..605275376 --- /dev/null +++ b/src/expr/case.rs @@ -0,0 +1,57 @@ +// 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. + +use crate::expr::PyExpr; +use datafusion_expr::Case; +use pyo3::prelude::*; + +#[pyclass(name = "Case", module = "datafusion.expr", subclass)] +#[derive(Clone)] +pub struct PyCase { + case: Case, +} + +impl From for Case { + fn from(case: PyCase) -> Self { + case.case + } +} + +impl From for PyCase { + fn from(case: Case) -> PyCase { + PyCase { case } + } +} + +#[pymethods] +impl PyCase { + fn expr(&self) -> Option { + self.case.expr.as_ref().map(|e| (**e).clone().into()) + } + + fn when_then_expr(&self) -> Vec<(PyExpr, PyExpr)> { + self.case + .when_then_expr + .iter() + .map(|e| ((*e.0).clone().into(), (*e.1).clone().into())) + .collect() + } + + fn else_expr(&self) -> Option { + self.case.else_expr.as_ref().map(|e| (**e).clone().into()) + } +} diff --git a/src/expr/cast.rs b/src/expr/cast.rs new file mode 100644 index 000000000..a72199876 --- /dev/null +++ b/src/expr/cast.rs @@ -0,0 +1,78 @@ +// 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. + +use crate::{common::data_type::PyDataType, expr::PyExpr}; +use datafusion_expr::{Cast, TryCast}; +use pyo3::prelude::*; + +#[pyclass(name = "Cast", module = "datafusion.expr", subclass)] +#[derive(Clone)] +pub struct PyCast { + cast: Cast, +} + +impl From for Cast { + fn from(cast: PyCast) -> Self { + cast.cast + } +} + +impl From for PyCast { + fn from(cast: Cast) -> PyCast { + PyCast { cast } + } +} + +#[pymethods] +impl PyCast { + fn expr(&self) -> PyResult { + Ok((*self.cast.expr).clone().into()) + } + + fn data_type(&self) -> PyResult { + Ok(self.cast.data_type.clone().into()) + } +} + +#[pyclass(name = "TryCast", module = "datafusion.expr", subclass)] +#[derive(Clone)] +pub struct PyTryCast { + try_cast: TryCast, +} + +impl From for TryCast { + fn from(try_cast: PyTryCast) -> Self { + try_cast.try_cast + } +} + +impl From for PyTryCast { + fn from(try_cast: TryCast) -> PyTryCast { + PyTryCast { try_cast } + } +} + +#[pymethods] +impl PyTryCast { + fn expr(&self) -> PyResult { + Ok((*self.try_cast.expr).clone().into()) + } + + fn data_type(&self) -> PyResult { + Ok(self.try_cast.data_type.clone().into()) + } +}