refactor: Update udwf method signature and simplify input handling

- Changed the type hint for the return type in the _create_window_udf_decorator method to use pa.DataType directly instead of a TypeVar.
- Simplified the handling of input types by removing redundant checks and directly using the input types list.
- Removed unnecessary comments and cleaned up the code for better readability.
- Updated the test for udwf to use parameterized tests for better coverage and maintainability.

Author: kosiew

python/datafusion/udf.py

@@ -22,20 +22,15 @@
 import functools
 from abc import ABCMeta, abstractmethod
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Callable, Optional, overload
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, overload
 
 import pyarrow as pa
 
 import datafusion._internal as df_internal
 from datafusion.expr import Expr
 
 if TYPE_CHECKING:
-    # for python 3.10 and above, we can use
-    # from typing import TypeAlias
-    # but for python 3.9, we use the following
-    from typing_extensions import TypeAlias
-
-    _R: TypeAlias = pa.DataType
+    _R = TypeVar("_R", bound=pa.DataType)
 
 
 class Volatility(Enum):
@@ -719,19 +714,13 @@ def _create_window_udf(
             msg = "`func` must implement the abstract base class WindowEvaluator"
             raise TypeError(msg)
 
-        if name is None:
-            name = WindowUDF._get_default_name(func)
-
-        input_types_list = WindowUDF._normalize_input_types(input_types)
-
-        return WindowUDF(
-            name=name,
-            func=func,
-            input_types=input_types_list,
-            return_type=return_type,
-            volatility=volatility,
+        name = name or func.__qualname__.lower()
+        input_types = (
+            [input_types] if isinstance(input_types, pa.DataType) else input_types
         )
 
+        return WindowUDF(name, func, input_types, return_type, volatility)
+
     @staticmethod
     def _get_default_name(func: Callable) -> str:
         """Get the default name for a function based on its attributes."""
@@ -751,10 +740,10 @@ def _normalize_input_types(
     @staticmethod
     def _create_window_udf_decorator(
         input_types: pa.DataType | list[pa.DataType],
-        return_type: _R,
+        return_type: pa.DataType,
         volatility: Volatility | str,
         name: Optional[str] = None,
-    ) -> Callable[..., Callable[..., Expr]]:
+    ) -> Callable[[Callable[[], WindowEvaluator]], Callable[..., Expr]]:
         """Create a decorator for a WindowUDF."""
 
         def decorator(func: Callable[[], WindowEvaluator]) -> Callable[..., Expr]:

python/tests/test_udwf.py

@@ -428,55 +428,24 @@ def test_udwf_functions(complex_window_df, name, expr, expected):
     assert result.column(0) == pa.array(expected)
 
 
-def test_udwf_overloads(count_window_df):
-    """Test different overload patterns for UDWF function."""
-    # Single input type syntax
-    single_input = udwf(
-        SimpleWindowCount, pa.int64(), pa.int64(), volatility="immutable"
-    )
-
-    # List of input types syntax
-    list_input = udwf(
-        SimpleWindowCount, [pa.int64()], pa.int64(), volatility="immutable"
-    )
-
-    # Decorator syntax with single input type
-    @udwf(pa.int64(), pa.int64(), "immutable")
-    def window_count_single() -> WindowEvaluator:
-        return SimpleWindowCount()
-
-    # Decorator syntax with list of input types
-    @udwf([pa.int64()], pa.int64(), "immutable")
-    def window_count_list() -> WindowEvaluator:
-        return SimpleWindowCount()
-
-    # Test all variants produce the same result
+@pytest.mark.parametrize(
+    "udwf_func",
+    [
+        udwf(SimpleWindowCount, pa.int64(), pa.int64(), "immutable"),
+        udwf(SimpleWindowCount, [pa.int64()], pa.int64(), "immutable"),
+        udwf([pa.int64()], pa.int64(), "immutable")(lambda: SimpleWindowCount()),
+        udwf(pa.int64(), pa.int64(), "immutable")(lambda: SimpleWindowCount()),
+    ],
+)
+def test_udwf_overloads(udwf_func, count_window_df):
     df = count_window_df.select(
-        single_input(column("a"))
+        udwf_func(column("a"))
         .window_frame(WindowFrame("rows", None, None))
         .build()
-        .alias("single"),
-        list_input(column("a"))
-        .window_frame(WindowFrame("rows", None, None))
-        .build()
-        .alias("list"),
-        window_count_single(column("a"))
-        .window_frame(WindowFrame("rows", None, None))
-        .build()
-        .alias("decorator_single"),
-        window_count_list(column("a"))
-        .window_frame(WindowFrame("rows", None, None))
-        .build()
-        .alias("decorator_list"),
+        .alias("count")
     )
-
     result = df.collect()[0]
-    expected = pa.array([0, 1, 2])
-
-    assert result.column(0) == expected
-    assert result.column(1) == expected
-    assert result.column(2) == expected
-    assert result.column(3) == expected
+    assert result.column(0) == pa.array([0, 1, 2])
 
 
 def test_udwf_named_function(ctx, count_window_df):