Add fallback overload for AggregateUDF.udaf

Implement fallback for PyCapsule-backed providers, ensuring
type checkers are satisfied without protocol-aware stubs.
Update typing imports and cast PyCapsule inputs in
AggregateUDF.from_pycapsule for precise constructor typing.

Author: kosiew

python/datafusion/user_defined.py

@@ -22,7 +22,16 @@
 import functools
 from abc import ABCMeta, abstractmethod
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol, TypeVar, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Optional,
+    Protocol,
+    TypeVar,
+    cast,
+    overload,
+)
 
 import pyarrow as pa
 
@@ -293,11 +302,11 @@ class AggregateUDF:
     def __init__(
         self,
         name: str,
-        accumulator: Callable[[], Accumulator],
-        input_types: list[pa.DataType],
-        return_type: pa.DataType,
-        state_type: list[pa.DataType],
-        volatility: Volatility | str,
+        accumulator: Callable[[], Accumulator] | AggregateUDFExportable,
+        input_types: list[pa.DataType] | None,
+        return_type: pa.DataType | None,
+        state_type: list[pa.DataType] | None,
+        volatility: Volatility | str | None,
     ) -> None:
         """Instantiate a user-defined aggregate function (UDAF).
 
@@ -307,6 +316,18 @@ def __init__(
         if hasattr(accumulator, "__datafusion_aggregate_udf__"):
             self._udaf = df_internal.AggregateUDF.from_pycapsule(accumulator)
             return
+        if (
+            input_types is None
+            or return_type is None
+            or state_type is None
+            or volatility is None
+        ):
+            msg = (
+                "`input_types`, `return_type`, `state_type`, and `volatility` "
+                "must be provided when `accumulator` is callable."
+            )
+            raise TypeError(msg)
+
         self._udaf = df_internal.AggregateUDF(
             name,
             accumulator,
@@ -350,6 +371,14 @@ def udaf(
         name: Optional[str] = None,
     ) -> AggregateUDF: ...
 
+    @overload
+    @staticmethod
+    def udaf(accum: AggregateUDFExportable) -> AggregateUDF: ...
+
+    @overload
+    @staticmethod
+    def udaf(accum: object) -> AggregateUDF: ...
+
     @staticmethod
     def udaf(*args: Any, **kwargs: Any):  # noqa: D417, C901
         """Create a new User-Defined Aggregate Function (UDAF).
@@ -480,16 +509,17 @@ def wrapper(*args: Any, **kwargs: Any) -> Expr:
         return _decorator(*args, **kwargs)
 
     @staticmethod
-    def from_pycapsule(func: AggregateUDFExportable) -> AggregateUDF:
+    def from_pycapsule(func: AggregateUDFExportable | object) -> AggregateUDF:
         """Create an Aggregate UDF from AggregateUDF PyCapsule object.
 
         This function will instantiate a Aggregate UDF that uses a DataFusion
         AggregateUDF that is exported via the FFI bindings.
         """
-        name = str(func.__class__)
+        capsule = cast(AggregateUDFExportable, func)
+        name = str(capsule.__class__)
         return AggregateUDF(
             name=name,
-            accumulator=func,
+            accumulator=capsule,
             input_types=None,
             return_type=None,
             state_type=None,