Skip to content

Commit ccdfbda

Browse files
authored
fix: Fix panic in window case (#17320)
1 parent 41610e3 commit ccdfbda

File tree

2 files changed

+41
-28
lines changed

2 files changed

+41
-28
lines changed

crates/polars-mem-engine/src/executors/projection_utils.rs

Lines changed: 33 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -117,36 +117,41 @@ fn execute_projection_cached_window_fns(
117117
// first we partition the window function by the values they group over.
118118
// the group_by values should be cached
119119
exprs.iter().enumerate_u32().for_each(|(index, phys)| {
120-
let e = phys.as_expression().unwrap();
121-
122120
let mut is_window = false;
123-
for e in e.into_iter() {
124-
if let Expr::Window {
125-
partition_by,
126-
options,
127-
order_by,
128-
..
129-
} = e
130-
{
131-
let entry = match options {
132-
WindowType::Over(_) => {
133-
let mut key = format!("{:?}", partition_by.as_slice());
134-
if let Some((e, k)) = order_by {
135-
polars_expr::prelude::window_function_format_order_by(
136-
&mut key,
137-
e.as_ref(),
138-
k,
139-
)
140-
}
141-
windows.entry(key).or_insert_with(Vec::new)
142-
},
143-
#[cfg(feature = "dynamic_group_by")]
144-
WindowType::Rolling(options) => rolling.entry(options).or_insert_with(Vec::new),
145-
};
146-
entry.push((index, phys.clone()));
147-
is_window = true;
148-
break;
121+
if let Some(e) = phys.as_expression() {
122+
for e in e.into_iter() {
123+
if let Expr::Window {
124+
partition_by,
125+
options,
126+
order_by,
127+
..
128+
} = e
129+
{
130+
let entry = match options {
131+
WindowType::Over(_) => {
132+
let mut key = format!("{:?}", partition_by.as_slice());
133+
if let Some((e, k)) = order_by {
134+
polars_expr::prelude::window_function_format_order_by(
135+
&mut key,
136+
e.as_ref(),
137+
k,
138+
)
139+
}
140+
windows.entry(key).or_insert_with(Vec::new)
141+
},
142+
#[cfg(feature = "dynamic_group_by")]
143+
WindowType::Rolling(options) => {
144+
rolling.entry(options).or_insert_with(Vec::new)
145+
},
146+
};
147+
entry.push((index, phys.clone()));
148+
is_window = true;
149+
break;
150+
}
149151
}
152+
} else {
153+
// Window physical expressions always have the `Expr`.
154+
is_window = false;
150155
}
151156
if !is_window {
152157
other.push((index, phys.as_ref()))

py-polars/tests/unit/operations/test_window.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -503,3 +503,11 @@ def test_window_chunked_std_17102() -> None:
503503
df = pl.concat([c1, c2], rechunk=False)
504504
out = df.select(pl.col("B").std().over("A").alias("std"))
505505
assert out.unique().item() == 0.7071067811865476
506+
507+
508+
def test_window_17308() -> None:
509+
df = pl.DataFrame({"A": [1, 2], "B": [3, 4], "grp": ["A", "B"]})
510+
511+
assert df.select(pl.col("A").sum(), pl.col("B").sum().over("grp")).to_dict(
512+
as_series=False
513+
) == {"A": [3, 3], "B": [3, 4]}

0 commit comments

Comments
 (0)