Merge fix from #387

Quantco · Sep 13, 2024 · 1d58498 · 1d58498
1 parent 34c4789
commit 1d58498
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/tabmat/categorical_matrix.py b/src/tabmat/categorical_matrix.py
@@ -261,8 +261,8 @@ def _extract_codes_and_categories(
             indices = pd_vec.cat.codes.to_numpy()
         elif namespace.__name__ == "polars":
             pl_vec: pl.Series = nw.to_native(cat_vec)
-            categories = pl_vec.cat.get_categories().to_numpy()
-            indices = pl_vec.to_physical().fill_null(-1).to_numpy()
+            categories = pl_vec.cat.to_local().cat.get_categories().to_numpy()
+            indices = pl_vec.cat.to_local().to_physical().fill_null(-1).to_numpy()
 
     return indices, categories, namespace
 

diff --git a/tests/test_categorical_matrix.py b/tests/test_categorical_matrix.py
@@ -271,5 +271,5 @@ def test_polars_non_contiguous_codes():
         _ = pl.Series(["beagle", "poodle", "labrador"], dtype=pl.Categorical)
         cat_series = pl.Series(str_series, dtype=pl.Categorical)
 
-    indices, categories = _extract_codes_and_categories(cat_series)
+    indices, categories, _ = _extract_codes_and_categories(cat_series)
     np.testing.assert_array_equal(str_series, categories[indices].tolist())