fix sorted window expression (#4349)

polars/polars-lazy/src/physical_plan/expressions/window.rs

@@ -349,10 +349,7 @@ impl PhysicalExpr for WindowExpr {
                 // The example below shows the naive version without group tuple mapping
 
                 // columns
-                // a
-                // b
-                // a
-                // a
+                // a b a a
                 //
                 // agg list
                 // [0, 2, 3]
@@ -368,6 +365,9 @@ impl PhysicalExpr for WindowExpr {
                 //
                 // take by argsorted indexes and voila groups mapped
                 // [0, 1, 2, 3]
+
+                // TODO!
+                // investigate if sorted arrays can be return directly
                 let out_column = ac.aggregated();
                 let mut original_idx = Vec::with_capacity(out_column.len());
                 match gb.get_groups() {
@@ -378,7 +378,7 @@ impl PhysicalExpr for WindowExpr {
                     }
                     GroupsProxy::Slice { groups, .. } => {
                         for g in groups {
-                            original_idx.extend(g[0]..g[0] + 1)
+                            original_idx.extend(g[0]..g[0] + g[1])
                         }
                     }
                 };

polars/polars-utils/src/sort.rs

@@ -5,7 +5,7 @@ use rayon::{prelude::*, ThreadPool};
 /// The second argsort sorts indices from `0` to `len` so can be just assigned to the
 /// new index location.
 ///
-/// Besides that we know that all indices are unique ang thus not alias so we can parallelize.
+/// Besides that we know that all indices are unique and thus not alias so we can parallelize.
 ///
 /// This sort does not sort in place and will allocate.
 ///

py-polars/tests/test_window.py

@@ -212,3 +212,19 @@ def test_window_functions_list_types() -> None:
         "col_list": [[1], [1], [2], [2]],
         "list_shifted": [[[], [1]], [[], [1]], [[], [2]], [[], [2]]],
     }
+
+
+def test_sorted_window_expression() -> None:
+    size = 10
+    df = pl.DataFrame(
+        {"a": np.random.randint(10, size=size), "b": np.random.randint(10, size=size)}
+    )
+    expr = (pl.col("a") + pl.col("b")).over("b").alias("computed")
+
+    out1 = df.with_column(expr).sort("b")
+
+    # explicit sort
+    df = df.sort("b")
+    out2 = df.with_column(expr)
+
+    assert out1.frame_equal(out2)