fix: reenable num_centroids argument for approx_percentile_cont

python/datafusion/functions.py

@@ -1473,19 +1473,18 @@ def approx_median(arg: Expr, distinct: bool = False) -> Expr:
 def approx_percentile_cont(
     expression: Expr,
     percentile: Expr,
+    num_centroids: Expr | None = None,
     distinct: bool = False,
 ) -> Expr:
     """Returns the value that is approximately at a given percentile of ``expr``."""
-    # Re-enable num_centroids: https://github.com/apache/datafusion-python/issues/777
-    num_centroids = None
     if num_centroids is None:
         return Expr(
-            f.approx_percentile_cont(expression.expr, percentile.expr, distinct=distinct)
+            f.approx_percentile_cont(expression.expr, percentile.expr, distinct=distinct, num_centroids=None)
         )
 
     return Expr(
         f.approx_percentile_cont(
-            expression.expr, percentile.expr, distinct=distinct
+            expression.expr, percentile.expr, distinct=distinct, num_centroids=num_centroids.expr
         )
     )
 

src/functions.rs

@@ -57,9 +57,15 @@ pub fn approx_percentile_cont(
     expression: PyExpr,
     percentile: PyExpr,
     distinct: bool,
+    num_centroids: Option<PyExpr>, // enforces optional arguments at the end, currently
 ) -> PyResult<PyExpr> {
-    let expr =
-        functions_aggregate::expr_fn::approx_percentile_cont(expression.expr, percentile.expr);
+    let args = if let Some(num_centroids) = num_centroids {
+        vec![expression.expr, percentile.expr, num_centroids.expr]
+    } else {
+        vec![expression.expr, percentile.expr]
+    };
+    let udaf = functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf();
+    let expr = udaf.call(args);
     if distinct {
         Ok(expr.distinct().build()?.into())
     } else {