Refactor RustAccumulator to support pyarrow array types and improve type checking for list types

Author: kosiew

python/tests/test_udaf.py

@@ -75,8 +75,8 @@ def merge(self, states: list[pa.Array]) -> None:
             if state is not None:
                 self._values.extend(state)
 
-    def evaluate(self) -> pa.Array:
-        return pa.array(self._values, type=pa.timestamp("ns"))
+    def evaluate(self) -> pa.Scalar:
+        return pa.scalar(self._values, type=pa.list_(pa.timestamp("ns")))
 
 
 @pytest.fixture

src/udaf.rs

@@ -38,11 +38,52 @@ use crate::utils::{parse_volatility, validate_pycapsule};
 struct RustAccumulator {
     accum: Py<PyAny>,
     return_type: DataType,
+    pyarrow_array_type: Option<Py<PyType>>,
+    pyarrow_chunked_array_type: Option<Py<PyType>>,
 }
 
 impl RustAccumulator {
     fn new(accum: Py<PyAny>, return_type: DataType) -> Self {
-        Self { accum, return_type }
+        Self {
+            accum,
+            return_type,
+            pyarrow_array_type: None,
+            pyarrow_chunked_array_type: None,
+        }
+    }
+
+    fn ensure_pyarrow_types(
+        &mut self,
+        py: Python<'_>,
+    ) -> PyResult<(Bound<'_, PyType>, Bound<'_, PyType>)> {
+        if self.pyarrow_array_type.is_none() || self.pyarrow_chunked_array_type.is_none() {
+            let pyarrow = PyModule::import(py, "pyarrow")?;
+            let array_attr = pyarrow.getattr("Array")?;
+            let array_type = array_attr.downcast::<PyType>()?;
+            let chunked_array_attr = pyarrow.getattr("ChunkedArray")?;
+            let chunked_array_type = chunked_array_attr.downcast::<PyType>()?;
+            self.pyarrow_array_type = Some(array_type.unbind());
+            self.pyarrow_chunked_array_type = Some(chunked_array_type.unbind());
+        }
+        Ok((
+            self.pyarrow_array_type
+                .as_ref()
+                .expect("array type set")
+                .bind(py),
+            self.pyarrow_chunked_array_type
+                .as_ref()
+                .expect("chunked array type set")
+                .bind(py),
+        ))
+    }
+
+    fn is_pyarrow_array_like(
+        &mut self,
+        py: Python<'_>,
+        value: &Bound<'_, PyAny>,
+    ) -> PyResult<bool> {
+        let (array_type, chunked_array_type) = self.ensure_pyarrow_types(py)?;
+        Ok(value.is_instance(&array_type)? || value.is_instance(&chunked_array_type)?)
     }
 }
 
@@ -65,7 +106,7 @@ impl Accumulator for RustAccumulator {
                 self.return_type,
                 DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _)
             );
-            if is_list_type && is_pyarrow_array_like(py, &value)? {
+            if is_list_type && self.is_pyarrow_array_like(py, &value)? {
                 let pyarrow = PyModule::import(py, "pyarrow")?;
                 let list_value = value.call_method0("to_pylist")?;
                 let py_type = self.return_type.to_pyarrow(py)?;
@@ -171,15 +212,6 @@ pub fn to_rust_accumulator(accum: Py<PyAny>) -> AccumulatorFactoryFunction {
     })
 }
 
-fn is_pyarrow_array_like(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult<bool> {
-    let pyarrow = PyModule::import(py, "pyarrow")?;
-    let array_attr = pyarrow.getattr("Array")?;
-    let array_type = array_attr.downcast::<PyType>()?;
-    let chunked_array_attr = pyarrow.getattr("ChunkedArray")?;
-    let chunked_array_type = chunked_array_attr.downcast::<PyType>()?;
-    Ok(value.is_instance(array_type)? || value.is_instance(chunked_array_type)?)
-}
-
 fn aggregate_udf_from_capsule(capsule: &Bound<'_, PyCapsule>) -> PyDataFusionResult<AggregateUDF> {
     validate_pycapsule(capsule, "datafusion_aggregate_udf")?;