Speed up datashader rendering of points (#379)

Datashader was consistently slower than matplotlib for points due to
five performance bottlenecks:

1. Dask DataFrame passed to cvs.points() instead of pandas (~137x
   scheduler overhead on already-computed data)
2. Double extent computation (get_extent on dask, then .compute again)
3. Per-point _hex_no_alpha() calls in O(n) list comprehension
4. _build_datashader_color_key iterated all points instead of
   early-exiting after finding all categories
5. _want_decorations created O(n) Python set from color vector

After fixes, datashader is 1.2-1.4x faster than matplotlib for plain
points and up to 1.6x faster for categorical coloring at 500K+ points.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: timtreis

src/spatialdata_plot/pl/_datashader.py

@@ -61,15 +61,22 @@ 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)
+    # Pre-extract as numpy arrays for fast access.  Count only categories that
+    # actually appear in the data (sentinel categories from _inject_ds_nan_sentinel
+    # may never appear in codes) so the early-exit triggers.  See #379.
+    categories_arr = np.asarray(cat_series.categories, dtype=str)
+    codes = np.asarray(cat_series.codes)
+    n_present = len(np.unique(codes[codes >= 0]))
     first_color: dict[str, str] = {}
-    for code, color in zip(cat_series.codes, colors_arr, strict=False):
+    for code, color in zip(codes, color_vector, strict=False):
         if code < 0:
             continue
-        cat_name = str(cat_series.categories[code])
+        cat_name = categories_arr[code]
         if cat_name not in first_color:
             first_color[cat_name] = _hex_no_alpha(color) if isinstance(color, str) and color.startswith("#") else color
-    return {str(c): first_color.get(str(c), na_hex) for c in cat_series.categories}
+            if len(first_color) == n_present:
+                break
+    return {c: first_color.get(c, na_hex) for c in categories_arr}
 
 
 def _inject_ds_nan_sentinel(series: pd.Series, sentinel: str = _DS_NAN_CATEGORY) -> pd.Series:

src/spatialdata_plot/pl/render.py

