Add option for passing in constructor arguments to the udaf

Author: timsaucer

python/datafusion/udf.py

@@ -23,7 +23,7 @@
 from datafusion.expr import Expr
 from typing import Callable, TYPE_CHECKING, TypeVar
 from abc import ABCMeta, abstractmethod
-from typing import List
+from typing import List, Any, Optional
 from enum import Enum
 import pyarrow
 
@@ -186,14 +186,21 @@ def __init__(
         return_type: _R,
         state_type: list[pyarrow.DataType],
         volatility: Volatility | str,
+        arguments: list[Any],
     ) -> None:
         """Instantiate a user-defined aggregate function (UDAF).
 
         See :py:func:`udaf` for a convenience function and argument
         descriptions.
         """
         self._udaf = df_internal.AggregateUDF(
-            name, accumulator, input_types, return_type, state_type, str(volatility)
+            name,
+            accumulator,
+            input_types,
+            return_type,
+            state_type,
+            str(volatility),
+            arguments,
         )
 
     def __call__(self, *args: Expr) -> Expr:
@@ -212,6 +219,7 @@ def udaf(
         return_type: _R,
         state_type: list[pyarrow.DataType],
         volatility: Volatility | str,
+        arguments: Optional[list[Any]] = None,
         name: str | None = None,
     ) -> AggregateUDF:
         """Create a new User-Defined Aggregate Function.
@@ -224,6 +232,7 @@ def udaf(
             return_type: The data type of the return value.
             state_type: The data types of the intermediate accumulation.
             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:
@@ -238,13 +247,15 @@ def udaf(
             name = accum.__qualname__.lower()
         if isinstance(input_types, pyarrow.lib.DataType):
             input_types = [input_types]
+        arguments = [] if arguments is None else arguments
         return AggregateUDF(
             name=name,
             accumulator=accum,
             input_types=input_types,
             return_type=return_type,
             state_type=state_type,
             volatility=volatility,
+            arguments=arguments,
         )
 
 

python/tests/test_udaf.py

@@ -27,8 +27,8 @@
 class Summarize(Accumulator):
     """Interface of a user-defined accumulation."""
 
-    def __init__(self):
-        self._sum = pa.scalar(0.0)
+    def __init__(self, initial_value: float = 0.0):
+        self._sum = pa.scalar(initial_value)
 
     def state(self) -> List[pa.Scalar]:
         return [self._sum]
@@ -97,7 +97,7 @@ def test_errors(df):
         df.collect()
 
 
-def test_aggregate(df):
+def test_udaf_aggregate(df):
     summarize = udaf(
         Summarize,
         pa.float64(),
@@ -106,14 +106,41 @@ def test_aggregate(df):
         volatility="immutable",
     )
 
-    df = df.aggregate([], [summarize(column("a"))])
+    df1 = df.aggregate([], [summarize(column("a"))])
 
     # execute and collect the first (and only) batch
-    result = df.collect()[0]
+    result = df1.collect()[0]
+
+    assert result.column(0) == pa.array([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([1.0 + 2.0 + 3.0])
 
 
+def test_udaf_aggregate_with_arguments(df):
+    bias = 10.0
+
+    summarize = udaf(
+        Summarize,
+        pa.float64(),
+        pa.float64(),
+        [pa.float64()],
+        volatility="immutable",
+        arguments=[bias],
+    )
+
+    df1 = df.aggregate([], [summarize(column("a"))])
+
+    # execute and collect the first (and only) batch
+    result = df1.collect()[0]
+
+    assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0])
+
+
 def test_group_by(df):
     summarize = udaf(
         Summarize,

src/udaf.rs

@@ -128,11 +128,15 @@ impl Accumulator for RustAccumulator {
     }
 }
 
-pub fn to_rust_accumulator(accum: PyObject) -> AccumulatorFactoryFunction {
+pub fn to_rust_accumulator(
+    accum: PyObject,
+    arguments: Vec<PyObject>,
+) -> AccumulatorFactoryFunction {
     Arc::new(move |_| -> Result<Box<dyn Accumulator>> {
         let accum = Python::with_gil(|py| {
+            let py_args = PyTuple::new_bound(py, arguments.iter());
             accum
-                .call0(py)
+                .call1(py, py_args)
                 .map_err(|e| DataFusionError::Execution(format!("{e}")))
         })?;
         Ok(Box::new(RustAccumulator::new(accum)))
@@ -149,21 +153,22 @@ pub struct PyAggregateUDF {
 #[pymethods]
 impl PyAggregateUDF {
     #[new]
-    #[pyo3(signature=(name, accumulator, input_type, return_type, state_type, volatility))]
+    #[pyo3(signature=(name, accumulator, input_type, return_type, state_type, volatility, arguments))]
     fn new(
         name: &str,
         accumulator: PyObject,
         input_type: PyArrowType<Vec<DataType>>,
         return_type: PyArrowType<DataType>,
         state_type: PyArrowType<Vec<DataType>>,
         volatility: &str,
+        arguments: Vec<PyObject>,
     ) -> PyResult<Self> {
         let function = create_udaf(
             name,
             input_type.0,
             Arc::new(return_type.0),
             parse_volatility(volatility)?,
-            to_rust_accumulator(accumulator),
+            to_rust_accumulator(accumulator, arguments),
             Arc::new(state_type.0),
         );
         Ok(Self { function })