diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py
index 848ab4cee..e44f9cec4 100644
--- a/python/datafusion/user_defined.py
+++ b/python/datafusion/user_defined.py
@@ -33,7 +33,7 @@
 if TYPE_CHECKING:
     from _typeshed import CapsuleType as _PyCapsule
 
-    _R = TypeVar("_R", bound=pa.DataType)
+    _R = TypeVar("_R", bound=pa.Array)
     from collections.abc import Callable, Sequence
 
 
@@ -125,7 +125,7 @@ def __init__(
         name: str,
         func: Callable[..., _R],
         input_fields: list[pa.Field],
-        return_field: _R,
+        return_field: pa.Field,
         volatility: Volatility | str,
     ) -> None:
         """Instantiate a scalar user-defined function (UDF).
@@ -202,7 +202,8 @@ def udf(*args: Any, **kwargs: Any):  # noqa: D417
             input_fields (list[pa.Field | pa.DataType]): The data types or Fields
                 of the arguments to ``func``. This list must be of the same length
                 as the number of arguments.
-            return_field (_R): The field of the return value from the function.
+            return_field (pa.DataType | pa.Field): The field of the return value
+                from the function.
             volatility (Volatility | str): See `Volatility` for allowed values.
             name (Optional[str]): A descriptive name for the function.
 
@@ -264,7 +265,7 @@ def _function(
 
         def _decorator(
             input_fields: Sequence[pa.DataType | pa.Field] | pa.DataType | pa.Field,
-            return_field: _R,
+            return_field: pa.DataType | pa.Field,
             volatility: Volatility | str,
             name: str | None = None,
         ) -> Callable: