perf(color): extract one aligned column instead of copying the whole table

src/spatialdata_plot/pl/utils.py

@@ -1232,6 +1232,37 @@ def _build_alignment_dtype_hint(
     return ""
 
 
+def _extract_color_column(
+    table: AnnData,
+    value_key: str,
+    *,
+    origin: str,
+    element: GeoDataFrame,
+    element_name: str,
+    table_layer: str | None = None,
+) -> pd.Series:
+    """Read one color column from ``table`` aligned to ``element`` order, without copying the table.
+
+    Equivalent to ``get_values(value_key, sdata=..., element_name=..., table_name=...)[value_key]`` but
+    skips the table->element join, whose ``table[indices, :].copy()`` does an expensive out-of-order
+    sparse CSR row-gather. Restricts to rows annotating ``element_name`` (via ``region_key``), then
+    reindexes to the element's instance order (``NaN`` for instances with no table row), preserving the
+    categorical dtype of ``obs`` columns so the downstream legend path is unchanged.
+    """
+    attrs = table.uns["spatialdata_attrs"]
+    region_key, instance_key = attrs["region_key"], attrs["instance_key"]
+    mask = table.obs[region_key].to_numpy() == element_name
+    inst = table.obs[instance_key].to_numpy()[mask]
+    if origin == "var":
+        source = table.layers[table_layer] if table_layer is not None else table.X
+        col = source[:, table.var_names.get_loc(value_key)]
+        col = np.asarray(col.todense()).ravel() if hasattr(col, "todense") else np.asarray(col).ravel()
+        values = pd.Series(col[mask], index=inst)
+    else:  # obs column; .values keeps a Categorical categorical so the legend path still sees one
+        values = pd.Series(table.obs[value_key].values[mask], index=inst)
+    return values.reindex(element.index)
+
+
 def _set_color_source_vec(
     sdata: sd.SpatialData,
     element: SpatialElement | None,
@@ -1283,6 +1314,23 @@ def _set_color_source_vec(
             )
         if preloaded_color_data is not None:
             color_source_vector = preloaded_color_data
+        elif (
+            isinstance(element, GeoDataFrame)
+            and isinstance(element_name, str)
+            and table_name is not None
+            and table_name in sdata.tables
+            and origins[0].origin in ("obs", "var")
+        ):
+            # Fast path: read the single aligned column directly instead of joining/copying the
+            # whole annotating table (the join's out-of-order sparse row-gather dominates large renders).
+            color_source_vector = _extract_color_column(
+                sdata[table_name],
+                value_to_plot,
+                origin=origins[0].origin,
+                element=element,
+                element_name=element_name,
+                table_layer=table_layer,
+            )
         elif explicit_table_shadows_df:
             # Pass the table as `element` so upstream `get_values` skips the
             # element-column lookup and avoids the multi-origin error.

tests/pl/test_utils.py

@@ -679,3 +679,64 @@ def test_element_none_measures_single_table_elements(self, sdata_blobs: SpatialD
         # default blobs: only blobs_labels has a single annotating table
         measure_obs(sdata_blobs)
         assert "spatial" in sdata_blobs["table"].obsm
+
+
+class TestExtractColorColumn:
+    """`_extract_color_column` matches spatialdata's `get_values` bit-identically without copying the table."""
+
+    @staticmethod
+    def _annotated_shapes(n: int = 30, *, shuffle: bool = False, drop: int = 0, seed: int = 0) -> SpatialData:
+        rng = np.random.default_rng(seed)
+        coords = rng.random((n, 2)) * 100
+        geom = gpd.GeoDataFrame(
+            {"geometry": [Point(*xy) for xy in coords], "radius": np.ones(n)}, index=pd.Index(range(n))
+        )
+        inst = (rng.permutation(n) if shuffle else np.arange(n))[drop:]
+        adata = AnnData(
+            X=rng.random((len(inst), 4)).astype("float32"),
+            obs=pd.DataFrame(
+                {
+                    "region": pd.Categorical(["shapes"] * len(inst)),
+                    "instance_id": inst,
+                    "num": rng.random(len(inst)),
+                    "cat": pd.Categorical(rng.choice(list("abc"), len(inst))),
+                }
+            ),
+        )
+        adata.var_names = [f"g{i}" for i in range(4)]
+        table = TableModel.parse(adata, region="shapes", region_key="region", instance_key="instance_id")
+        return SpatialData(shapes={"shapes": ShapesModel.parse(geom)}, tables={"table": table})
+
+    @pytest.mark.parametrize(("key", "origin"), [("g0", "var"), ("g3", "var"), ("num", "obs"), ("cat", "obs")])
+    def test_matches_get_values(self, key: str, origin: str):
+        from spatialdata import get_values
+
+        from spatialdata_plot.pl.utils import _extract_color_column
+
+        sdata = self._annotated_shapes()
+        old = pd.Series(get_values(value_key=key, sdata=sdata, element_name="shapes", table_name="table")[key])
+        new = _extract_color_column(sdata["table"], key, origin=origin, element=sdata["shapes"], element_name="shapes")
+        assert (old.index == new.index).all()
+        if pd.api.types.is_numeric_dtype(old):
+            np.testing.assert_allclose(old.to_numpy(float), new.to_numpy(float))
+        else:
+            assert old.astype(str).equals(new.astype(str))
+            assert isinstance(new.dtype, pd.CategoricalDtype)  # preserved for the legend path
+
+    def test_shuffled_table_order_realigns(self):
+        from spatialdata import get_values
+
+        from spatialdata_plot.pl.utils import _extract_color_column
+
+        sdata = self._annotated_shapes(shuffle=True)
+        old = pd.Series(get_values(value_key="g0", sdata=sdata, element_name="shapes", table_name="table")["g0"])
+        new = _extract_color_column(sdata["table"], "g0", origin="var", element=sdata["shapes"], element_name="shapes")
+        np.testing.assert_allclose(old.to_numpy(float), new.to_numpy(float))
+
+    def test_missing_instances_become_nan(self):
+        from spatialdata_plot.pl.utils import _extract_color_column
+
+        sdata = self._annotated_shapes(drop=5)  # 5 shapes have no annotating table row
+        new = _extract_color_column(sdata["table"], "g0", origin="var", element=sdata["shapes"], element_name="shapes")
+        assert len(new) == 30
+        assert int(new.isna().sum()) == 5