Switching to use factory methods for udaf and udwf

Author: timsaucer

python/datafusion/tests/test_udaf.py

@@ -79,22 +79,19 @@ def test_errors(df):
             volatility="immutable",
         )
 
-    accum = udaf(
-        MissingMethods,
-        pa.int64(),
-        pa.int64(),
-        [pa.int64()],
-        volatility="immutable",
-    )
-    df = df.aggregate([], [accum(column("a"))])
-
     msg = (
         "Can't instantiate abstract class MissingMethods (without an implementation "
         "for abstract methods 'evaluate', 'merge', 'update'|with abstract methods "
         "evaluate, merge, update)"
     )
     with pytest.raises(Exception, match=msg):
-        df.collect()
+        accum = udaf(  # noqa F841
+            MissingMethods,
+            pa.int64(),
+            pa.int64(),
+            [pa.int64()],
+            volatility="immutable",
+        )
 
 
 def test_udaf_aggregate(df):
@@ -125,12 +122,11 @@ def test_udaf_aggregate_with_arguments(df):
     bias = 10.0
 
     summarize = udaf(
-        Summarize,
+        lambda: Summarize(bias),
         pa.float64(),
         pa.float64(),
         [pa.float64()],
         volatility="immutable",
-        arguments=[bias],
     )
 
     df1 = df.aggregate([], [summarize(column("a"))])
@@ -140,6 +136,13 @@ def test_udaf_aggregate_with_arguments(df):
 
     assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0])
 
+    df2 = df.aggregate([], [summarize(column("a"))])
+
+    # Run a second time to ensure the state is properly reset
+    result = df2.collect()[0]
+
+    assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0])
+
 
 def test_group_by(df):
     summarize = udaf(

python/datafusion/tests/test_udwf.py

@@ -24,7 +24,7 @@
 
 
 class ExponentialSmoothDefault(WindowEvaluator):
-    def __init__(self, alpha: float = 0.8) -> None:
+    def __init__(self, alpha: float = 0.9) -> None:
         self.alpha = alpha
 
     def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array:
@@ -44,7 +44,7 @@ def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array:
 
 
 class ExponentialSmoothBounded(WindowEvaluator):
-    def __init__(self, alpha: float) -> None:
+    def __init__(self, alpha: float = 0.9) -> None:
         self.alpha = alpha
 
     def supports_bounded_execution(self) -> bool:
@@ -75,7 +75,7 @@ def evaluate(
 
 
 class ExponentialSmoothRank(WindowEvaluator):
-    def __init__(self, alpha: float) -> None:
+    def __init__(self, alpha: float = 0.9) -> None:
         self.alpha = alpha
 
     def include_rank(self) -> bool:
@@ -101,7 +101,7 @@ def evaluate_all_with_rank(
 
 
 class ExponentialSmoothFrame(WindowEvaluator):
-    def __init__(self, alpha: float) -> None:
+    def __init__(self, alpha: float = 0.9) -> None:
         self.alpha = alpha
 
     def uses_window_frame(self) -> bool:
@@ -134,7 +134,7 @@ class SmoothTwoColumn(WindowEvaluator):
     the previous and next rows.
     """
 
-    def __init__(self, alpha: float) -> None:
+    def __init__(self, alpha: float = 0.9) -> None:
         self.alpha = alpha
 
     def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array:
@@ -195,11 +195,10 @@ def test_udwf_errors(df):
     pa.float64(),
     pa.float64(),
     volatility="immutable",
-    arguments=[0.9],
 )
 
-smooth_no_arugments = udwf(
-    ExponentialSmoothDefault,
+smooth_w_arguments = udwf(
+    lambda: ExponentialSmoothDefault(0.8),
     pa.float64(),
     pa.float64(),
     volatility="immutable",
@@ -210,42 +209,38 @@ def test_udwf_errors(df):
     pa.float64(),
     pa.float64(),
     volatility="immutable",
-    arguments=[0.9],
 )
 
 smooth_rank = udwf(
     ExponentialSmoothRank,
     pa.utf8(),
     pa.float64(),
     volatility="immutable",
-    arguments=[0.9],
 )
 
 smooth_frame = udwf(
     ExponentialSmoothFrame,
     pa.float64(),
     pa.float64(),
     volatility="immutable",
-    arguments=[0.9],
 )
 
 smooth_two_col = udwf(
     SmoothTwoColumn,
     [pa.int64(), pa.int64()],
     pa.float64(),
     volatility="immutable",
-    arguments=[0.9],
 )
 
 data_test_udwf_functions = [
     (
-        "default_udwf",
+        "default_udwf_no_arguments",
         smooth_default(column("a")),
         [0, 0.9, 1.89, 2.889, 3.889, 4.889, 5.889],
     ),
     (
-        "default_udwf_no_arguments",
-        smooth_no_arugments(column("a")),
+        "default_udwf_w_arguments",
+        smooth_w_arguments(column("a")),
         [0, 0.8, 1.76, 2.752, 3.75, 4.75, 5.75],
     ),
     (