diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py index d2513995..27f8332e 100644 --- a/python/datafusion/functions.py +++ b/python/datafusion/functions.py @@ -372,7 +372,7 @@ def col(name: str) -> Expr: def count_star() -> Expr: """Create a COUNT(1) aggregate expression.""" - return Expr(f.count_star()) + return Expr(f.count(Expr.literal(1))) def case(expr: Expr) -> CaseBuilder: diff --git a/src/functions.rs b/src/functions.rs index 78712a8b..4ffe2557 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -432,25 +432,6 @@ fn col(name: &str) -> PyResult { }) } -// TODO: should we just expose this in python? -/// Create a COUNT(1) aggregate expression -#[pyfunction] -fn count_star() -> PyExpr { - functions_aggregate::expr_fn::count(lit(1)).into() -} - -/// Wrapper for [`functions_aggregate::expr_fn::count`] -/// Count the number of non-null values in the column -#[pyfunction] -fn count(expr: PyExpr, distinct: bool) -> PyResult { - let expr = functions_aggregate::expr_fn::count(expr.expr); - if distinct { - Ok(expr.distinct().build()?.into()) - } else { - Ok(expr.into()) - } -} - /// Create a CASE WHEN statement with literal WHEN expressions for comparison to the base expression. #[pyfunction] fn case(expr: PyExpr) -> PyResult { @@ -830,6 +811,7 @@ aggregate_function!(bit_xor); aggregate_function!(bool_and); aggregate_function!(bool_or); aggregate_function!(corr, y x); +aggregate_function!(count); fn add_builder_fns_to_window( window_fn: Expr, @@ -978,7 +960,6 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(cosh))?; m.add_wrapped(wrap_pyfunction!(cot))?; m.add_wrapped(wrap_pyfunction!(count))?; - m.add_wrapped(wrap_pyfunction!(count_star))?; m.add_wrapped(wrap_pyfunction!(covar_pop))?; m.add_wrapped(wrap_pyfunction!(covar_samp))?; m.add_wrapped(wrap_pyfunction!(current_date))?;