Skip to content

Commit 070b595

Browse files
committed
Make first_value and last_value identical in the interface
1 parent bc3be5d commit 070b595

File tree

3 files changed

+53
-15
lines changed

3 files changed

+53
-15
lines changed

python/datafusion/functions.py

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1710,29 +1710,47 @@ def regr_syy(y: Expr, x: Expr, distinct: bool = False) -> Expr:
17101710
def first_value(
17111711
arg: Expr,
17121712
distinct: bool = False,
1713-
filter: bool = None,
1714-
order_by: Expr | None = None,
1715-
null_treatment: common.NullTreatment | None = None,
1713+
filter: Optional[bool] = None,
1714+
order_by: Optional[list[Expr]] = None,
1715+
null_treatment: Optional[common.NullTreatment] = None,
17161716
) -> Expr:
17171717
"""Returns the first value in a group of values."""
1718+
order_by_cols = [e.expr for e in order_by] if order_by is not None else None
1719+
17181720
return Expr(
17191721
f.first_value(
17201722
arg.expr,
17211723
distinct=distinct,
17221724
filter=filter,
1723-
order_by=order_by,
1725+
order_by=order_by_cols,
17241726
null_treatment=null_treatment,
17251727
)
17261728
)
17271729

17281730

1729-
def last_value(arg: Expr) -> Expr:
1731+
def last_value(
1732+
arg: Expr,
1733+
distinct: bool = False,
1734+
filter: Optional[bool] = None,
1735+
order_by: Optional[list[Expr]] = None,
1736+
null_treatment: Optional[common.NullTreatment] = None,
1737+
) -> Expr:
17301738
"""Returns the last value in a group of values.
17311739
17321740
To set parameters on this expression, use ``.order_by()``, ``.distinct()``,
17331741
``.filter()``, or ``.null_treatment()``.
17341742
"""
1735-
return Expr(f.last_value(arg.expr))
1743+
order_by_cols = [e.expr for e in order_by] if order_by is not None else None
1744+
1745+
return Expr(
1746+
f.last_value(
1747+
arg.expr,
1748+
distinct=distinct,
1749+
filter=filter,
1750+
order_by=order_by_cols,
1751+
null_treatment=null_treatment,
1752+
)
1753+
)
17361754

17371755

17381756
def bit_and(arg: Expr, distinct: bool = False) -> Expr:

python/datafusion/tests/test_functions.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -963,6 +963,7 @@ def test_first_last_value(df):
963963
assert result.column(3) == pa.array(["!"])
964964
assert result.column(4) == pa.array([6])
965965
assert result.column(5) == pa.array([datetime(2020, 7, 2)])
966+
df.show()
966967

967968

968969
def test_binary_string_functions(df):

src/functions.rs

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -319,18 +319,15 @@ pub fn regr_syy(expr_y: PyExpr, expr_x: PyExpr, distinct: bool) -> PyResult<PyEx
319319
}
320320
}
321321

322-
#[pyfunction]
323-
pub fn first_value(
324-
expr: PyExpr,
322+
fn add_builder_fns_to_aggregate(
323+
agg_fn: Expr,
325324
distinct: bool,
326325
filter: Option<PyExpr>,
327326
order_by: Option<Vec<PyExpr>>,
328327
null_treatment: Option<NullTreatment>,
329328
) -> PyResult<PyExpr> {
330-
// If we initialize the UDAF with order_by directly, then it gets over-written by the builder
331-
let agg_fn = functions_aggregate::expr_fn::first_value(expr.expr, None);
332-
333-
// luckily, I can guarantee initializing a builder with an `order_by` default of empty vec
329+
// Since ExprFuncBuilder::new() is private, we can guarantee initializing
330+
// a builder with an `order_by` default of empty vec
334331
let order_by = order_by
335332
.map(|x| x.into_iter().map(|x| x.expr).collect::<Vec<_>>())
336333
.unwrap_or_default();
@@ -351,8 +348,30 @@ pub fn first_value(
351348
}
352349

353350
#[pyfunction]
354-
pub fn last_value(expr: PyExpr) -> PyExpr {
355-
functions_aggregate::expr_fn::last_value(vec![expr.expr]).into()
351+
pub fn first_value(
352+
expr: PyExpr,
353+
distinct: bool,
354+
filter: Option<PyExpr>,
355+
order_by: Option<Vec<PyExpr>>,
356+
null_treatment: Option<NullTreatment>,
357+
) -> PyResult<PyExpr> {
358+
// If we initialize the UDAF with order_by directly, then it gets over-written by the builder
359+
let agg_fn = functions_aggregate::expr_fn::first_value(expr.expr, None);
360+
361+
add_builder_fns_to_aggregate(agg_fn, distinct, filter, order_by, null_treatment)
362+
}
363+
364+
#[pyfunction]
365+
pub fn last_value(
366+
expr: PyExpr,
367+
distinct: bool,
368+
filter: Option<PyExpr>,
369+
order_by: Option<Vec<PyExpr>>,
370+
null_treatment: Option<NullTreatment>,
371+
) -> PyResult<PyExpr> {
372+
let agg_fn = functions_aggregate::expr_fn::last_value(vec![expr.expr]);
373+
374+
add_builder_fns_to_aggregate(agg_fn, distinct, filter, order_by, null_treatment)
356375
}
357376

358377
#[pyfunction]

0 commit comments

Comments
 (0)