Fix NaN sentinel, groups filtering, and code quality issues

- Replace hardcoded "nan" sentinel with _DS_NAN_CATEGORY = "ds_nan" to
  avoid collision with user category names
- Fix groups filtering crash: wrap pd.Categorical in pd.Series before
  calling .reset_index() (shapes and points paths)
- Add groups + na_color=None filtering: when na_color is fully
  transparent, non-matching elements are removed instead of rendered
  invisibly (shapes and points)
- Extract _build_datashader_color_key() to deduplicate identical
  color_key dict building loops in shapes and points
- Extract _build_alignment_dtype_hint() from inline 57-line diagnostic
  block in _set_color_source_vec
- Fix "cola" typo in shape NaN test (should be "col_a")
- Remove dead duplicate instance_id assignments in test fixtures
- Update groups docstring to document na_color=None filtering behavior
- Add tests for groups + na_color=None filtering path

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: timtreis

src/spatialdata_plot/pl/basic.py

@@ -213,8 +213,9 @@ def render_shapes(
             `fill_alpha` will overwrite the value present in the cmap.
         groups : list[str] | str | None
             When using `color` and the key represents discrete labels, `groups` can be used to show only a subset of
-            them. Other values are set to NA. If elment is None, broadcasting behaviour is attempted (use the same
-            values for all elements).
+            them. Other values are set to NA. When ``na_color=None``, non-matching elements are filtered out entirely
+            (shapes and points only). If element is None, broadcasting behaviour is attempted (use the same values for
+            all elements).
         palette :  list[str] | str | None
             Palette for discrete annotations. List of valid color names that should be used for the categories. Must
             match the number of groups. If element is None, broadcasting behaviour is attempted (use the same values for

src/spatialdata_plot/pl/render.py

@@ -64,6 +64,10 @@
 
 _Normalize = Normalize | abc.Sequence[Normalize]
 
+# Sentinel category name used in datashader categorical paths to represent
+# missing (NaN) values.  Must not collide with realistic user category names.
+_DS_NAN_CATEGORY = "ds_nan"
+
 
 def _coerce_categorical_source(cat_source: Any) -> pd.Categorical:
     """Return a pandas Categorical from known, concrete sources only."""
@@ -82,6 +86,26 @@ def _coerce_categorical_source(cat_source: Any) -> pd.Categorical:
     return pd.Categorical(pd.Series(cat_source))
 
 
+def _build_datashader_color_key(
+    cat_series: pd.Categorical,
+    color_vector: Any,
+    na_color_hex: str,
+) -> dict[str, str]:
+    """Build a datashader ``color_key`` dict from a categorical series and its color vector."""
+    colors_arr = np.asarray(color_vector, dtype=object)
+    color_key: dict[str, str] = {}
+    for cat in cat_series.categories:
+        if cat == _DS_NAN_CATEGORY:
+            key_color = na_color_hex
+        else:
+            idx = np.flatnonzero(cat_series == cat)
+            key_color = colors_arr[idx[0]] if idx.size else na_color_hex
+        if isinstance(key_color, str) and key_color.startswith("#"):
+            key_color = _hex_no_alpha(key_color)
+        color_key[str(cat)] = key_color
+    return color_key
+
+
 def _split_colorbar_params(params: dict[str, object] | None) -> tuple[dict[str, object], dict[str, object], str | None]:
     """Split colorbar params into layout hints, Matplotlib kwargs, and label override."""
     layout: dict[str, object] = {}
@@ -185,6 +209,20 @@ def _render_shapes(
 
     values_are_categorical = color_source_vector is not None
 
+    # When groups are specified and na_color is fully transparent (na_color=None),
+    # filter out non-matching elements instead of showing them as invisible geometry.
+    if groups is not None and values_are_categorical and render_params.cmap_params.na_color.alpha == "00":
+        csv_series = pd.Series(color_source_vector)
+        keep = csv_series.isin(groups).values
+        shapes = shapes[keep].reset_index(drop=True)
+        sdata_filt[element] = shapes
+        color_source_vector = pd.Categorical(csv_series[keep].reset_index(drop=True))
+        color_vector = (
+            np.asarray(color_vector)[keep]
+            if not hasattr(color_vector, "reset_index")
+            else (color_vector[keep].reset_index(drop=True))
+        )
+
     # color_source_vector is None when the values aren't categorical
     if values_are_categorical and render_params.transfunc is not None:
         color_vector = render_params.transfunc(color_vector)
@@ -322,9 +360,9 @@ def _render_shapes(
         continuous_nan_shapes = None
         if col_for_color is not None and (render_params.groups is None or len(render_params.groups) > 1):
             if color_by_categorical:
-                # add nan as a category so that shapes with nan value are colored in the nan color
+                # add a sentinel category so that shapes with NaN value are colored in the na_color
                 transformed_element[col_for_color] = (
-                    transformed_element[col_for_color].cat.add_categories("nan").fillna("nan")
+                    transformed_element[col_for_color].cat.add_categories(_DS_NAN_CATEGORY).fillna(_DS_NAN_CATEGORY)
                 )
                 agg = cvs.polygons(
                     transformed_element,
@@ -391,17 +429,9 @@ def _render_shapes(
         color_key: dict[str, str] | None = None
         if color_by_categorical and col_for_color is not None:
             cat_series = _coerce_categorical_source(transformed_element[col_for_color])
-            colors_arr = np.asarray(color_vector, dtype=object)
-            color_key = {}
-            for cat in cat_series.categories:
-                if cat == "nan":
-                    key_color = render_params.cmap_params.na_color.get_hex()
-                else:
-                    idx = np.flatnonzero(cat_series == cat)
-                    key_color = colors_arr[idx[0]] if idx.size else render_params.cmap_params.na_color.get_hex()
-                if isinstance(key_color, str) and key_color.startswith("#"):
-                    key_color = _hex_no_alpha(key_color)
-                color_key[str(cat)] = key_color
+            color_key = _build_datashader_color_key(
+                cat_series, color_vector, render_params.cmap_params.na_color.get_hex()
+            )
 
         if color_by_categorical or col_for_color is None:
             ds_cmap = None
@@ -812,6 +842,27 @@ def _render_points(
         )
         points_dd = points_with_color_dd
 
+    # When groups are specified and na_color is fully transparent (na_color=None),
+    # filter out non-matching points instead of rendering invisible geometry.
+    if groups is not None and color_source_vector is not None and render_params.cmap_params.na_color.alpha == "00":
+        csv_series = pd.Series(color_source_vector)
+        keep = csv_series.isin(groups).values
+        color_source_vector = pd.Categorical(csv_series[keep].reset_index(drop=True))
+        color_vector = (
+            np.asarray(color_vector)[keep]
+            if not hasattr(color_vector, "reset_index")
+            else (color_vector[keep].reset_index(drop=True))
+        )
+        # re-register filtered points in sdata_filt
+        points_dd = dask.dataframe.from_pandas(points[keep].reset_index(drop=True), npartitions=1)
+        sdata_filt.points[element] = PointsModel.parse(points_dd, coordinates={"x": "x", "y": "y"})
+        set_transformation(
+            element=sdata_filt.points[element],
+            transformation=transformation_in_cs,
+            to_coordinate_system=coordinate_system,
+        )
+        n_points = int(keep.sum())
+
     # color_source_vector is None when the values aren't categorical
     if color_source_vector is None and render_params.transfunc is not None:
         color_vector = render_params.transfunc(color_vector)
@@ -895,9 +946,9 @@ def _render_points(
                     cat_series = cat_series.astype("category")
                 if hasattr(cat_series.cat, "as_known"):
                     cat_series = cat_series.cat.as_known()
-                if "nan" not in cat_series.cat.categories:
-                    cat_series = cat_series.cat.add_categories("nan")
-                transformed_element[col_for_color] = cat_series.fillna("nan")
+                if _DS_NAN_CATEGORY not in cat_series.cat.categories:
+                    cat_series = cat_series.cat.add_categories(_DS_NAN_CATEGORY)
+                transformed_element[col_for_color] = cat_series.fillna(_DS_NAN_CATEGORY)
                 agg = cvs.points(transformed_element, "x", "y", agg=ds.by(col_for_color, ds.count()))
             else:
                 reduction_name = render_params.ds_reduction if render_params.ds_reduction is not None else "sum"
@@ -942,17 +993,9 @@ def _render_points(
         color_key: dict[str, str] | None = None
         if color_by_categorical and col_for_color is not None:
             cat_series = _coerce_categorical_source(transformed_element[col_for_color])
-            colors_arr = np.asarray(color_vector, dtype=object)
-            color_key = {}
-            for cat in cat_series.categories:
-                if cat == "nan":
-                    key_color = render_params.cmap_params.na_color.get_hex()
-                else:
-                    idx = np.flatnonzero(cat_series == cat)
-                    key_color = colors_arr[idx[0]] if idx.size else render_params.cmap_params.na_color.get_hex()
-                if isinstance(key_color, str) and key_color.startswith("#"):
-                    key_color = _hex_no_alpha(key_color)
-                color_key[str(cat)] = key_color
+            color_key = _build_datashader_color_key(
+                cat_series, color_vector, render_params.cmap_params.na_color.get_hex()
+            )
 
         if (
             color_vector is not None

src/spatialdata_plot/pl/utils.py

@@ -983,6 +983,40 @@ def _infer_color_data_kind(
     return "numeric", pd.to_numeric(series, errors="coerce")
 
 
+def _build_alignment_dtype_hint(
+    sdata: sd.SpatialData | None,
+    element: object,
+    color_series: pd.Series,
+    table_name: str | None,
+) -> str:
+    """Build a diagnostic hint string for dtype mismatches between element and table indices."""
+    hints: list[str] = []
+    color_index_dtype = getattr(color_series.index, "dtype", None)
+    element_index_dtype = getattr(getattr(element, "index", None), "dtype", None) if element is not None else None
+
+    table_instance_dtype = None
+    instance_key = None
+    if table_name is not None and sdata is not None and table_name in sdata.tables:
+        table = sdata.tables[table_name]
+        try:
+            _, _, instance_key = get_table_keys(table)
+        except (KeyError, ValueError, TypeError, AttributeError):
+            instance_key = None
+        if instance_key is not None and hasattr(table, "obs") and instance_key in table.obs:
+            table_instance_dtype = table.obs[instance_key].dtype
+
+    if (
+        element_index_dtype is not None
+        and table_instance_dtype is not None
+        and element_index_dtype != table_instance_dtype
+    ):
+        hints.append(f"element index dtype is {element_index_dtype}, '{instance_key}' dtype is {table_instance_dtype}")
+    if color_index_dtype is not None and element_index_dtype is not None and color_index_dtype != element_index_dtype:
+        hints.append(f"color index dtype is {color_index_dtype}, element index dtype is {element_index_dtype}")
+
+    return f" (hint: {'; '.join(hints)})" if hints else ""
+
+
 def _set_color_source_vec(
     sdata: sd.SpatialData,
     element: SpatialElement | None,
@@ -1039,54 +1073,7 @@ def _set_color_source_vec(
         if color_series.isna().all():
             element_label = _format_element_name(element_name)
             location = f"table '{table_name}'" if table_name is not None else "the element"
-            # Provide dtype hints to help diagnose index alignment issues
-            dtype_hints: list[str] = []
-            color_index_dtype = getattr(color_series.index, "dtype", None)
-            element_index_dtype = (
-                getattr(getattr(element, "index", None), "dtype", None) if element is not None else None
-            )
-
-            table_instance_dtype = None
-            table_index_dtype = None
-            instance_key = None
-            if table_name is not None and sdata is not None and table_name in sdata.tables:
-                table = sdata.tables[table_name]
-                table_index_dtype = getattr(getattr(table, "obs", None), "index", None)
-                if table_index_dtype is not None:
-                    table_index_dtype = getattr(table_index_dtype, "dtype", None)
-                try:
-                    _, _, instance_key = get_table_keys(table)
-                except (KeyError, ValueError, TypeError, AttributeError):
-                    instance_key = None
-                if instance_key is not None and hasattr(table, "obs") and instance_key in table.obs:
-                    table_instance_dtype = table.obs[instance_key].dtype
-
-            if (
-                element_index_dtype is not None
-                and table_instance_dtype is not None
-                and element_index_dtype != table_instance_dtype
-            ):
-                dtype_hints.append(
-                    f"element index dtype is {element_index_dtype}, '{instance_key}' dtype is {table_instance_dtype}"
-                )
-            if (
-                table_index_dtype is not None
-                and table_instance_dtype is not None
-                and table_index_dtype != table_instance_dtype
-            ):
-                dtype_hints.append(
-                    f"table index dtype is {table_index_dtype}, '{instance_key}' dtype is {table_instance_dtype}"
-                )
-            if (
-                color_index_dtype is not None
-                and element_index_dtype is not None
-                and color_index_dtype != element_index_dtype
-            ):
-                dtype_hints.append(
-                    f"color index dtype is {color_index_dtype}, element index dtype is {element_index_dtype}"
-                )
-
-            dtype_hint = f" (hint: {'; '.join(dtype_hints)})" if dtype_hints else ""
+            dtype_hint = _build_alignment_dtype_hint(sdata, element, color_series, table_name)
             raise ValueError(
                 f"Column '{value_to_plot}' for element '{element_label}' contains only missing values after aligning "
                 f"with {location}. This usually means the instance ids/indices could not be aligned or converted, so "

tests/conftest.py

@@ -200,7 +200,6 @@ def sdata_blobs_points_with_nans_in_table() -> SpatialData:
     adata.var = pd.DataFrame({}, index=["col1", "col2"])
     adata.obs = pd.DataFrame(get_standard_RNG().normal(size=(n_obs, 3)), columns=["col_a", "col_b", "col_c"])
     adata.obs.iloc[0:30, adata.obs.columns.get_loc("col_a")] = np.nan
-    adata.obs["instance_id"] = np.arange(adata.n_obs)
     cat_pattern = ["a", "b", np.nan]
     repeats = (n_obs + len(cat_pattern) - 1) // len(cat_pattern)
     adata.obs["category"] = pd.Categorical((cat_pattern * repeats)[:n_obs])
@@ -221,7 +220,6 @@ def sdata_blobs_shapes_with_nans_in_table() -> SpatialData:
     adata.var = pd.DataFrame({}, index=["col1", "col2"])
     adata.obs = pd.DataFrame(get_standard_RNG().normal(size=(n_obs, 3)), columns=["col_a", "col_b", "col_c"])
     adata.obs.iloc[0, adata.obs.columns.get_loc("col_a")] = np.nan
-    adata.obs["instance_id"] = np.arange(adata.n_obs)
     cat_pattern = ["a", "b", np.nan, "c", "a"]
     repeats = (n_obs + len(cat_pattern) - 1) // len(cat_pattern)
     adata.obs["category"] = pd.Categorical((cat_pattern * repeats)[:n_obs])

tests/pl/test_render_points.py

@@ -576,6 +576,11 @@ def test_plot_can_annotate_points_with_nan_in_df_continuous_datashader(self, sda
         sdata_blobs["blobs_points"]["cont_color"] = pd.Series([np.nan, 2, 9, 13] * 50)
         sdata_blobs.pl.render_points("blobs_points", color="cont_color", size=40, method="datashader").pl.show()
 
+    def test_plot_groups_na_color_none_filters_points(self, sdata_blobs: SpatialData):
+        """When groups is set and na_color=None, non-matching points are filtered out entirely."""
+        sdata_blobs["blobs_points"]["cat_color"] = pd.Series(["a", "b", "c", "a"] * 50, dtype="category")
+        sdata_blobs.pl.render_points("blobs_points", color="cat_color", groups=["a"], na_color=None, size=30).pl.show()
+
 
 def test_raises_when_table_does_not_annotate_element(sdata_blobs: SpatialData):
     # Work on an independent copy since we mutate tables

tests/pl/test_render_shapes.py

@@ -946,13 +946,13 @@ def test_plot_can_annotate_shapes_with_nan_in_table_obs_categorical_datashader(
     def test_plot_can_annotate_shapes_with_nan_in_table_obs_continuous(
         self, sdata_blobs_shapes_with_nans_in_table: SpatialData
     ):
-        sdata_blobs_shapes_with_nans_in_table.pl.render_shapes("blobs_polygons", color="cola").pl.show()
+        sdata_blobs_shapes_with_nans_in_table.pl.render_shapes("blobs_polygons", color="col_a").pl.show()
 
     def test_plot_can_annotate_shapes_with_nan_in_table_obs_continuous_datashader(
         self, sdata_blobs_shapes_with_nans_in_table: SpatialData
     ):
         sdata_blobs_shapes_with_nans_in_table.pl.render_shapes(
-            "blobs_polygons", color="cola", method="datashader"
+            "blobs_polygons", color="col_a", method="datashader"
         ).pl.show()
 
     def test_plot_can_annotate_shapes_with_nan_in_table_X_continuous(
@@ -983,6 +983,11 @@ def test_plot_can_annotate_shapes_with_nan_in_df_continuous_datashader(self, sda
         sdata_blobs["blobs_polygons"]["cont_color"] = [np.nan, 2, 3, 4, 5]
         sdata_blobs.pl.render_shapes("blobs_polygons", color="cont_color", method="datashader").pl.show()
 
+    def test_plot_groups_na_color_none_filters_shapes(self, sdata_blobs: SpatialData):
+        """When groups is set and na_color=None, non-matching shapes are filtered out entirely."""
+        sdata_blobs["blobs_polygons"]["cat_color"] = pd.Series(["a", "b", "a", "b", "a"], dtype="category")
+        sdata_blobs.pl.render_shapes("blobs_polygons", color="cat_color", groups=["a"], na_color=None).pl.show()
+
 
 def test_plot_can_handle_nan_values_in_color_data(sdata_blobs: SpatialData, caplog):
     """Test that NaN values in color data are handled gracefully and logged."""