From 7b9994b86e2c898ef5c59f236fba0e30cf7b4c5c Mon Sep 17 00:00:00 2001 From: jiangzhx Date: Tue, 12 Sep 2023 15:12:02 +0800 Subject: [PATCH] add first_value last_value --- datafusion/tests/test_functions.py | 23 +++++++++++++++++++++++ src/functions.rs | 4 ++++ 2 files changed, 27 insertions(+) diff --git a/datafusion/tests/test_functions.py b/datafusion/tests/test_functions.py index f1f64c30a..e504cc498 100644 --- a/datafusion/tests/test_functions.py +++ b/datafusion/tests/test_functions.py @@ -479,6 +479,29 @@ def test_case(df): assert result.column(2) == pa.array(["Hola", "Mundo", None]) +def test_first_last_value(df): + df = df.aggregate( + [], + [ + f.first_value(column("a")), + f.first_value(column("b")), + f.first_value(column("d")), + f.last_value(column("a")), + f.last_value(column("b")), + f.last_value(column("d")), + ], + ) + + result = df.collect() + result = result[0] + assert result.column(0) == pa.array(["Hello"]) + assert result.column(1) == pa.array([4]) + assert result.column(2) == pa.array([datetime(2022, 12, 31)]) + assert result.column(3) == pa.array(["!"]) + assert result.column(4) == pa.array([6]) + assert result.column(5) == pa.array([datetime(2020, 7, 2)]) + + def test_binary_string_functions(df): df = df.select( f.encode(column("a"), literal("base64")), diff --git a/src/functions.rs b/src/functions.rs index ef26240fe..87e5f5b1d 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -362,6 +362,8 @@ aggregate_function!(stddev_samp, Stddev); aggregate_function!(var, Variance); aggregate_function!(var_pop, VariancePop); aggregate_function!(var_samp, Variance); +aggregate_function!(first_value, FirstValue); +aggregate_function!(last_value, LastValue); pub(crate) fn init_module(m: &PyModule) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(abs))?; @@ -489,6 +491,8 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(var_pop))?; m.add_wrapped(wrap_pyfunction!(var_samp))?; m.add_wrapped(wrap_pyfunction!(window))?; + m.add_wrapped(wrap_pyfunction!(first_value))?; + m.add_wrapped(wrap_pyfunction!(last_value))?; //Binary String Functions m.add_wrapped(wrap_pyfunction!(encode))?;