Add bindings for case, cast, and trycast

datafusion/__init__.py

@@ -62,6 +62,9 @@
     IsNotFalse,
     IsNotUnknown,
     Negative,
+    Case,
+    Cast,
+    TryCast,
     Between,
 )
 
@@ -100,6 +103,9 @@
     "IsNotFalse",
     "IsNotUnknown",
     "Negative",
+    "Case",
+    "Cast",
+    "TryCast",
     "Between",
 ]
 

datafusion/tests/test_imports.py

@@ -58,6 +58,9 @@
     IsNotFalse,
     IsNotUnknown,
     Negative,
+    Case,
+    Cast,
+    TryCast,
     Between,
 )
 
@@ -106,6 +109,9 @@ def test_class_module_is_datafusion():
         IsNotFalse,
         IsNotUnknown,
         Negative,
+        Case,
+        Cast,
+        TryCast,
         Between,
     ]:
         assert klass.__module__ == "datafusion.expr"

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;
@@ -231,6 +233,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>()?;
     // operators
     m.add_class::<table_scan::PyTableScan>()?;

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

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