[arrow-datafusion-python] branch main updated: add first_value last_value (#498)

[agrove]

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/arrow-datafusion-python.git


The following commit(s) were added to refs/heads/main by this push:
     new a91188c  add first_value last_value (#498)
a91188c is described below

commit a91188c51019a1f9c190a6596970914763415965
Author: zhenxing jiang <jiang...@gmail.com>
AuthorDate: Fri Oct 13 17:18:55 2023 -0500

    add first_value last_value (#498)
    
    Co-authored-by: Andy Grove <andygrov...@gmail.com>
---
 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 f1f64c3..e504cc4 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 eed2815..2f2f34e 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);
 aggregate_function!(bit_and, BitAnd);
 aggregate_function!(bit_or, BitOr);
 aggregate_function!(bit_xor, BitXor);
@@ -494,6 +496,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))?;
     m.add_wrapped(wrap_pyfunction!(bit_and))?;
     m.add_wrapped(wrap_pyfunction!(bit_or))?;
     m.add_wrapped(wrap_pyfunction!(bit_xor))?;