feat: Add overloads for udwf function to support multiple input types and decorator syntax

Author: kosiew

python/datafusion/udf.py

@@ -626,6 +626,27 @@ def __call__(self, *args: Expr) -> Expr:
         args_raw = [arg.expr for arg in args]
         return Expr(self._udwf.__call__(*args_raw))
 
+    @overload
+    @staticmethod
+    def udwf(
+        input_type: pa.DataType | list[pa.DataType],
+        return_type: pa.DataType,
+        state_type: list[pa.DataType],
+        volatility: str,
+        name: Optional[str] = None,
+    ) -> Callable[..., WindowUDF]: ...
+
+    @overload
+    @staticmethod
+    def udwf(
+        windown: Callable[[], WindowEvaluator],
+        input_type: pa.DataType | list[pa.DataType],
+        return_type: pa.DataType,
+        state_type: list[pa.DataType],
+        volatility: str,
+        name: Optional[str] = None,
+    ) -> WindowUDF: ...
+
     @staticmethod
     def udwf(
         *args: Any, **kwargs: Any

python/tests/test_udwf.py

@@ -422,3 +422,71 @@ def test_udwf_functions(complex_window_df, name, expr, expected):
     result = df.sort(column("a")).select(column(name)).collect()[0]
 
     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
+    df = count_window_df.select(
+        single_input(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"),
+    )
+
+    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
+
+
+def test_udwf_named_function(ctx, count_window_df):
+    """Test UDWF with explicit name parameter."""
+    window_count = udwf(
+        SimpleWindowCount,
+        pa.int64(),
+        pa.int64(),
+        volatility="immutable",
+        name="my_custom_counter",
+    )
+
+    ctx.register_udwf(window_count)
+    result = ctx.sql(
+        "SELECT my_custom_counter(a) OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM test_table"
+    ).collect()[0]
+    assert result.column(0) == pa.array([0, 1, 2])