perf(render_points): keep categorical colour as codes+palette (memory at scale)

The resolved per-point colour is already a pd.Categorical (int codes + a small hex
palette), but every consumer re-expanded it to a per-point object array. Keep it
compact instead:

- _datashader._build_datashader_color_key: index color_vector at the one first-occurrence
  per category instead of np.asarray-ing the whole per-point vector.
- _datashader strip: strip alpha on the k categories and remap codes (no expand +
  refactorize + re-expand); array fallback unchanged.
- render._scatter_points: for non-uniform categorical hex, pass int codes +
  ListedColormap(categories) + BoundaryNorm instead of a per-point hex array; single
  category keeps the scalar color= fast path. NaN is already baked into the na_color
  category upstream (codes >= 0); a stray -1 falls back to the old path.

Byte-identical output: a main-vs-branch render of {categorical, continuous, uniform} x
{matplotlib, datashader} matches to the byte (max RGBA diff 0). At 20M points the
datashader strip drops from 720 MB / 962 ms to 340 MB / 117 ms (2.1x mem, 8x faster);
the matplotlib c= payload shrinks ~8x.

Tests: codes-path RGBA == per-point hex; single category -> scalar color; color-key
equivalence with the previous expansion.

Author: timtreis

src/spatialdata_plot/pl/_datashader.py

@@ -97,11 +97,10 @@ def _build_datashader_color_key(
 ) -> dict[str, str]:
     """Build a datashader ``color_key`` dict from a categorical series and its color vector."""
     na_hex = _hex_no_alpha(na_color_hex) if na_color_hex.startswith("#") else na_color_hex
-    colors_arr = np.asarray(color_vector, dtype=object)
     categories = np.asarray(cat_series.categories, dtype=str)
     codes = np.asarray(cat_series.codes)
 
