Change udwf() to take an instance rather than a class so we can parameterize it

Author: timsaucer

python/datafusion/tests/test_udwf.py

@@ -25,8 +25,8 @@
 class ExponentialSmooth(WindowEvaluator):
     """Interface of a user-defined accumulation."""
 
-    def __init__(self) -> None:
-        self.alpha = 0.9
+    def __init__(self, alpha: float) -> None:
+        self.alpha = alpha
 
     def evaluate_all(self, values: pa.Array, num_rows: int) -> pa.Array:
         results = []
@@ -66,15 +66,15 @@ def df():
 def test_udwf_errors(df):
     with pytest.raises(TypeError):
         udwf(
-            NotSubclassOfWindowEvaluator,
+            NotSubclassOfWindowEvaluator(),
             pa.float64(),
             pa.float64(),
             volatility="immutable",
         )
 
 
 smooth = udwf(
-    ExponentialSmooth,
+    ExponentialSmooth(0.9),
     pa.float64(),
     pa.float64(),
     volatility="immutable",

python/datafusion/udf.py

@@ -21,7 +21,7 @@
 
 import datafusion._internal as df_internal
 from datafusion.expr import Expr
-from typing import Callable, TYPE_CHECKING, TypeVar, Type
+from typing import Callable, TYPE_CHECKING, TypeVar
 from abc import ABCMeta, abstractmethod
 from typing import List
 from enum import Enum
@@ -412,7 +412,7 @@ class WindowUDF:
     def __init__(
         self,
         name: str | None,
-        func: Type[WindowEvaluator],
+        func: WindowEvaluator,
         input_type: pyarrow.DataType,
         return_type: pyarrow.DataType,
         volatility: Volatility | str,
@@ -437,7 +437,7 @@ def __call__(self, *args: Expr) -> Expr:
 
     @staticmethod
     def udwf(
-        func: Type[WindowEvaluator],
+        func: WindowEvaluator,
         input_type: pyarrow.DataType,
         return_type: pyarrow.DataType,
         volatility: Volatility | str,
@@ -455,12 +455,12 @@ def udwf(
         Returns:
             A user defined window function.
         """
-        if not issubclass(func, WindowEvaluator):
+        if not isinstance(func, WindowEvaluator):
             raise TypeError(
                 "`func` must implement the abstract base class WindowEvaluator"
             )
         if name is None:
-            name = func.__qualname__.lower()
+            name = func.__class__.__qualname__.lower()
         return WindowUDF(
             name=name,
             func=func,

src/udwf.rs

@@ -186,14 +186,10 @@ impl PartitionEvaluator for RustPartitionEvaluator {
     }
 }
 
-pub fn to_rust_partition_evaluator(evalutor: PyObject) -> PartitionEvaluatorFactory {
+pub fn to_rust_partition_evaluator(evaluator: PyObject) -> PartitionEvaluatorFactory {
     Arc::new(move || -> Result<Box<dyn PartitionEvaluator>> {
-        let evalutor = Python::with_gil(|py| {
-            evalutor
-                .call0(py)
-                .map_err(|e| DataFusionError::Execution(format!("{e}")))
-        })?;
-        Ok(Box::new(RustPartitionEvaluator::new(evalutor)))
+        let evaluator = Python::with_gil(|py| evaluator.clone_ref(py));
+        Ok(Box::new(RustPartitionEvaluator::new(evaluator)))
     })
 }