From 28a45f2aa49dcda3f677f0c6fc1442e733875110 Mon Sep 17 00:00:00 2001 From: nameexhaustion Date: Fri, 17 Jan 2025 01:47:39 +1100 Subject: [PATCH] refactor: Fix flaky test (#20751) --- crates/polars-expr/src/expressions/apply.rs | 8 ++------ py-polars/tests/unit/operations/map/test_map_groups.py | 2 +- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/crates/polars-expr/src/expressions/apply.rs b/crates/polars-expr/src/expressions/apply.rs index 33c45043cb43..f4e6af091b12 100644 --- 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 index 5432f001a238..717b35303842 100644 --- 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]}), )