[arrow-datafusion-python] branch main updated: feat: add case function (#447) (#448)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb 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 1fde8e4  feat: add case function (#447) (#448)
1fde8e4 is described below

commit 1fde8e47ae6714f02c2d2a088577e92e814dfa34
Author: Daniel Mesejo <me...@gmail.com>
AuthorDate: Sat Aug 5 20:12:30 2023 +0200

    feat: add case function (#447) (#448)
---
 datafusion/tests/test_functions.py | 22 ++++++++++++++++
 src/expr.rs                        |  1 +
 src/expr/conditional_expr.rs       | 54 ++++++++++++++++++++++++++++++++++++++
 src/functions.rs                   | 15 ++++++++---
 4 files changed, 89 insertions(+), 3 deletions(-)

diff --git a/datafusion/tests/test_functions.py b/datafusion/tests/test_functions.py
index ec33474..c5ddea3 100644
--- a/datafusion/tests/test_functions.py
+++ b/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])
diff --git a/src/expr.rs b/src/expr.rs
index 2fd638a..d1022e9 100644
--- a/src/expr.rs
+++ b/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;
diff --git a/src/expr/conditional_expr.rs b/src/expr/conditional_expr.rs
new file mode 100644
index 0000000..96ef58f
--- /dev/null
+++ b/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())
+    }
+}
diff --git a/src/functions.rs b/src/functions.rs
index 8b60e64..06d35ef 100644
--- a/src/functions.rs
+++ b/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))?;