fix window function braodcasting

pola-rs · Sep 10, 2024 · e1735d9 · e1735d9
1 parent 96c8816
commit e1735d9
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/crates/polars-expr/src/expressions/window.rs b/crates/polars-expr/src/expressions/window.rs
@@ -519,6 +519,10 @@ impl PhysicalExpr for WindowExpr {
         match self.determine_map_strategy(ac.agg_state(), sorted_keys, &gb)? {
             Nothing => {
                 let mut out = ac.flat_naive().into_owned();
+
+                if ac.is_literal() {
+                    out = out.new_from_index(0, df.height())
+                }
                 cache_gb(gb, state, &cache_key);
                 if let Some(name) = &self.out_name {
                     out.rename(name.clone());

diff --git a/py-polars/tests/unit/operations/test_window.py b/py-polars/tests/unit/operations/test_window.py
@@ -511,3 +511,10 @@ def test_window_17308() -> None:
     assert df.select(pl.col("A").sum(), pl.col("B").sum().over("grp")).to_dict(
         as_series=False
     ) == {"A": [3, 3], "B": [3, 4]}
+
+
+def test_lit_window_broadcast() -> None:
+    # the broadcast should happen in the window function
+    assert pl.DataFrame({"a": [1, 1, 2]}).select(pl.lit(0).over("a").alias("a"))[
+        "a"
+    ].to_list() == [0, 0, 0]