feat(rust, python): improve error msg if window expressions length do… (

#5262)
pola-rs · Oct 19, 2022 · 4008174 · 4008174
1 parent a0f13bc
commit 4008174
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 5 deletions.
diff --git a/polars/polars-core/src/series/unstable.rs b/polars/polars-core/src/series/unstable.rs
@@ -4,7 +4,7 @@ use std::ptr::NonNull;
 use crate::prelude::*;
 
 /// A wrapper type that should make it a bit more clear that we should not clone Series
-#[derive(Debug, Copy, Clone)]
+#[derive(Copy, Clone)]
 #[cfg(feature = "private")]
 pub struct UnstableSeries<'a> {
     lifetime: PhantomData<&'a Series>,

diff --git a/polars/polars-lazy/src/physical_plan/expressions/window.rs b/polars/polars-lazy/src/physical_plan/expressions/window.rs
@@ -356,10 +356,42 @@ impl PhysicalExpr for WindowExpr {
                 let out_column = ac.aggregated();
                 let flattened = out_column.explode()?;
                 if flattened.len() != df.height() {
-                    return Err(PolarsError::ComputeError(
-                        "the length of the window expression did not match that of the group"
-                            .into(),
-                    ));
+                    let ca = out_column.list().unwrap();
+                    let non_matching_group =
+                        ca.into_iter()
+                            .zip(ac.groups().iter())
+                            .find(|(output, group)| {
+                                if let Some(output) = output {
+                                    output.as_ref().len() != group.len()
+                                } else {
+                                    false
+                                }
+                            });
+
+                    return if let Some((output, group)) = non_matching_group {
+                        let first = group.first();
+                        let group = groupby_columns
+                            .iter()
+                            .map(|s| format!("{}", s.get(first as usize)))
+                            .collect::<Vec<_>>();
+                        let err_msg = format!(
+                            "{}\n> Group: ",
+                            "The length of the window expression did not match that of the group."
+                        );
+                        let err_msg = column_delimited(err_msg, &group);
+                        let err_msg = format!(
+                            "{}\n> Group length: {}\n> Output: '{:?}'",
+                            err_msg,
+                            group.len(),
+                            output.unwrap()
+                        );
+                        Err(PolarsError::ComputeError(err_msg.into()))
+                    } else {
+                        Err(PolarsError::ComputeError(
+                            "The length of the window expression did not match that of the group."
+                                .into(),
+                        ))
+                    };
                 }
 
                 // idx (new-idx, original-idx)

diff --git a/py-polars/tests/unit/test_errors.py b/py-polars/tests/unit/test_errors.py
@@ -236,3 +236,19 @@ def test_is_nan_on_non_boolean() -> None:
         pl.Series([1, 2, 3]).fill_nan(0)
     with pytest.raises(pl.InvalidOperationError):
         pl.Series(["1", "2", "3"]).fill_nan("2")  # type: ignore[arg-type]
+
+
+def test_window_expression_different_group_length() -> None:
+    try:
+        pl.DataFrame({"groups": ["a", "a", "b", "a", "b"]}).select(
+            [pl.col("groups").apply(lambda _: pl.Series([1, 2])).over("groups")]
+        )
+    except pl.ComputeError as e:
+        msg = str(e)
+        assert (
+            "The length of the window expression did not match that of the group."
+            in msg
+        )
+        assert "Group:" in msg
+        assert "Group length:" in msg
+        assert "Output: 'shape:" in msg