add first_value last_value

Author: jiangzhx

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")),

src/functions.rs

@@ -54,7 +54,7 @@ fn digest(value: PyExpr, method: PyExpr) -> PyExpr {
 /// Concatenates the text representations of all the arguments.
 /// NULL arguments are ignored.
 #[pyfunction]
-#[pyo3(signature = (*args))]
+#[pyo3(signature = (* args))]
 fn concat(args: Vec<PyExpr>) -> PyResult<PyExpr> {
     let args = args.into_iter().map(|e| e.expr).collect::<Vec<_>>();
     Ok(datafusion_expr::concat(&args).into())
@@ -64,7 +64,7 @@ fn concat(args: Vec<PyExpr>) -> PyResult<PyExpr> {
 /// The first argument is used as the separator string, and should not be NULL.
 /// Other NULL arguments are ignored.
 #[pyfunction]
-#[pyo3(signature = (sep, *args))]
+#[pyo3(signature = (sep, * args))]
 fn concat_ws(sep: String, args: Vec<PyExpr>) -> PyResult<PyExpr> {
     let args = args.into_iter().map(|e| e.expr).collect::<Vec<_>>();
     Ok(datafusion_expr::concat_ws(lit(sep), args).into())
@@ -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))?;