Update corr python wrapper to expose only builder parameters used

Author: timsaucer

python/datafusion/functions.py

@@ -1643,9 +1643,21 @@ def avg(
     return Expr(f.avg(expression.expr, filter=filter_raw))
 
 
-def corr(value1: Expr, value2: Expr, distinct: bool = False) -> Expr:
-    """Returns the correlation coefficient between ``value1`` and ``value2``."""
-    return Expr(f.corr(value1.expr, value2.expr, distinct=distinct))
+def corr(value_y: Expr, value_x: Expr, filter: Optional[Expr] = None) -> Expr:
+    """Returns the correlation coefficient between ``value1`` and ``value2``.
+
+    This aggregate function expects both values to be numeric and will return a float.
+
+    If using the builder functions described in ref:`_aggregation` this function ignores
+    the options ``order_by``, ``null_treatment``, and ``distinct``.
+
+    Args:
+        value_y: The dependent variable for correlation
+        value_x: The independent variable for correlation
+        filter: If provided, only compute against rows for which the filter is true
+    """
+    filter_raw = filter.expr if filter is not None else None
+    return Expr(f.corr(value_y.expr, value_x.expr, filter=filter_raw))
 
 
 def count(args: Expr | list[Expr] | None = None, distinct: bool = False) -> Expr:

python/datafusion/tests/test_aggregation.py

@@ -169,6 +169,16 @@ def test_aggregation(df, agg_expr, expected, array_sort):
             ),
             [83, 68, 122, 124, 117],
         ),
+        (
+            "corr",
+            f.corr(column("c3"), column("c2")),
+            [-0.1056, -0.2808, 0.0023, 0.0022, -0.2473],
+        ),
+        (
+            "corr_w_filter",
+            f.corr(column("c3"), column("c2"), filter=column("c3") > lit(0)),
+            [-0.3298, 0.2925, 0.2467, -0.2269, 0.0358],
+        ),
     ],
 )
 def test_aggregate_100(df_aggregate_100, name, expr, expected):