feat: Support nested keys in window functions (#20837)

pola-rs · Jan 22, 2025 · 714b679 · 714b679
1 parent 5c261c0
commit 714b679
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 8 deletions.
diff --git a/crates/polars-expr/src/expressions/window.rs b/crates/polars-expr/src/expressions/window.rs
@@ -567,16 +567,38 @@ impl PhysicalExpr for WindowExpr {
 
                         let get_join_tuples = || {
                             if group_by_columns.len() == 1 {
+                                let mut left = group_by_columns[0].clone();
                                 // group key from right column
-                                let right = &keys[0];
-                                PolarsResult::Ok(Arc::new(
-                                    group_by_columns[0]
-                                        .as_materialized_series()
-                                        .hash_join_left(
-                                            right.as_materialized_series(),
-                                            JoinValidation::ManyToMany,
-                                            true,
+                                let mut right = keys[0].clone();
+
+                                let (left, right) = if left.dtype().is_nested() {
+                                    (
+                                        ChunkedArray::<BinaryOffsetType>::with_chunk(
+                                            "".into(),
+                                            row_encode::_get_rows_encoded_unordered(&[
+                                                left.clone()
+                                            ])?
+                                            .into_array(),
+                                        )
+                                        .into_series(),
+                                        ChunkedArray::<BinaryOffsetType>::with_chunk(
+                                            "".into(),
+                                            row_encode::_get_rows_encoded_unordered(&[
+                                                right.clone()
+                                            ])?
+                                            .into_array(),
                                         )
+                                        .into_series(),
+                                    )
+                                } else {
+                                    (
+                                        left.into_materialized_series().clone(),
+                                        right.into_materialized_series().clone(),
+                                    )
+                                };
+
+                                PolarsResult::Ok(Arc::new(
+                                    left.hash_join_left(&right, JoinValidation::ManyToMany, true)
                                         .unwrap()
                                         .1,
                                 ))

diff --git a/py-polars/tests/unit/operations/test_window.py b/py-polars/tests/unit/operations/test_window.py
@@ -536,3 +536,9 @@ def test_order_by_sorted_keys_18943() -> None:
 
     out = df.set_sorted("g").select(pl.col("x").cum_sum().over("g", order_by="t"))
     assert_frame_equal(out, expect)
+
+
+def test_nested_window_keys() -> None:
+    df = pl.DataFrame({"x": 1, "y": "two"})
+    assert df.select(pl.col("y").first().over(pl.struct("x").implode())).item() == "two"
+    assert df.select(pl.col("y").first().over(pl.struct("x"))).item() == "two"