fix: Fix panic in window case (#17320)

pola-rs · Jul 1, 2024 · ccdfbda · ccdfbda
1 parent 41610e3
commit ccdfbda
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 28 deletions.
diff --git a/crates/polars-mem-engine/src/executors/projection_utils.rs b/crates/polars-mem-engine/src/executors/projection_utils.rs
@@ -117,36 +117,41 @@ fn execute_projection_cached_window_fns(
     // first we partition the window function by the values they group over.
     // the group_by values should be cached
     exprs.iter().enumerate_u32().for_each(|(index, phys)| {
-        let e = phys.as_expression().unwrap();
-
         let mut is_window = false;
-        for e in e.into_iter() {
-            if let Expr::Window {
-                partition_by,
-                options,
-                order_by,
-                ..
-            } = e
-            {
-                let entry = match options {
-                    WindowType::Over(_) => {
-                        let mut key = format!("{:?}", partition_by.as_slice());
-                        if let Some((e, k)) = order_by {
-                            polars_expr::prelude::window_function_format_order_by(
-                                &mut key,
-                                e.as_ref(),
-                                k,
-                            )
-                        }
-                        windows.entry(key).or_insert_with(Vec::new)
-                    },
-                    #[cfg(feature = "dynamic_group_by")]
-                    WindowType::Rolling(options) => rolling.entry(options).or_insert_with(Vec::new),
-                };
-                entry.push((index, phys.clone()));
-                is_window = true;
-                break;
+        if let Some(e) = phys.as_expression() {
+            for e in e.into_iter() {
+                if let Expr::Window {
+                    partition_by,
+                    options,
+                    order_by,
+                    ..
+                } = e
+                {
+                    let entry = match options {
+                        WindowType::Over(_) => {
+                            let mut key = format!("{:?}", partition_by.as_slice());
+                            if let Some((e, k)) = order_by {
+                                polars_expr::prelude::window_function_format_order_by(
+                                    &mut key,
+                                    e.as_ref(),
+                                    k,
+                                )
+                            }
+                            windows.entry(key).or_insert_with(Vec::new)
+                        },
+                        #[cfg(feature = "dynamic_group_by")]
+                        WindowType::Rolling(options) => {
+                            rolling.entry(options).or_insert_with(Vec::new)
+                        },
+                    };
+                    entry.push((index, phys.clone()));
+                    is_window = true;
+                    break;
+                }
             }
+        } else {
+            // Window physical expressions always have the `Expr`.
+            is_window = false;
         }
         if !is_window {
             other.push((index, phys.as_ref()))

diff --git a/py-polars/tests/unit/operations/test_window.py b/py-polars/tests/unit/operations/test_window.py
@@ -503,3 +503,11 @@ def test_window_chunked_std_17102() -> None:
     df = pl.concat([c1, c2], rechunk=False)
     out = df.select(pl.col("B").std().over("A").alias("std"))
     assert out.unique().item() == 0.7071067811865476
+
+
+def test_window_17308() -> None:
+    df = pl.DataFrame({"A": [1, 2], "B": [3, 4], "grp": ["A", "B"]})
+
+    assert df.select(pl.col("A").sum(), pl.col("B").sum().over("grp")).to_dict(
+        as_series=False
+    ) == {"A": [3, 3], "B": [3, 4]}