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

apache · Dec 28, 2023 · 76d7fcf · 76d7fcf
1 parent 2889de0
commit 76d7fcf
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/datafusion/__init__.py b/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,

diff --git a/src/udaf.rs b/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),