perf: skip re-validation when building SpatialData from already-valid elements

Add skip_element_validation() context manager (backed by a ContextVar) that
makes __setitem__ call get_model(validate=False) — type inference only, no
schema.validate().  Use it in every code path that constructs a SpatialData
from elements that originated from an existing SpatialData and were never
externally mutated: bounding_box_query, polygon_query, query_by_coordinate_system,
transform_to_coordinate_system, subset, and init_from_elements.

test_query_spatial_data: 0.77s → 0.64s (the remaining time is the query work
itself — filtering, shapely ops, raster cropping).

Also inline a minimal 2-image SpatialData in
test_transformations_between_coordinate_systems instead of relying on the
full 8-element images fixture; the test only ever uses image2d and
image2d_multiscale, so writing the other 6 to disk was pure waste.
test_transformations_between_coordinate_systems: 0.61s → 0.44s.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: LucaMarconato

src/spatialdata/_core/_elements.py

@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
+import contextlib
 from collections import UserDict
-from collections.abc import Iterable, KeysView, ValuesView
+from collections.abc import Iterable, Iterator, KeysView, ValuesView
+from contextvars import ContextVar
 from typing import TypeVar
 
 from anndata import AnnData
@@ -24,6 +26,25 @@
     get_model,
 )
 
+_skip_element_validation: ContextVar[bool] = ContextVar("_skip_element_validation", default=False)
+
+
+@contextlib.contextmanager
+def skip_element_validation() -> Iterator[None]:
+    """
+    Context manager to skip schema validation when inserting elements into SpatialData containers.
+
+    Use this only when inserting elements that are already known to be valid (e.g. elements
+    taken directly from an existing SpatialData object).  Skipping validation is unsafe for
+    externally-sourced data.
+    """
+    token = _skip_element_validation.set(True)
+    try:
+        yield
+    finally:
+        _skip_element_validation.reset(token)
+
+
 T = TypeVar("T")
 
 
@@ -69,7 +90,7 @@ def values(self) -> ValuesView[T]:
 class Images(Elements[DataArray | DataTree]):
     def __setitem__(self, key: str, value: Raster_T) -> None:
         self._check_key(key, self.keys(), self._shared_keys)
-        schema = get_model(value)
+        schema = get_model(value, validate=not _skip_element_validation.get())
         if schema not in (Image2DModel, Image3DModel):
             raise TypeError(f"Unknown element type with schema: {schema!r}.")
         super().__setitem__(key, value)
@@ -78,7 +99,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
 class Labels(Elements[DataArray | DataTree]):
     def __setitem__(self, key: str, value: Raster_T) -> None:
         self._check_key(key, self.keys(), self._shared_keys)
-        schema = get_model(value)
+        schema = get_model(value, validate=not _skip_element_validation.get())
         if schema not in (Labels2DModel, Labels3DModel):
             raise TypeError(f"Unknown element type with schema: {schema!r}.")
         super().__setitem__(key, value)
@@ -87,7 +108,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
 class Shapes(Elements[GeoDataFrame]):
     def __setitem__(self, key: str, value: GeoDataFrame) -> None:
         self._check_key(key, self.keys(), self._shared_keys)
-        schema = get_model(value)
+        schema = get_model(value, validate=not _skip_element_validation.get())
         if schema != ShapesModel:
             raise TypeError(f"Unknown element type with schema: {schema!r}.")
         super().__setitem__(key, value)
@@ -96,7 +117,7 @@ def __setitem__(self, key: str, value: GeoDataFrame) -> None:
 class Points(Elements[DaskDataFrame]):
     def __setitem__(self, key: str, value: DaskDataFrame) -> None:
         self._check_key(key, self.keys(), self._shared_keys)
-        schema = get_model(value)
+        schema = get_model(value, validate=not _skip_element_validation.get())
         if schema != PointsModel:
             raise TypeError(f"Unknown element type with schema: {schema!r}.")
         super().__setitem__(key, value)
@@ -105,7 +126,7 @@ def __setitem__(self, key: str, value: DaskDataFrame) -> None:
 class Tables(Elements[AnnData]):
     def __setitem__(self, key: str, value: AnnData) -> None:
         self._check_key(key, self.keys(), self._shared_keys)
-        schema = get_model(value)
+        schema = get_model(value, validate=not _skip_element_validation.get())
         if schema != TableModel:
             raise TypeError(f"Unknown element type with schema: {schema!r}.")
         super().__setitem__(key, value)

