refactor: update order_by handling in Window class for improved type support

Author: kosiew

python/datafusion/expr.py

@@ -696,7 +696,7 @@ def over(self, window: Window) -> Expr:
             window: Window definition
         """
         partition_by_raw = expr_list_to_raw_expr_list(window._partition_by)
-        order_by_raw = sort_list_to_raw_sort_list(window._order_by)
+        order_by_raw = window._order_by
         window_frame_raw = (
             window._window_frame.window_frame
             if window._window_frame is not None
@@ -1182,7 +1182,7 @@ def __init__(
         self,
         partition_by: Optional[list[Expr] | Expr] = None,
         window_frame: Optional[WindowFrame] = None,
-        order_by: Optional[list[SortExpr | Expr] | Expr | SortExpr] = None,
+        order_by: Optional[list[SortExpr | Expr | str] | Expr | SortExpr | str] = None,
         null_treatment: Optional[NullTreatment] = None,
     ) -> None:
         """Construct a window definition.
@@ -1195,7 +1195,7 @@ def __init__(
         """
         self._partition_by = partition_by
         self._window_frame = window_frame
-        self._order_by = order_by
+        self._order_by = sort_list_to_raw_sort_list(order_by)
         self._null_treatment = null_treatment
 
 

python/tests/test_dataframe.py

@@ -33,7 +33,6 @@
     WindowFrame,
     column,
     literal,
-    col,
 )
 from datafusion import (
     functions as f,
@@ -292,6 +291,8 @@ def test_sort_unsupported(df):
 
 
 def test_aggregate_string_and_expression_equivalent(df):
+    from datafusion import col
+
     result_str = df.aggregate("a", [f.count()]).sort("a").to_pydict()
     result_expr = df.aggregate(col("a"), [f.count()]).sort("a").to_pydict()
     assert result_str == result_expr
@@ -778,6 +779,13 @@ def test_distinct():
         ),
         [1, 1, 1, 1, 5, 5, 5],
     ),
+    (
+        "first_value_order_by_string",
+        f.first_value(column("a")).over(
+            Window(partition_by=[column("c")], order_by="b")
+        ),
+        [1, 1, 1, 1, 5, 5, 5],
+    ),
     (
         "last_value",
         f.last_value(column("a")).over(