Broaden AggregateUDF typing for PyCapsule support

Enhance AggregateUDF to accept PyCapsule providers.
Validate that callable accumulators provide type metadata
prior to UDAF construction. Add overload for udaf for
static type checkers to recognize PyCapsule-backed
aggregate functions.

Author: kosiew

python/datafusion/user_defined.py

@@ -293,11 +293,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 +307,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 +362,10 @@ def udaf(
         name: Optional[str] = None,
     ) -> AggregateUDF: ...
 
+    @overload
+    @staticmethod
+    def udaf(accum: AggregateUDFExportable) -> AggregateUDF: ...
+
     @staticmethod
     def udaf(*args: Any, **kwargs: Any):  # noqa: D417, C901
         """Create a new User-Defined Aggregate Function (UDAF).