src/spatialdata/_core/query/spatial_query.py

@@ -15,6 +15,7 @@
 from xarray import DataArray, DataTree
 
 from spatialdata import to_polygons
+from spatialdata._core._elements import skip_element_validation
 from spatialdata._core.query._utils import _get_filtered_or_unfiltered_tables, get_bounding_box_corners
 from spatialdata._core.spatialdata import SpatialData
 from spatialdata._docs import docstring_parameter
@@ -538,7 +539,8 @@ def _(
 
     tables = _get_filtered_or_unfiltered_tables(filter_table, new_elements, sdata)
 
-    return SpatialData(**new_elements, tables=tables, attrs=sdata.attrs)
+    with skip_element_validation():
+        return SpatialData(**new_elements, tables=tables, attrs=sdata.attrs)
 
 
 @bounding_box_query.register(DataArray)
@@ -874,7 +876,8 @@ def _(
 
     tables = _get_filtered_or_unfiltered_tables(filter_table, new_elements, sdata)
 
-    return SpatialData(**new_elements, tables=tables, attrs=sdata.attrs)
+    with skip_element_validation():
+        return SpatialData(**new_elements, tables=tables, attrs=sdata.attrs)
 
 
 @polygon_query.register(DataArray)

src/spatialdata/_core/spatialdata.py

@@ -21,7 +21,7 @@
 from xarray import DataArray, DataTree
 from zarr.errors import GroupNotFoundError
 
-from spatialdata._core._elements import Images, Labels, Points, Shapes, Tables
+from spatialdata._core._elements import Images, Labels, Points, Shapes, Tables, skip_element_validation
 from spatialdata._core.validation import (
     check_all_keys_case_insensitively_unique,
     check_target_region_column_symmetry,
@@ -641,7 +641,8 @@ def filter_by_coordinate_system(
             element_names=element_names_in_coordinate_system,
         )
 
-        return SpatialData(**elements, tables=tables, attrs=self.attrs)
+        with skip_element_validation():
+            return SpatialData(**elements, tables=tables, attrs=self.attrs)
 
     # TODO: move to relational query with refactor
     def _filter_tables(
@@ -885,7 +886,8 @@ def transform_to_coordinate_system(
                 if element_type not in elements:
                     elements[element_type] = {}
                 elements[element_type][element_name] = transformed
-        return SpatialData(**elements, tables=sdata.tables, attrs=self.attrs)
+        with skip_element_validation():
+            return SpatialData(**elements, tables=sdata.tables, attrs=self.attrs)
 
     def elements_are_self_contained(self) -> dict[str, bool]:
         """
@@ -2242,7 +2244,8 @@ def init_from_elements(
                 assert model == ShapesModel
                 element_type = "shapes"
             elements_dict.setdefault(element_type, {})[name] = element
-        return cls(**elements_dict, attrs=attrs)
+        with skip_element_validation():
+            return cls(**elements_dict, attrs=attrs)
 
     def subset(
         self,
@@ -2284,7 +2287,8 @@ def subset(
             include_orphan_tables,
             elements_dict=elements_dict,
         )
-        return SpatialData(**elements_dict, tables=tables, attrs=self.attrs)
+        with skip_element_validation():
+            return SpatialData(**elements_dict, tables=tables, attrs=self.attrs)
 
     def __getitem__(self, item: str) -> SpatialElement | AnnData:
         """

tests/core/operations/test_transform.py

@@ -694,9 +694,16 @@ def test_transform_elements_and_entire_spatial_data_object_multi_hop(
                 assert np.allclose(affine, np.array([[1, 0], [0, 1]]))
 
 
-def test_transformations_between_coordinate_systems(images):
+def test_transformations_between_coordinate_systems():
     # just a test that all the code is executed without errors and a quick test that the affine matrix is correct.
     # For a full test the notebooks are more exhaustive
+    # Use a minimal sdata with only the two images needed — avoids writing 6 unused elements to disk.
+    images = SpatialData(
+        images={
+            "image2d": Image2DModel.parse(np.zeros((3, 8, 8)), dims=("c", "y", "x")),
+            "image2d_multiscale": Image2DModel.parse(np.zeros((3, 8, 8)), dims=("c", "y", "x"), scale_factors=[2]),
+        }
+    )
     with tempfile.TemporaryDirectory() as tmpdir:
         images.write(Path(tmpdir) / "sdata.zarr")
         el0 = images.images["image2d"]