feat: udaf: enable multiple column input (#546)

Author: dlovell

datafusion/__init__.py

@@ -213,6 +213,8 @@ def udaf(accum, input_type, return_type, state_type, volatility, name=None):
         )
     if name is None:
         name = accum.__qualname__.lower()
+    if isinstance(input_type, pa.lib.DataType):
+        input_type = [input_type]
     return AggregateUDF(
         name=name,
         accumulator=accum,

src/udaf.rs

@@ -148,14 +148,14 @@ impl PyAggregateUDF {
     fn new(
         name: &str,
         accumulator: PyObject,
-        input_type: PyArrowType<DataType>,
+        input_type: PyArrowType<Vec<DataType>>,
         return_type: PyArrowType<DataType>,
         state_type: PyArrowType<Vec<DataType>>,
         volatility: &str,
     ) -> PyResult<Self> {
         let function = create_udaf(
             name,
-            vec![input_type.0],
+            input_type.0,
             Arc::new(return_type.0),
             parse_volatility(volatility)?,
             to_rust_accumulator(accumulator),