This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion-python.git
The following commit(s) were added to refs/heads/main by this push:
new 3ca46b0 Add bindings for case, cast, and trycast (#232)
3ca46b0 is described below
commit 3ca46b070b68a998bf540ad9cbb8342465bfd3c8
Author: Jeremy Dyer <[email protected]>
AuthorDate: Thu Feb 23 16:36:33 2023 -0500
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(+)
diff --git a/datafusion/__init__.py b/datafusion/__init__.py
index 2a81a4d..3a6482d 100644
--- a/datafusion/__init__.py
+++ b/datafusion/__init__.py
@@ -63,6 +63,9 @@ from .expr import (
IsNotFalse,
IsNotUnknown,
Negative,
+ Case,
+ Cast,
+ TryCast,
Between,
)
@@ -102,6 +105,9 @@ __all__ = [
"IsNotFalse",
"IsNotUnknown",
"Negative",
+ "Case",
+ "Cast",
+ "TryCast",
"Between",
]
diff --git a/datafusion/tests/test_imports.py b/datafusion/tests/test_imports.py
index ccdfeff..3d4a0d9 100644
--- a/datafusion/tests/test_imports.py
+++ b/datafusion/tests/test_imports.py
@@ -59,6 +59,9 @@ from datafusion.expr import (
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 3d184fa..350eb66 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::<PySimilarTo>()?;
m.add_class::<PyScalarVariable>()?;
m.add_class::<alias::PyAlias>()?;
+ m.add_class::<case::PyCase>()?;
+ m.add_class::<cast::PyCast>()?;
+ m.add_class::<cast::PyTryCast>()?;
m.add_class::<between::PyBetween>()?;
m.add_class::<indexed_field::PyGetIndexedField>()?;
// operators
diff --git a/src/expr/case.rs b/src/expr/case.rs
new file mode 100644
index 0000000..6052753
--- /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<PyCase> for Case {
+ fn from(case: PyCase) -> Self {
+ case.case
+ }
+}
+
+impl From<Case> for PyCase {
+ fn from(case: Case) -> PyCase {
+ PyCase { case }
+ }
+}
+
+#[pymethods]
+impl PyCase {
+ fn expr(&self) -> Option<PyExpr> {
+ 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<PyExpr> {
+ 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 0000000..a721998
--- /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<PyCast> for Cast {
+ fn from(cast: PyCast) -> Self {
+ cast.cast
+ }
+}
+
+impl From<Cast> for PyCast {
+ fn from(cast: Cast) -> PyCast {
+ PyCast { cast }
+ }
+}
+
+#[pymethods]
+impl PyCast {
+ fn expr(&self) -> PyResult<PyExpr> {
+ Ok((*self.cast.expr).clone().into())
+ }
+
+ fn data_type(&self) -> PyResult<PyDataType> {
+ Ok(self.cast.data_type.clone().into())
+ }
+}
+
+#[pyclass(name = "TryCast", module = "datafusion.expr", subclass)]
+#[derive(Clone)]
+pub struct PyTryCast {
+ try_cast: TryCast,
+}
+
+impl From<PyTryCast> for TryCast {
+ fn from(try_cast: PyTryCast) -> Self {
+ try_cast.try_cast
+ }
+}
+
+impl From<TryCast> for PyTryCast {
+ fn from(try_cast: TryCast) -> PyTryCast {
+ PyTryCast { try_cast }
+ }
+}
+
+#[pymethods]
+impl PyTryCast {
+ fn expr(&self) -> PyResult<PyExpr> {
+ Ok((*self.try_cast.expr).clone().into())
+ }
+
+ fn data_type(&self) -> PyResult<PyDataType> {
+ Ok(self.try_cast.data_type.clone().into())
+ }
+}