fix: Fix Series.n_unique raising for list of struct (#20724)

pola-rs · Jan 15, 2025 · 958d00f · 958d00f
1 parent 46b2714
commit 958d00f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/crates/polars-core/src/series/implementations/list.rs b/crates/polars-core/src/series/implementations/list.rs
@@ -200,9 +200,6 @@ impl SeriesTrait for SeriesWrap<ListChunked> {
 
     #[cfg(feature = "algorithm_group_by")]
     fn n_unique(&self) -> PolarsResult<usize> {
-        if !self.inner_dtype().is_primitive_numeric() {
-            polars_bail!(opq = n_unique, self.dtype());
-        }
         // this can be called in aggregation, so this fast path can be worth a lot
         match self.len() {
             0 => Ok(0),

diff --git a/py-polars/tests/unit/operations/unique/test_n_unique.py b/py-polars/tests/unit/operations/unique/test_n_unique.py
@@ -46,3 +46,17 @@ def test_n_unique_null() -> None:
 )
 def test_n_unique_categorical(input: list[str | None], output: int) -> None:
     assert pl.Series(input, dtype=pl.Categorical).n_unique() == output
+
+
+def test_n_unique_list_of_struct_20341() -> None:
+    df = pl.DataFrame(
+        {
+            "a": [
+                [{"a": 1, "b": 2}, {"a": 10, "b": 20}],
+                [{"a": 1, "b": 2}, {"a": 10, "b": 20}],
+                [{"a": 3, "b": 4}],
+            ]
+        }
+    )
+    assert df.select("a").n_unique() == 2
+    assert df["a"].n_unique() == 2