Skip to content

Commit 8d74219

Browse files
timtreistimtreis
authored
perf(color): extract one aligned column instead of copying the whole table (#709)
1 parent cb91f41 commit 8d74219

2 files changed

Lines changed: 109 additions & 0 deletions

File tree

‎src/spatialdata_plot/pl/utils.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1232,6 +1232,37 @@ def _build_alignment_dtype_hint(
12321232
return ""
12331233

12341234

1235+
def _extract_color_column(
1236+
table: AnnData,
1237+
value_key: str,
1238+
*,
1239+
origin: str,
1240+
element: GeoDataFrame,
1241+
element_name: str,
1242+
table_layer: str | None = None,
1243+
) -> pd.Series:
1244+
"""Read one color column from ``table`` aligned to ``element`` order, without copying the table.
1245+
1246+
Equivalent to ``get_values(value_key, sdata=..., element_name=..., table_name=...)[value_key]`` but
1247+
skips the table->element join, whose ``table[indices, :].copy()`` does an expensive out-of-order
1248+
sparse CSR row-gather. Restricts to rows annotating ``element_name`` (via ``region_key``), then
1249+
reindexes to the element's instance order (``NaN`` for instances with no table row), preserving the
1250+
categorical dtype of ``obs`` columns so the downstream legend path is unchanged.
1251+
"""
1252+
attrs = table.uns["spatialdata_attrs"]
1253+
region_key, instance_key = attrs["region_key"], attrs["instance_key"]
1254+
mask = table.obs[region_key].to_numpy() == element_name
1255+
inst = table.obs[instance_key].to_numpy()[mask]
1256+
if origin == "var":
1257+
source = table.layers[table_layer] if table_layer is not None else table.X
1258+
col = source[:, table.var_names.get_loc(value_key)]
1259+
col = np.asarray(col.todense()).ravel() if hasattr(col, "todense") else np.asarray(col).ravel()
1260+
values = pd.Series(col[mask], index=inst)
1261+
else: # obs column; .values keeps a Categorical categorical so the legend path still sees one
1262+
values = pd.Series(table.obs[value_key].values[mask], index=inst)
1263+
return values.reindex(element.index)
1264+
1265+
12351266
def _set_color_source_vec(
12361267
sdata: sd.SpatialData,
12371268
element: SpatialElement | None,
@@ -1283,6 +1314,23 @@ def _set_color_source_vec(
12831314
)
12841315
if preloaded_color_data is not None:
12851316
color_source_vector = preloaded_color_data
1317+
elif (
1318+
isinstance(element, GeoDataFrame)
1319+
and isinstance(element_name, str)
1320+
and table_name is not None
1321+
and table_name in sdata.tables
1322+
and origins[0].origin in ("obs", "var")
1323+
):
1324+
# Fast path: read the single aligned column directly instead of joining/copying the
1325+
# whole annotating table (the join's out-of-order sparse row-gather dominates large renders).
1326+
color_source_vector = _extract_color_column(
1327+
sdata[table_name],
1328+
value_to_plot,
1329+
origin=origins[0].origin,
1330+
element=element,
1331+
element_name=element_name,
1332+
table_layer=table_layer,
1333+
)
12861334
elif explicit_table_shadows_df:
12871335
# Pass the table as `element` so upstream `get_values` skips the
12881336
# element-column lookup and avoids the multi-origin error.

‎tests/pl/test_utils.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -679,3 +679,64 @@ def test_element_none_measures_single_table_elements(self, sdata_blobs: SpatialD
679679
# default blobs: only blobs_labels has a single annotating table
680680
measure_obs(sdata_blobs)
681681
assert "spatial" in sdata_blobs["table"].obsm
682+
683+
684+
class TestExtractColorColumn:
685+
"""`_extract_color_column` matches spatialdata's `get_values` bit-identically without copying the table."""
686+
687+
@staticmethod
688+
def _annotated_shapes(n: int = 30, *, shuffle: bool = False, drop: int = 0, seed: int = 0) -> SpatialData:
689+
rng = np.random.default_rng(seed)
690+
coords = rng.random((n, 2)) * 100
691+
geom = gpd.GeoDataFrame(
692+
{"geometry": [Point(*xy) for xy in coords], "radius": np.ones(n)}, index=pd.Index(range(n))
693+
)
694+
inst = (rng.permutation(n) if shuffle else np.arange(n))[drop:]
695+
adata = AnnData(
696+
X=rng.random((len(inst), 4)).astype("float32"),
697+
obs=pd.DataFrame(
698+
{
699+
"region": pd.Categorical(["shapes"] * len(inst)),
700+
"instance_id": inst,
701+
"num": rng.random(len(inst)),
702+
"cat": pd.Categorical(rng.choice(list("abc"), len(inst))),
703+
}
704+
),
705+
)
706+
adata.var_names = [f"g{i}" for i in range(4)]
707+
table = TableModel.parse(adata, region="shapes", region_key="region", instance_key="instance_id")
708+
return SpatialData(shapes={"shapes": ShapesModel.parse(geom)}, tables={"table": table})
709+
710+
@pytest.mark.parametrize(("key", "origin"), [("g0", "var"), ("g3", "var"), ("num", "obs"), ("cat", "obs")])
711+
def test_matches_get_values(self, key: str, origin: str):
712+
from spatialdata import get_values
713+
714+
from spatialdata_plot.pl.utils import _extract_color_column
715+
716+
sdata = self._annotated_shapes()
717+
old = pd.Series(get_values(value_key=key, sdata=sdata, element_name="shapes", table_name="table")[key])
718+
new = _extract_color_column(sdata["table"], key, origin=origin, element=sdata["shapes"], element_name="shapes")
719+
assert (old.index == new.index).all()
720+
if pd.api.types.is_numeric_dtype(old):
721+
np.testing.assert_allclose(old.to_numpy(float), new.to_numpy(float))
722+
else:
723+
assert old.astype(str).equals(new.astype(str))
724+
assert isinstance(new.dtype, pd.CategoricalDtype) # preserved for the legend path
725+
726+
def test_shuffled_table_order_realigns(self):
727+
from spatialdata import get_values
728+
729+
from spatialdata_plot.pl.utils import _extract_color_column
730+
731+
sdata = self._annotated_shapes(shuffle=True)
732+
old = pd.Series(get_values(value_key="g0", sdata=sdata, element_name="shapes", table_name="table")["g0"])
733+
new = _extract_color_column(sdata["table"], "g0", origin="var", element=sdata["shapes"], element_name="shapes")
734+
np.testing.assert_allclose(old.to_numpy(float), new.to_numpy(float))
735+
736+
def test_missing_instances_become_nan(self):
737+
from spatialdata_plot.pl.utils import _extract_color_column
738+
739+
sdata = self._annotated_shapes(drop=5) # 5 shapes have no annotating table row
740+
new = _extract_color_column(sdata["table"], "g0", origin="var", element=sdata["shapes"], element_name="shapes")
741+
assert len(new) == 30
742+
assert int(new.isna().sum()) == 5

0 commit comments

Comments
 (0)