Set up UDWF to take arguments as constructor just like UDAF to ensure we get a clean state when functions are reused

Author: timsaucer

python/datafusion/udf.py

@@ -429,18 +429,19 @@ class WindowUDF:
     def __init__(
         self,
         name: str | None,
-        func: WindowEvaluator,
+        func: Type[WindowEvaluator],
         input_types: list[pyarrow.DataType],
         return_type: pyarrow.DataType,
         volatility: Volatility | str,
+        arguments: list[Any],
     ) -> None:
         """Instantiate a user-defined window function (UDWF).
 
         See :py:func:`udwf` for a convenience function and argument
         descriptions.
         """
         self._udwf = df_internal.WindowUDF(
-            name, func, input_types, return_type, str(volatility)
+            name, func, input_types, return_type, str(volatility), arguments
         )
 
     def __call__(self, *args: Expr) -> Expr:
@@ -454,10 +455,11 @@ def __call__(self, *args: Expr) -> Expr:
 
     @staticmethod
     def udwf(
-        func: WindowEvaluator,
+        func: Type[WindowEvaluator],
         input_types: pyarrow.DataType | list[pyarrow.DataType],
         return_type: pyarrow.DataType,
         volatility: Volatility | str,
+        arguments: Optional[list[Any]] = None,
         name: str | None = None,
     ) -> WindowUDF:
         """Create a new User-Defined Window Function.
@@ -467,23 +469,26 @@ def udwf(
             input_types: The data types of the arguments to ``func``.
             return_type: The data type of the return value.
             volatility: See :py:class:`Volatility` for allowed values.
+            arguments: A list of arguments to pass in to the __init__ method for accum.
             name: A descriptive name for the function.
 
         Returns:
             A user-defined window function.
         """
-        if not isinstance(func, WindowEvaluator):
+        if not issubclass(func, WindowEvaluator):
             raise TypeError(
                 "`func` must implement the abstract base class WindowEvaluator"
             )
         if name is None:
             name = func.__class__.__qualname__.lower()
         if isinstance(input_types, pyarrow.DataType):
             input_types = [input_types]
+        arguments = [] if arguments is None else arguments
         return WindowUDF(
             name=name,
             func=func,
             input_types=input_types,
             return_type=return_type,
             volatility=volatility,
+            arguments=arguments,
         )

python/tests/test_udwf.py

@@ -24,7 +24,7 @@
 
 
 class ExponentialSmoothDefault(WindowEvaluator):
-    def __init__(self, alpha: float) -> None:
+    def __init__(self, alpha: float = 0.8) -> None:
         self.alpha = alpha
 
     def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array:
@@ -183,46 +183,58 @@ def df():
 def test_udwf_errors(df):
     with pytest.raises(TypeError):
         udwf(
-            NotSubclassOfWindowEvaluator(),
+            NotSubclassOfWindowEvaluator,
             pa.float64(),
             pa.float64(),
             volatility="immutable",
         )
 
 
 smooth_default = udwf(
-    ExponentialSmoothDefault(0.9),
+    ExponentialSmoothDefault,
+    pa.float64(),
+    pa.float64(),
+    volatility="immutable",
+    arguments=[0.9],
+)
+
+smooth_no_arugments = udwf(
+    ExponentialSmoothDefault,
     pa.float64(),
     pa.float64(),
     volatility="immutable",
 )
 
 smooth_bounded = udwf(
-    ExponentialSmoothBounded(0.9),
+    ExponentialSmoothBounded,
     pa.float64(),
     pa.float64(),
     volatility="immutable",
+    arguments=[0.9],
 )
 
 smooth_rank = udwf(
-    ExponentialSmoothRank(0.9),
+    ExponentialSmoothRank,
     pa.utf8(),
     pa.float64(),
     volatility="immutable",
+    arguments=[0.9],
 )
 
 smooth_frame = udwf(
-    ExponentialSmoothFrame(0.9),
+    ExponentialSmoothFrame,
     pa.float64(),
     pa.float64(),
     volatility="immutable",
+    arguments=[0.9],
 )
 
 smooth_two_col = udwf(
-    SmoothTwoColumn(0.9),
+    SmoothTwoColumn,
     [pa.int64(), pa.int64()],
     pa.float64(),
     volatility="immutable",
+    arguments=[0.9],
 )
 
 data_test_udwf_functions = [
@@ -231,6 +243,11 @@ def test_udwf_errors(df):
         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")),
+        [0, 0.8, 1.76, 2.752, 3.75, 4.75, 5.75],
+    ),
     (
         "default_udwf_partitioned",
         smooth_default(column("a")).partition_by(column("c")).build(),

src/udwf.rs

@@ -195,9 +195,17 @@ impl PartitionEvaluator for RustPartitionEvaluator {
     }
 }
 
-pub fn to_rust_partition_evaluator(evaluator: PyObject) -> PartitionEvaluatorFactory {
+pub fn to_rust_partition_evaluator(
+    evaluator: PyObject,
+    arguments: Vec<PyObject>,
+) -> PartitionEvaluatorFactory {
     Arc::new(move || -> Result<Box<dyn PartitionEvaluator>> {
-        let evaluator = Python::with_gil(|py| evaluator.clone_ref(py));
+        let evaluator = Python::with_gil(|py| {
+            let py_args = PyTuple::new_bound(py, arguments.iter());
+            evaluator
+                .call1(py, py_args)
+                .map_err(|e| DataFusionError::Execution(e.to_string()))
+        })?;
         Ok(Box::new(RustPartitionEvaluator::new(evaluator)))
     })
 }
@@ -212,13 +220,14 @@ pub struct PyWindowUDF {
 #[pymethods]
 impl PyWindowUDF {
     #[new]
-    #[pyo3(signature=(name, evaluator, input_types, return_type, volatility))]
+    #[pyo3(signature=(name, evaluator, input_types, return_type, volatility, arguments))]
     fn new(
         name: &str,
         evaluator: PyObject,
         input_types: Vec<PyArrowType<DataType>>,
         return_type: PyArrowType<DataType>,
         volatility: &str,
+        arguments: Vec<PyObject>,
     ) -> PyResult<Self> {
         let return_type = return_type.0;
         let input_types = input_types.into_iter().map(|t| t.0).collect();
@@ -228,7 +237,7 @@ impl PyWindowUDF {
             input_types,
             return_type,
             parse_volatility(volatility)?,
-            to_rust_partition_evaluator(evaluator),
+            to_rust_partition_evaluator(evaluator, arguments),
         ));
         Ok(Self { function })
     }