Skip to content

Commit

Permalink
add first_value last_value
Browse files Browse the repository at this point in the history
  • Loading branch information
jiangzhx committed Oct 7, 2023
1 parent 9ef0a57 commit 25f815b
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 0 deletions.
23 changes: 23 additions & 0 deletions datafusion/tests/test_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")),
Expand Down
4 changes: 4 additions & 0 deletions src/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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))?;
Expand Down Expand Up @@ -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))?;
Expand Down

0 comments on commit 25f815b

Please sign in to comment.