Feature/expose when function (#836)

Author: timsaucer

python/datafusion/functions.py

@@ -245,6 +245,7 @@
     "var",
     "var_pop",
     "var_samp",
+    "when",
     "window",
 ]
 
@@ -364,6 +365,16 @@ def case(expr: Expr) -> CaseBuilder:
     return CaseBuilder(f.case(expr.expr))
 
 
+def when(when: Expr, then: Expr) -> CaseBuilder:
+    """Create a case expression that has no base expression.
+
+    Create a :py:class:`~datafusion.expr.CaseBuilder` to match cases for the
+    expression ``expr``. See :py:class:`~datafusion.expr.CaseBuilder` for
+    detailed usage.
+    """
+    return CaseBuilder(f.when(when.expr, then.expr))
+
+
 def window(
     name: str,
     args: list[Expr],

python/datafusion/tests/test_functions.py

@@ -836,6 +836,25 @@ def test_case(df):
     assert result.column(2) == pa.array(["Hola", "Mundo", None])
 
 
+def test_when_with_no_base(df):
+    df.show()
+    df = df.select(
+        column("b"),
+        f.when(column("b") > literal(5), literal("too big"))
+        .when(column("b") < literal(5), literal("too small"))
+        .otherwise(literal("just right"))
+        .alias("goldilocks"),
+        f.when(column("a") == literal("Hello"), column("a")).end().alias("greeting"),
+    )
+    df.show()
+
+    result = df.collect()
+    result = result[0]
+    assert result.column(0) == pa.array([4, 5, 6])
+    assert result.column(1) == pa.array(["too small", "just right", "too big"])
+    assert result.column(2) == pa.array(["Hello", None, None])
+
+
 def test_regr_funcs_sql(df):
     # test case base on
     # https://github.com/apache/arrow-datafusion/blob/d1361d56b9a9e0c165d3d71a8df6795d2a5f51dd/datafusion/core/tests/sqllogictests/test_files/aggregate.slt#L2330

src/functions.rs

@@ -533,6 +533,14 @@ fn case(expr: PyExpr) -> PyResult<PyCaseBuilder> {
     })
 }
 
+/// Create a CASE WHEN statement with literal WHEN expressions for comparison to the base expression.
+#[pyfunction]
+fn when(when: PyExpr, then: PyExpr) -> PyResult<PyCaseBuilder> {
+    Ok(PyCaseBuilder {
+        case_builder: datafusion_expr::when(when.expr, then.expr),
+    })
+}
+
 /// Helper function to find the appropriate window function.
 ///
 /// Search procedure:
@@ -910,6 +918,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(char_length))?;
     m.add_wrapped(wrap_pyfunction!(coalesce))?;
     m.add_wrapped(wrap_pyfunction!(case))?;
+    m.add_wrapped(wrap_pyfunction!(when))?;
     m.add_wrapped(wrap_pyfunction!(col))?;
     m.add_wrapped(wrap_pyfunction!(concat_ws))?;
     m.add_wrapped(wrap_pyfunction!(concat))?;