Skip to content

Commit

Permalink
migrate grouping to UDAF
Browse files Browse the repository at this point in the history
  • Loading branch information
Michael-J-Ward committed Jul 25, 2024
1 parent 98498e9 commit f1717a2
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
2 changes: 1 addition & 1 deletion python/datafusion/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -1326,7 +1326,7 @@ def grouping(arg: Expr, distinct: bool = False) -> Expr:
Returns 1 if the value of the argument is aggregated, 0 if not.
"""
return Expr(f.grouping([arg.expr], distinct=distinct))
return Expr(f.grouping(arg.expr, distinct=distinct))


def max(arg: Expr, distinct: bool = False) -> Expr:
Expand Down
11 changes: 10 additions & 1 deletion src/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,16 @@ pub fn corr(y: PyExpr, x: PyExpr, distinct: bool) -> PyResult<PyExpr> {
}
}

#[pyfunction]
pub fn grouping(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
let expr = functions_aggregate::expr_fn::grouping(expression.expr);
if distinct {
Ok(expr.distinct().build()?.into())
} else {
Ok(expr.into())
}
}

#[pyfunction]
pub fn sum(args: PyExpr) -> PyExpr {
functions_aggregate::expr_fn::sum(args.expr).into()
Expand Down Expand Up @@ -799,7 +809,6 @@ array_fn!(flatten, array);
array_fn!(range, start stop step);

aggregate_function!(array_agg, ArrayAgg);
aggregate_function!(grouping, Grouping);
aggregate_function!(max, Max);
aggregate_function!(mean, Avg);
aggregate_function!(min, Min);
Expand Down

0 comments on commit f1717a2

Please sign in to comment.