-    if len(colors_arr) != len(codes):
+    if len(color_vector) != len(codes):
         logger.warning(
             f"color_vector length ({len(color_vector)}) does not match categorical series length "
             f"({len(codes)}); some categories may receive the na_color fallback."
@@ -111,11 +110,13 @@ def _build_datashader_color_key(
     # avoiding a Python loop over all points.  See #379.
     unique_codes, first_indices = np.unique(codes, return_index=True)
 
+    # Index color_vector only at those first occurrences (one per category) rather than expanding the
+    # whole per-point vector to an object array — color_vector is a compact pd.Categorical at scale.
     first_color: dict[str, str] = {}
     for code, idx in zip(unique_codes, first_indices, strict=True):
-        if code < 0 or idx >= len(colors_arr):
+        if code < 0 or idx >= len(color_vector):
             continue
-        c = colors_arr[idx]
+        c = color_vector[idx]
         first_color[categories[code]] = _hex_no_alpha(c) if isinstance(c, str) and c.startswith("#") else c
 
     return {cat: first_color.get(cat, na_hex) for cat in categories}
@@ -403,10 +404,17 @@ def _shade_datashader_aggregate(
         and isinstance(color_vector[0], str)
         and color_vector[0].startswith("#")
     ):
-        # Strip alpha on the unique colours and map back rather than parsing once per point; pd.factorize
-        # dedups in O(n) (hash, no sort) where np.unique would sort millions of strings.
-        codes, uniques = pd.factorize(np.asarray(color_vector))
-        color_vector = np.asarray([_hex_no_alpha(c) for c in uniques])[codes]
+        # Strip alpha on the unique colours, never on the per-point vector. color_vector is already a
+        # pd.Categorical (codes + a small palette) at scale: strip its k categories and remap the codes,
+        # staying compact. (Plain-array fallback factorizes in O(n) — hash, no sort.)
+        if isinstance(color_vector, pd.Categorical):
+            stripped = np.asarray([_hex_no_alpha(c) for c in color_vector.categories])
+            uniq_codes, uniques = pd.factorize(stripped)  # dedup over k categories, not n points
+            new_codes = np.where(color_vector.codes >= 0, uniq_codes[color_vector.codes], -1)
+            color_vector = pd.Categorical.from_codes(new_codes, categories=uniques)
+        else:
+            codes, uniques = pd.factorize(np.asarray(color_vector))
+            color_vector = np.asarray([_hex_no_alpha(c) for c in uniques])[codes]
 
     # density without a color column collapses to a sequential count gradient; everything else with no
     # explicit continuous value (categorical or no color) goes through the categorical shader.

src/spatialdata_plot/pl/render.py

@@ -19,7 +19,7 @@
 import xarray as xr
 from matplotlib import patheffects
 from matplotlib.cm import ScalarMappable
-from matplotlib.colors import Colormap, ListedColormap, Normalize
+from matplotlib.colors import BoundaryNorm, Colormap, ListedColormap, Normalize, to_rgba_array
 from scanpy._settings import settings as sc_settings
 from scanpy.plotting._tools.scatterplots import _add_categorical_legend
 from spatialdata import get_extent, get_values
@@ -1056,12 +1056,28 @@ def _scatter_points(
     # scalar ``color=`` instead of a per-point ``c=`` array: matplotlib then skips its per-point colour
     # machinery — the dominant cost at scale (10M points: ~9s -> ~3.7s) — for a visually identical result.
     # Numeric vectors keep the ``c=``/``cmap``/``norm`` path (they need the colormap).
-    cv = np.asarray(color_vector)
     color_kwargs: dict[str, Any]
-    if cv.ndim == 1 and cv.dtype.kind in "US" and _color_vector_is_uniform(cv):
-        color_kwargs = {"color": str(cv[0])}
+    if isinstance(color_vector, pd.Categorical) and (color_vector.codes >= 0).all():
+        # Categorical hex colours: pass the int codes + a ListedColormap of the (few) category colours
+        # instead of expanding to a per-point hex object array (~8x more memory at scale). BoundaryNorm
+        # edges at the half-integers map code i -> colormap entry i exactly, so the RGBA is identical.
+        # (_color.py bakes NaN into the na_color category, so codes >= 0; a stray -1 falls back below.)
+        categories = list(color_vector.categories)
+        if len(categories) == 1:
+            color_kwargs = {"color": str(categories[0])}  # uniform: scalar color=, skip per-point machinery
+        else:
+            n_cat = len(categories)
+            color_kwargs = {
+                "c": color_vector.codes,
+                "cmap": ListedColormap(to_rgba_array(categories)),
+                "norm": BoundaryNorm(np.arange(-0.5, n_cat, 1.0), ncolors=n_cat),
+            }
     else:
-        color_kwargs = {"c": color_vector, "cmap": cmap, "norm": norm}
+        cv = np.asarray(color_vector)
+        if cv.ndim == 1 and cv.dtype.kind in "US" and _color_vector_is_uniform(cv):
+            color_kwargs = {"color": str(cv[0])}
+        else:
+            color_kwargs = {"c": color_vector, "cmap": cmap, "norm": norm}
     return ax.scatter(
         x,
         y,

tests/pl/test_render_points.py

@@ -31,7 +31,7 @@
     _ds_shade_categorical,
     _pad_degenerate_extent,
 )
-from spatialdata_plot.pl.render import _marker_spread_px, _warn_groups_ignored_continuous
+from spatialdata_plot.pl.render import _marker_spread_px, _scatter_points, _warn_groups_ignored_continuous
 from tests.conftest import (
     CANVAS_HEIGHT,
     CANVAS_WIDTH,
@@ -1338,3 +1338,76 @@ def test_datashader_zero_extent_renders(coords):
     sdata = SpatialData(points={"points": PointsModel.parse(df)})
     sdata.pl.render_points("points", method="datashader").pl.show()
     plt.close("all")
+
+
+def _categorical_hex_vector(n: int, k: int, seed: int = 0) -> pd.Categorical:
+    """Per-point hex colour vector as _set_color_source_vec builds it: Categorical(source.map(palette))."""
+    rng = np.random.default_rng(seed)
+    cats = [f"ct{i}" for i in range(k)]
+    src = pd.Categorical(rng.choice(cats, n), categories=cats)
+    palette = {c: f"#{(i * 9973) % 0xFFFFFF:06x}ff" for i, c in enumerate(cats)}
+    return pd.Categorical(pd.Series(src).map(palette))
+
+
+def test_scatter_points_categorical_codes_match_per_point_hex():
+    # The codes+ListedColormap path must produce byte-identical RGBA to passing the per-point hex array.
+    cv = _categorical_hex_vector(400, k=6, seed=1)
+    rng = np.random.default_rng(2)
+    x, y = rng.random(400), rng.random(400)
+    common = {
+        "size": 10.0,
+        "cmap": plt.get_cmap("viridis"),
+        "norm": Normalize(),
+        "alpha": 1.0,
+        "trans_data": None,
+        "zorder": 1,
+    }
+
+    fig, ax = plt.subplots()
+    sc_codes = _scatter_points(ax, x, y, cv, **common)  # exercises the new codes path
+    sc_codes.update_scalarmappable()
+    fc_codes = sc_codes.get_facecolors()
+    plt.close(fig)
+
+    fig, ax = plt.subplots()
+    sc_hex = ax.scatter(x, y, c=np.asarray(cv), s=10.0)  # reference: literal per-point hex
+    fc_hex = sc_hex.get_facecolors()
+    plt.close(fig)
+
+    assert np.abs(fc_codes - fc_hex).max() == 0.0
+
+
+def test_scatter_points_single_category_uses_scalar_color():
+    # A single resolved colour must keep the scalar color= fast path (one facecolor, not a per-point array).
+    cv = _categorical_hex_vector(300, k=1, seed=3)
+    rng = np.random.default_rng(5)
+    fig, ax = plt.subplots()
+    sc = _scatter_points(
+        ax,
+        rng.random(300),
+        rng.random(300),
+        cv,
+        size=10.0,
+        cmap=plt.get_cmap("viridis"),
+        norm=Normalize(),
+        alpha=1.0,
+        trans_data=None,
+        zorder=1,
+    )
+    assert sc.get_facecolors().shape[0] == 1  # scalar color -> single facecolor
+    plt.close(fig)
+
+
+def test_build_datashader_color_key_categorical_indexes_palette():
+    # The color key must map each category to its alpha-stripped colour without expanding the full vector.
+    cv = _categorical_hex_vector(500, k=5, seed=4)
+    source = pd.Categorical([f"g{i % 5}" for i in range(500)], categories=[f"g{i}" for i in range(5)])
+    key = _build_datashader_color_key(source, cv, "#808080ff")
+    # equivalence with the previous object-array expansion:
+    colors_arr = np.asarray(cv, dtype=object)
+    codes = source.codes
+    expected = {}
+    for code, idx in zip(*np.unique(codes, return_index=True), strict=True):
+        c = colors_arr[idx]
+        expected[str(np.asarray(source.categories)[code])] = c[:7] if c.startswith("#") else c
+    assert key == expected