Keep category colors in sync when filtering tables

docs/changelog.md

@@ -13,4 +13,5 @@ For developers, please consult the [contributing guide](https://github.com/scver
 
 ### Fixed
 
+- Querying with `filter_table=True` now keeps `uns["<column>_colors"]` in sync with the categories of the filtered table; previously the colors of categories without any remaining rows were dropped while the categories were restored, so colors and categories no longer matched by position.
 - Querying a `SpatialData` object with multiple (batched) bounding boxes now raises an explicit `NotImplementedError` instead of failing with an opaque `AssertionError`. Batched queries are supported when querying a `SpatialElement` directly, but not when querying a `SpatialData` object, since that would require returning one `SpatialData` object per bounding box.

src/spatialdata/_utils.py

@@ -5,6 +5,7 @@
 import warnings
 from collections.abc import Callable, Generator
 from contextlib import contextmanager
+from copy import copy
 from itertools import islice
 from typing import Any, TypeVar
 
@@ -242,6 +243,14 @@ def _inplace_fix_subset_categorical_obs(subset_adata: AnnData | None, original_a
             c = obs[column].cat.set_categories(original_adata.obs[column].cat.categories)
             obs[column] = c
     subset_adata.obs = obs
+    # anndata drops the colors of unused categories together with the categories,
+    # restore them so that colors and categories are again matched by position
+    for column in obs.columns:
+        colors_key = f"{column}_colors"
+        if isinstance(obs[column].dtype, pd.CategoricalDtype) and colors_key in original_adata.uns:
+            original_colors = original_adata.uns[colors_key]
+            if len(original_colors) == len(obs[column].cat.categories):
+                subset_adata.uns[colors_key] = copy(original_colors)
 
 
 # TODO: change to paramspec as soon as we drop support for python 3.9, see https://stackoverflow.com/a/68290080

tests/core/query/test_spatial_query.py

@@ -1126,3 +1126,38 @@ def test_query_multiple_boxes_len_one(sdata_blobs):
         target_coordinate_system="global",
     )
     assert_spatial_data_objects_are_identical(queried0, queried1)
+
+
+def test_query_filter_table_keeps_category_colors():
+    # three cells of types A, C and B; B lies outside the query box
+    shapes = GeoDataFrame(
+        geometry=[Point(x, 0) for x in (0, 10, 100)],
+        index=["a", "c", "b"],
+    ).assign(radius=1.0)
+    obs = pd.DataFrame(
+        {
+            "region": pd.Categorical(["cells"] * 3),
+            "cell_id": ["a", "c", "b"],
+            "cell_type": pd.Categorical(["A", "C", "B"], categories=["A", "B", "C"]),
+        },
+        index=["a", "c", "b"],
+    )
+    table = AnnData(obs=obs)
+    table.uns["cell_type_colors"] = ["red", "green", "blue"]
+    sdata = SpatialData(
+        shapes={"cells": ShapesModel.parse(shapes)},
+        tables={"table": TableModel.parse(table, region="cells", region_key="region", instance_key="cell_id")},
+    )
+
+    queried = bounding_box_query(
+        sdata,
+        axes=("x", "y"),
+        min_coordinate=[-5, -5],
+        max_coordinate=[15, 5],
+        target_coordinate_system="global",
+        filter_table=True,
+    )
+    queried_table = queried["table"]
+    assert queried_table.n_obs == 2
+    assert queried_table.obs["cell_type"].cat.categories.tolist() == ["A", "B", "C"]
+    assert list(map(str, queried_table.uns["cell_type_colors"])) == ["red", "green", "blue"]