fix(rust): allow get access to list of categoricals (#14015)

pola-rs · Jan 26, 2024 · 3651ecd · 3651ecd
1 parent 298e71f
commit 3651ecd
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git a/crates/polars-ops/src/chunked_array/list/namespace.rs b/crates/polars-ops/src/chunked_array/list/namespace.rs
@@ -326,9 +326,12 @@ pub trait ListNameSpaceImpl: AsList {
             .downcast_iter()
             .map(|arr| sublist_get(arr, idx))
             .collect::<Vec<_>>();
-        Series::try_from((ca.name(), chunks))
-            .unwrap()
-            .cast(&ca.inner_dtype())
+        // Safety: every element in list has dtype equal to its inner type
+        unsafe {
+            Series::try_from((ca.name(), chunks))
+                .unwrap()
+                .cast_unchecked(&ca.inner_dtype())
+        }
     }
 
     #[cfg(feature = "list_gather")]

diff --git a/py-polars/tests/unit/namespaces/test_list.py b/py-polars/tests/unit/namespaces/test_list.py
@@ -69,6 +69,18 @@ def test_list_arr_get() -> None:
     ) == {"lists": [None, None, 4]}
 
 
+def test_list_categorical_get() -> None:
+    df = pl.DataFrame(
+        {
+            "actions": pl.Series(
+                [["a", "b"], ["c"], [None], None], dtype=pl.List(pl.Categorical)
+            ),
+        }
+    )
+    expected = pl.Series("actions", ["a", "c", None, None], dtype=pl.Categorical)
+    assert_series_equal(df["actions"].list.get(0), expected, categorical_as_str=True)
+
+
 def test_contains() -> None:
     a = pl.Series("a", [[1, 2, 3], [2, 5], [6, 7, 8, 9]])
     out = a.list.contains(2)