From 6509d6ddd0008b561b2244f0e7a82d1821071146 Mon Sep 17 00:00:00 2001
From: Michael-J-Ward <ward.michael.j@gmail.com>
Date: Wed, 24 Jul 2024 13:05:00 -0500
Subject: [PATCH] migrate avg to UDAF

Ref: https://github.com/apache/datafusion/pull/10964
---
 src/functions.rs | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/functions.rs b/src/functions.rs
index 7bab918b..345045b8 100644
--- a/src/functions.rs
+++ b/src/functions.rs
@@ -87,6 +87,16 @@ pub fn approx_percentile_cont_with_weight(
     }
 }
 
+#[pyfunction]
+pub fn avg(expression: PyExpr, distinct: bool) -> PyResult<PyExpr> {
+    let expr = functions_aggregate::expr_fn::avg(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()
@@ -749,7 +759,6 @@ array_fn!(flatten, array);
 array_fn!(range, start stop step);
 
 aggregate_function!(array_agg, ArrayAgg);
-aggregate_function!(avg, Avg);
 aggregate_function!(corr, Correlation);
 aggregate_function!(grouping, Grouping);
 aggregate_function!(max, Max);