From 053262ad3599cedad796f6dcf4b10bd33a99f78a Mon Sep 17 00:00:00 2001 From: anon Date: Wed, 10 Jun 2026 15:52:49 +0200 Subject: [PATCH] perf(color): extract one aligned column instead of copying the whole table Coloring shapes by a table column resolved the value via get_values(sdata, table_name), which joins the table to the element. That join's table[indices, :].copy() does an out-of-order sparse CSR row-gather that dominates large renders (~370 ms on Visium/Xenium-width tables). Add _extract_color_column: region-mask the annotating table, read the single column (var from X / layers, or obs preserving categorical dtype) and reindex to the element's instance order (NaN for unannotated instances). Wire it into _set_color_source_vec for the shapes (GeoDataFrame) + table-origin (obs/var) case; points already use the preloaded shortcut, labels keep get_values. Bit-identical to get_values (verified on visium_hne + curio; 6 unit tests across var / obs / categorical / shuffled-order / missing-instances), 14-1000x faster on the extraction itself. Also fixes a latent bug: the previous element=sdata[table_name] shortcut did not realign rows, silently mis-coloring when table order != element order. Phase 1 of the table-copy investigation; slimming the structural _join_table_for_element is a separate follow-up. --- src/spatialdata_plot/pl/utils.py | 48 +++++++++++++++++++++++++ tests/pl/test_utils.py | 61 ++++++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+) diff --git a/src/spatialdata_plot/pl/utils.py b/src/spatialdata_plot/pl/utils.py index 25720fb5..c1ef029f 100644 --- a/src/spatialdata_plot/pl/utils.py +++ b/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. diff --git a/tests/pl/test_utils.py b/tests/pl/test_utils.py index 87b41495..36db5729 100644 --- a/tests/pl/test_utils.py +++ b/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