@@ -51,6 +51,7 @@
 from spatialdata_plot.pl.utils import (
     _ax_show_and_transform,
     _convert_shapes,
+    _datashader_canvas_from_dataframe,
     _decorate_axs,
     _get_collection_shape,
     _get_colors_for_categorical_obs,
@@ -81,14 +82,15 @@ def _want_decorations(color_vector: Any, na_color: Color) -> bool:
     cv = np.asarray(color_vector)
     if cv.size == 0:
         return False
-    unique_vals = set(cv.tolist())
-    if len(unique_vals) != 1:
+    # Fast check: if any value differs from the first, there is variety → show decorations.
+    first = cv.flat[0]
+    if not (cv == first).all():
         return True
-    only_val = next(iter(unique_vals))
+    # All values are the same — suppress decorations when that value is the NA color.
     na_hex = na_color.get_hex()
-    if isinstance(only_val, str) and only_val.startswith("#") and na_hex.startswith("#"):
-        return _hex_no_alpha(only_val) != _hex_no_alpha(na_hex)
-    return bool(only_val != na_hex)
+    if isinstance(first, str) and first.startswith("#") and na_hex.startswith("#"):
+        return _hex_no_alpha(first) != _hex_no_alpha(na_hex)
+    return bool(first != na_hex)
 
 
 def _reparse_points(
@@ -828,15 +830,18 @@ def _render_points(
         # use dpi/100 as a factor for cases where dpi!=100
         px = int(np.round(np.sqrt(render_params.size) * (fig_params.fig.dpi / 100)))
 
-        # apply transformations
+        # Apply transformations and materialize to pandas immediately.
+        # Keeping the data as a pandas DataFrame avoids dask scheduler overhead
+        # in both extent calculation and datashader aggregation (~137x faster
+        # for cvs.points on pre-computed data).  See #379.
         transformed_element = PointsModel.parse(
             trans.transform(sdata_filt.points[element][["x", "y"]]),
             annotation=sdata_filt.points[element][sdata_filt.points[element].columns.drop(["x", "y"])],
             transformations={coordinate_system: Identity()},
-        )
+        ).compute()
 
-        plot_width, plot_height, x_ext, y_ext, factor = _get_extent_and_range_for_datashader_canvas(
-            transformed_element, coordinate_system, ax, fig_params
+        plot_width, plot_height, x_ext, y_ext, factor = _datashader_canvas_from_dataframe(
+            transformed_element, ax, fig_params
         )
 
         # use datashader for the visualization of points
@@ -901,7 +906,10 @@ def _render_points(
             and isinstance(color_vector[0], str)
             and color_vector[0].startswith("#")
         ):
-            color_vector = np.asarray([_hex_no_alpha(x) for x in color_vector])
+            # Strip alpha from hex colors, deduplicating to avoid per-point work (#379).
+            unique_map = {c: _hex_no_alpha(c) for c in set(color_vector)}
+            if any(v != k for k, v in unique_map.items()):
+                color_vector = np.asarray([unique_map[c] for c in color_vector])
 
         nan_shaded = None
         if color_by_categorical or col_for_color is None:

src/spatialdata_plot/pl/utils.py

@@ -2973,15 +2973,16 @@ def set_zero_in_cmap_to_transparent(cmap: Colormap | str, steps: int | None = No
     return ListedColormap(colors)
 
 
-def _get_extent_and_range_for_datashader_canvas(
-    spatial_element: SpatialElement,
-    coordinate_system: str,
+def _compute_datashader_canvas_params(
+    x_ext: list[Any],
+    y_ext: list[Any],
     ax: Axes,
     fig_params: FigParams,
 ) -> tuple[Any, Any, list[Any], list[Any], Any]:
-    extent = get_extent(spatial_element, coordinate_system=coordinate_system)
-    x_ext = [min(0, extent["x"][0]), extent["x"][1]]
-    y_ext = [min(0, extent["y"][0]), extent["y"][1]]
+    """Compute datashader canvas dimensions from spatial extents.
+
+    Shared logic used by both the dask-based and pandas-based entry points.
+    """
     previous_xlim = ax.get_xlim()
     previous_ylim = ax.get_ylim()
     # increase range if sth larger was rendered on the axis before
@@ -3015,6 +3016,33 @@ def _get_extent_and_range_for_datashader_canvas(
     return plot_width, plot_height, x_ext, y_ext, factor
 
 
+def _get_extent_and_range_for_datashader_canvas(
+    spatial_element: SpatialElement,
+    coordinate_system: str,
+    ax: Axes,
+    fig_params: FigParams,
+) -> tuple[Any, Any, list[Any], list[Any], Any]:
+    extent = get_extent(spatial_element, coordinate_system=coordinate_system)
+    x_ext = [min(0, extent["x"][0]), extent["x"][1]]
+    y_ext = [min(0, extent["y"][0]), extent["y"][1]]
+    return _compute_datashader_canvas_params(x_ext, y_ext, ax, fig_params)
+
+
+def _datashader_canvas_from_dataframe(
+    df: pd.DataFrame,
+    ax: Axes,
+    fig_params: FigParams,
+) -> tuple[Any, Any, list[Any], list[Any], Any]:
+    """Compute datashader canvas params directly from a pandas DataFrame.
+
+    Avoids the overhead of ``get_extent()`` (which requires a dask-backed
+    SpatialElement) by reading min/max from the already-materialised data.
+    """
+    x_ext = [min(0, float(df["x"].min())), float(df["x"].max())]
+    y_ext = [min(0, float(df["y"].min())), float(df["y"].max())]
+    return _compute_datashader_canvas_params(x_ext, y_ext, ax, fig_params)
+
+
 def _create_image_from_datashader_result(
     ds_result: ds.transfer_functions.Image | np.ndarray[Any, np.dtype[np.uint8]],
     factor: float,