feat: add case function (#447) (#448)

Author: mesejo

datafusion/tests/test_functions.py

@@ -411,3 +411,25 @@ def test_temporal_functions(df):
     assert result.column(9) == pa.array(
         [datetime(2023, 9, 7, 5, 6, 14, 523952)] * 3, type=pa.timestamp("us")
     )
+
+
+def test_case(df):
+    df = df.select(
+        f.case(column("b"))
+        .when(literal(4), literal(10))
+        .otherwise(literal(8)),
+        f.case(column("a"))
+        .when(literal("Hello"), literal("Hola"))
+        .when(literal("World"), literal("Mundo"))
+        .otherwise(literal("!!")),
+        f.case(column("a"))
+        .when(literal("Hello"), literal("Hola"))
+        .when(literal("World"), literal("Mundo"))
+        .end(),
+    )
+
+    result = df.collect()
+    result = result[0]
+    assert result.column(0) == pa.array([10, 8, 8])
+    assert result.column(1) == pa.array(["Hola", "Mundo", "!!"])
+    assert result.column(2) == pa.array(["Hola", "Mundo", None])

src/expr.rs

@@ -59,6 +59,7 @@ pub mod bool_expr;
 pub mod case;
 pub mod cast;
 pub mod column;
+pub mod conditional_expr;
 pub mod create_memory_table;
 pub mod create_view;
 pub mod cross_join;

src/expr/conditional_expr.rs

@@ -0,0 +1,54 @@
+// 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::conditional_expressions::CaseBuilder;
+use pyo3::prelude::*;
+
+#[pyclass(name = "CaseBuilder", module = "datafusion.expr", subclass)]
+pub struct PyCaseBuilder {
+    pub case_builder: CaseBuilder,
+}
+
+impl From<PyCaseBuilder> for CaseBuilder {
+    fn from(case_builder: PyCaseBuilder) -> Self {
+        case_builder.case_builder
+    }
+}
+
+impl From<CaseBuilder> for PyCaseBuilder {
+    fn from(case_builder: CaseBuilder) -> PyCaseBuilder {
+        PyCaseBuilder { case_builder }
+    }
+}
+
+#[pymethods]
+impl PyCaseBuilder {
+    fn when(&mut self, when: PyExpr, then: PyExpr) -> PyCaseBuilder {
+        PyCaseBuilder {
+            case_builder: self.case_builder.when(when.expr, then.expr),
+        }
+    }
+
+    fn otherwise(&mut self, else_expr: PyExpr) -> PyResult<PyExpr> {
+        Ok(self.case_builder.otherwise(else_expr.expr)?.clone().into())
+    }
+
+    fn end(&mut self) -> PyResult<PyExpr> {
+        Ok(self.case_builder.end()?.clone().into())
+    }
+}

src/functions.rs

@@ -17,6 +17,9 @@
 
 use pyo3::{prelude::*, wrap_pyfunction};
 
+use crate::errors::DataFusionError;
+use crate::expr::conditional_expr::PyCaseBuilder;
+use crate::expr::PyExpr;
 use datafusion_common::Column;
 use datafusion_expr::expr::Alias;
 use datafusion_expr::{
@@ -27,9 +30,6 @@ use datafusion_expr::{
     BuiltinScalarFunction, Expr, WindowFrame,
 };
 
-use crate::errors::DataFusionError;
-use crate::expr::PyExpr;
-
 #[pyfunction]
 fn in_list(expr: PyExpr, value: Vec<PyExpr>, negated: bool) -> PyExpr {
     datafusion_expr::in_list(
@@ -115,6 +115,14 @@ fn count_star() -> PyResult<PyExpr> {
     })
 }
 
+/// Create a CASE WHEN statement with literal WHEN expressions for comparison to the base expression.
+#[pyfunction]
+fn case(expr: PyExpr) -> PyResult<PyCaseBuilder> {
+    Ok(PyCaseBuilder {
+        case_builder: datafusion_expr::case(expr.expr),
+    })
+}
+
 /// Creates a new Window function expression
 #[pyfunction]
 fn window(
@@ -355,6 +363,7 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(chr))?;
     m.add_wrapped(wrap_pyfunction!(char_length))?;
     m.add_wrapped(wrap_pyfunction!(coalesce))?;
+    m.add_wrapped(wrap_pyfunction!(case))?;
     m.add_wrapped(wrap_pyfunction!(col))?;
     m.add_wrapped(wrap_pyfunction!(concat_ws))?;
     m.add_wrapped(wrap_pyfunction!(concat))?;