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/datafusion-python.git
The following commit(s) were added to refs/heads/main by this push:
new 003eea8 Feature/expose when function (#836)
003eea8 is described below
commit 003eea8fe75bed9c12200ce296a8c7fa5212b534
Author: Tim Saucer <[email protected]>
AuthorDate: Fri Aug 30 17:05:00 2024 -0400
Feature/expose when function (#836)
---
python/datafusion/functions.py | 11 +++++++++++
python/datafusion/tests/test_functions.py | 19 +++++++++++++++++++
src/functions.rs | 9 +++++++++
3 files changed, 39 insertions(+)
diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py
index 59a1974..ec0c110 100644
--- a/python/datafusion/functions.py
+++ b/python/datafusion/functions.py
@@ -245,6 +245,7 @@ __all__ = [
"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],
diff --git a/python/datafusion/tests/test_functions.py
b/python/datafusion/tests/test_functions.py
index 732136e..e5429bd 100644
--- a/python/datafusion/tests/test_functions.py
+++ b/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
diff --git a/src/functions.rs b/src/functions.rs
index c53d4ad..2525636 100644
--- a/src/functions.rs
+++ b/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))?;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
