feat: Expose descending and nulls last in window order-by (#20919)

pola-rs · Jan 26, 2025 · 64a01e7 · 64a01e7
1 parent 2352c1e
commit 64a01e7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/py-polars/polars/expr/expr.py b/py-polars/polars/expr/expr.py
@@ -3412,6 +3412,8 @@ def over(
         partition_by: IntoExpr | Iterable[IntoExpr],
         *more_exprs: IntoExpr,
         order_by: IntoExpr | Iterable[IntoExpr] | None = None,
+        order_by_descending: bool = False,
+        order_by_nulls_last: bool = False,
         mapping_strategy: WindowMappingStrategy = "group_to_rows",
     ) -> Expr:
         """
@@ -3431,9 +3433,15 @@ def over(
             column names.
         *more_exprs
             Additional columns to group by, specified as positional arguments.
-        order_by:
+        order_by
             Order the window functions/aggregations with the partitioned groups by the
             result of the expression passed to `order_by`.
+        order_by_descending
+            In case 'order_by' is given, indicate whether to order in
+            ascending or descending order.
+        order_by_nulls_last
+            In case 'order_by' is given, indicate whether to order
+            the nulls in last position.
         mapping_strategy: {'group_to_rows', 'join', 'explode'}
             - group_to_rows
                 If the aggregation results in multiple values, assign them back to their
@@ -3571,7 +3579,7 @@ def over(
             self._pyexpr.over(
                 partition_by,
                 order_by=order_by,
-                order_by_descending=False,  # does not work yet
+                order_by_descending=order_by_descending,
                 order_by_nulls_last=False,  # does not work yet
                 mapping_strategy=mapping_strategy,
             )

diff --git a/py-polars/tests/unit/operations/test_window.py b/py-polars/tests/unit/operations/test_window.py
@@ -487,12 +487,14 @@ def test_window_order_by_8662() -> None:
     assert df.with_columns(
         x_lag0=pl.col("x").shift(1).over("g"),
         x_lag1=pl.col("x").shift(1).over("g", order_by="t"),
+        x_lag2=pl.col("x").shift(1).over("g", order_by="t", order_by_descending=True),
     ).to_dict(as_series=False) == {
         "g": [1, 1, 1, 1, 2, 2, 2, 2],
         "t": [1, 2, 3, 4, 4, 1, 2, 3],
         "x": [10, 20, 30, 40, 10, 20, 30, 40],
         "x_lag0": [None, 10, 20, 30, None, 10, 20, 30],
         "x_lag1": [None, 10, 20, 30, 40, None, 20, 30],
+        "x_lag2": [20, 30, 40, None, None, 30, 40, 10],
     }