refactor: Fix flaky test (#20751)

pola-rs · Jan 16, 2025 · 28a45f2 · 28a45f2
1 parent cf4b9bd
commit 28a45f2
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/crates/polars-expr/src/expressions/apply.rs b/crates/polars-expr/src/expressions/apply.rs
@@ -301,14 +301,10 @@ impl ApplyExpr {
 fn all_unit_length(ca: &ListChunked) -> bool {
     assert_eq!(ca.chunks().len(), 1);
 
-    // Handles the Null dtype - in that case the offsets can be (e.g. [0,0,0 ...])
-    if ca.null_count() == ca.len() {
-        return true;
-    }
-
     let list_arr = ca.downcast_iter().next().unwrap();
     let offset = list_arr.offsets().as_slice();
-    (offset[offset.len() - 1] as usize) == list_arr.len()
+    // Note: Checking offset.last() == 0 handles the Null dtype - in that case the offsets can be (e.g. [0,0,0 ...])
+    (offset[offset.len() - 1] as usize) == list_arr.len() || offset[offset.len() - 1] == 0
 }
 
 fn check_map_output_len(input_len: usize, output_len: usize, expr: &Expr) -> PolarsResult<()> {

diff --git a/py-polars/tests/unit/operations/map/test_map_groups.py b/py-polars/tests/unit/operations/map/test_map_groups.py
@@ -162,6 +162,6 @@ def foo(x: pl.Series) -> pl.Series:
         pl.DataFrame({"key": [0, 0, 1], "a": [None, None, None]})
         .group_by("key")
         .agg(pl.map_groups(exprs=["a"], function=foo))  # type: ignore[arg-type]
-        .sort("a"),
+        .sort("key"),
         pl.DataFrame({"key": [0, 1], "a": [None, None]}),
     )