Fix facet grid categorical missing IDs

Author: EltonChang1

plotnine/_utils/__init__.py

@@ -372,6 +372,9 @@ def _id_var(x: AnyArrayLike, drop: bool = False) -> list[int]:
         if drop:
             x = x.cat.remove_unused_categories()
             lst = list(x.cat.codes + 1)
+            if 0 in lst:
+                new_nan_code = max(lst) + 1
+                lst = [val if val != 0 else new_nan_code for val in lst]
         else:
             has_nan = any(np.isnan(i) for i in x if isinstance(i, float))
             if has_nan:

tests/test_facets.py

@@ -222,6 +222,25 @@ def test_facetting_with_unused_categories():
     p.draw_test()  # pyright: ignore
 
 
+def test_facet_grid_with_missing_categorical_values():
+    data = pd.DataFrame(
+        {
+            "x": [0.0] * 6,
+            "y": [0.0] * 6,
+            "row": ["a", "a", "a", "b", "b", "b"],
+            "col": pd.Categorical(
+                ["x", "y", None, "x", "y", None],
+                categories=["x", "y"],
+            ),
+        }
+    )
+
+    p = ggplot(data, aes("x", "y")) + geom_point() + facet_grid("row ~ col")
+
+    # No exception
+    p.draw_test()  # pyright: ignore
+
+
 def test_invalid_scales_value_raises():
     with pytest.raises(ValueError):
         # note the missing underscore

tests/test_utils.py

@@ -131,6 +131,19 @@ def test_ninteraction():
     assert ninteraction(data) == [1]
 
 
+def test_ninteraction_drops_unused_categorical_levels_with_missing():
+    data = pd.DataFrame(
+        {
+            "a": pd.Categorical(
+                ["x", "y", None, "x"],
+                categories=["x", "unused", "y"],
+            )
+        }
+    )
+
+    assert ninteraction(data, drop=True) == [1, 2, 3, 1]
+
+
 def test_ninteraction_datetime_series():
     # When a pandas datetime is converted Numpy datetime, the two
     # no longer compare as equal! This test ensures that case is