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())
+    }
+}

Reply via email to