Reduce time cost of tests

.github/workflows/test.yaml

@@ -20,13 +20,10 @@ jobs:
  matrix:
  include:
  - {os: windows-latest, python: "3.11", dask-version: "2025.12.0", name: "min dask"}
- - {os: windows-latest, python: "3.13", dask-version: "latest"}
  - {os: windows-latest, python: "3.14", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.11", dask-version: "latest"}
- - {os: ubuntu-latest, python: "3.13", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.14", dask-version: "latest"}
  - {os: macos-latest, python: "3.11", dask-version: "latest"}
- - {os: macos-latest, python: "3.13", prerelease: "allow", name: "prerelease"}
  - {os: macos-latest, python: "3.14", prerelease: "allow", name: "prerelease"}
  env:
  OS: ${{ matrix.os }}
@@ -62,7 +59,7 @@ jobs:
  PLATFORM: ${{ matrix.os }}
  DISPLAY: :42
  run: |
- uv run pytest --cov --color=yes --cov-report=xml
+ uv run pytest --cov --color=yes --cov-report=xml -n auto --dist worksteal
  - name: Upload coverage to Codecov
  uses: codecov/codecov-action@v5
  with:

pyproject.toml

@@ -69,6 +69,7 @@ test = [
  "pytest",
  "pytest-cov",
  "pytest-mock",
+ "pytest-xdist",
  "torch",
 ]
 docs = [
@@ -102,6 +103,7 @@ strict = true
 addopts = [
  "--import-mode=importlib", # allow using test files with same name
  "-s", # print output from tests
+ "-p no:napari", # napari registers a pytest plugin via its entry point; disable it here since spatialdata tests don't need it
 ]
 # These are all markers coming from xarray, dask or anndata. Added here to silence warnings.
 markers = [

src/spatialdata/_core/_elements.py

@@ -21,7 +21,6 @@
  PointsModel,
  ShapesModel,
  TableModel,
- get_axes_names,
  get_model,
 )
 
@@ -50,6 +49,7 @@ def _check_key(key: str, element_keys: Iterable[str], shared_keys: set[str | Non
  raise KeyError(*e.args) from e
 
  def __setitem__(self, key: str, value: T) -> None:
+ # note that each __setitem__ in the subclasses calls get_model(), which performs data validation
  self._add_shared_key(key)
  super().__setitem__(key, value)
 
@@ -72,15 +72,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Image2DModel, Image3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 3:
- Image2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 4:
- Image3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 4.")
+ super().__setitem__(key, value)
 
 
 class Labels(Elements[DataArray | DataTree]):
@@ -89,15 +81,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Labels2DModel, Labels3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 2:
- Labels2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 3:
- Labels3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 3.")
+ super().__setitem__(key, value)
 
 
 class Shapes(Elements[GeoDataFrame]):
@@ -106,7 +90,6 @@ def __setitem__(self, key: str, value: GeoDataFrame) -> None:
  schema = get_model(value)
  if schema != ShapesModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ShapesModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -116,7 +99,6 @@ def __setitem__(self, key: str, value: DaskDataFrame) -> None:
  schema = get_model(value)
  if schema != PointsModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- PointsModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -126,5 +108,4 @@ def __setitem__(self, key: str, value: AnnData) -> None:
  schema = get_model(value)
  if schema != TableModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- TableModel.validate(value)
  super().__setitem__(key, value)

src/spatialdata/_core/centroids.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from functools import singledispatch
 
-import dask.array as da
+import numpy as np
 import pandas as pd
 import xarray as xr
 from dask.dataframe import DataFrame as DaskDataFrame
@@ -56,40 +55,29 @@ def get_centroids(
  raise ValueError(f"The object type {type(e)} is not supported.")
 
 
-def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
+def _get_centroids_for_labels(xdata: xr.DataArray) -> pd.DataFrame:
  """
- Compute the component "axis" of the centroid of each label as a weighted average of the xarray coordinates.
+ Compute centroids for all labels in a DataArray in a single O(n_voxels) pass.
 
- Parameters
- ----------
- xdata
- The xarray DataArray containing the labels.
- axis
- The axis for which the centroids are computed.
-
- Returns
- -------
- pd.DataFrame
- A DataFrame containing one column, named after "axis", with the centroids of the labels along that axis.
- The index of the DataFrame is the collection of label values, sorted in ascending order.
+ Works for any number of spatial dimensions (2D and 3D labels).
  """
-centroids: dict[int, float] = defaultdict(float)
-for i in xdata[axis]:
- portion = xdata.sel(**{axis: i}).data
- u = da.unique(portion, return_counts=True)
- labels_values = u[0].compute()
- counts = u[1].compute()
- for j in range(len(labels_values)):
- label_value = labels_values[j]
- count = counts[j]
- centroids[label_value] += count * i.values.item()
-
-all_labels_values, all_labels_counts = da.unique(xdata.data, return_counts=True)
-all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute(), strict=True))
-for label_value in centroids:
-centroids[label_value] /= all_labels[label_value]
- centroids = dict(sorted(centroids.items(), key=lambda x: x[0]))
- return pd.DataFrame({axis: centroids.values()}, index=list(centroids.keys()))
+arr = xdata.data.compute()
+axes = list(xdata.dims)
+
+# Map label values to a contiguous range for bincount efficiency.
+label_ids, inverse = np.unique(arr, return_inverse=True)
+flat_inverse = inverse.ravel()
+counts = np.bincount(flat_inverse) # per-label pixel counts
+
+# indexing="ij" (matrix convention) ensures the i-th grid varies along the i-th
+# dimension of the output, correctly aligning with xdata.dims for any number of axes.
+ coord_grids = np.meshgrid(*[xdata[ax].values for ax in axes], indexing="ij")
+data: dict[str, np.ndarray] = {}
+for ax, grid in zip(axes, coord_grids, strict=True):
+ coord_sums = np.bincount(flat_inverse, weights=grid.ravel().astype(float))
+data[ax] = coord_sums / counts # counts > 0 by construction (unique guarantees this)
+
+ return pd.DataFrame(data, index=label_ids)
 
 
 @get_centroids.register(DataArray)
@@ -109,10 +97,7 @@ def _(
  assert len(e["scale0"]) == 1
  e = next(iter(e["scale0"].values()))
 
- dfs = []
- for axis in get_axes_names(e):
- dfs.append(_get_centroids_for_axis(e, axis))
- df = pd.concat(dfs, axis=1)
+ df = _get_centroids_for_labels(e)
  if not return_background and 0 in df.index:
  df = df.drop(index=0) # drop the background label
  t = get_transformation(e, coordinate_system)

src/spatialdata/models/models.py

@@ -1250,14 +1250,17 @@ def parse(
 
 def get_model(
  e: SpatialElement,
+ validate: bool = True,
 ) -> Schema_t:
  """
- Get the model for the given element.
+ Get the model for the given element. Validate using the model if `validate` is `True`.
 
  Parameters
  ----------
  e
  The element.
+ validate
+ Whether to validate the element using the model.
 
  Returns
  -------
@@ -1268,7 +1271,8 @@ def _validate_and_return(
  schema: Schema_t,
  e: SpatialElement,
  ) -> Schema_t:
- schema.validate(e)
+ if validate:
+ schema.validate(e)
  return schema
 
  if isinstance(e, DataArray | DataTree):

tests/conftest.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+
+# Disable numba JIT for the test suite (the test data is small so initializing the JIT is slower than using plain
+# Python)
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
+
+import copy as _copy
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any
 
@@ -33,6 +40,25 @@
  TableModel,
 )
 
+
+def _fast_deepcopy_sdata(sd: SpatialData) -> SpatialData:
+ """
+ Fast deepcopy for SpatialData objects in tests.
+
+ Uses copy.deepcopy (which skips model re-validation) and manually restores
+ the attrs that copy.deepcopy loses for DaskDataFrame (issue #503) and
+ GeoDataFrame (issue #286).
+ """
+ points_attrs = {k: _copy.deepcopy(v._attrs) for k, v in sd.points.items()}
+ shapes_attrs = {k: _copy.deepcopy(v.attrs) for k, v in sd.shapes.items()}
+ sd_copy = _copy.deepcopy(sd)
+ for k, attrs in points_attrs.items():
+ sd_copy.points[k]._attrs = attrs
+ for k, attrs in shapes_attrs.items():
+ sd_copy.shapes[k].attrs = attrs
+ return sd_copy
+
+
 SEED = 0
 RNG = default_rng(seed=SEED)
 
@@ -41,16 +67,26 @@
 POINT_PATH = Path(__file__).parent / "data/points.json"
 
 
-@pytest.fixture()
-def images() -> SpatialData:
+@pytest.fixture(scope="session")
+def _images_session() -> SpatialData:
  return SpatialData(images=_get_images())
 
 
 @pytest.fixture()
-def labels() -> SpatialData:
+def images(_images_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_images_session)
+
+
+@pytest.fixture(scope="session")
+def _labels_session() -> SpatialData:
  return SpatialData(labels=_get_labels())
 
 
+@pytest.fixture()
+def labels(_labels_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_labels_session)
+
+
 @pytest.fixture()
 def shapes() -> SpatialData:
  return SpatialData(shapes=_get_shapes())
@@ -87,8 +123,8 @@ def tables() -> list[AnnData]:
  return _tables
 
 
-@pytest.fixture()
-def full_sdata() -> SpatialData:
+@pytest.fixture(scope="session")
+def _full_sdata_session() -> SpatialData:
  return SpatialData(
  images=_get_images(),
  labels=_get_labels(),
@@ -98,6 +134,22 @@ def full_sdata() -> SpatialData:
  )
 
 
+@pytest.fixture()
+def full_sdata(_full_sdata_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_full_sdata_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_full_session() -> SpatialData:
+ return SpatialData(
+ images=_get_images(),
+ labels=_get_labels(),
+ shapes=_get_shapes(),
+ points=_get_points(),
+ tables=_get_tables(region="labels2d"),
+ )
+
+
 @pytest.fixture(
  # params=["labels"]
  params=["full", "empty"]
@@ -110,15 +162,9 @@ def full_sdata() -> SpatialData:
  ]
  # + ["empty_" + x for x in ["table"]] # TODO: empty table not supported yet
 )
-def sdata(request) -> SpatialData:
+def sdata(request, _sdata_full_session: SpatialData) -> SpatialData:
  if request.param == "full":
- return SpatialData(
- images=_get_images(),
- labels=_get_labels(),
- shapes=_get_shapes(),
- points=_get_points(),
- tables=_get_tables(region="labels2d"),
- )
+ return _fast_deepcopy_sdata(_sdata_full_session)
  if request.param == "empty":
  return SpatialData()
  return request.getfixturevalue(request.param)
@@ -304,18 +350,38 @@ def _get_new_table(spatial_element: None | str | Sequence[str], instance_id: Non
  return TableModel.parse(adata=adata, spatial_element=spatial_element, instance_id=instance_id)
 
 
-@pytest.fixture()
-def labels_blobs() -> ArrayLike:
- """Create a 2D labels."""
+@pytest.fixture(scope="session")
+def _labels_blobs_session() -> ArrayLike:
  return BlobsDataset()._labels_blobs()
 
 
 @pytest.fixture()
-def sdata_blobs() -> SpatialData:
+def labels_blobs(_labels_blobs_session: ArrayLike) -> ArrayLike:
  """Create a 2D labels."""
+ return deepcopy(_labels_blobs_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_blobs_session() -> SpatialData:
  from spatialdata.datasets import blobs
 
- return deepcopy(blobs(256, 300, 3))
+ return blobs(256, 300, 3)
+
+
+@pytest.fixture()
+def sdata_blobs(_sdata_blobs_session: SpatialData) -> SpatialData:
+ """Create a 2D labels."""
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+
+@pytest.fixture()
+def blobs_factory(_sdata_blobs_session: SpatialData) -> Callable[[], SpatialData]:
+ """Return a factory that creates cheap fresh copies of the session-scoped blobs dataset."""
+
+ def _make() -> SpatialData:
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+ return _make
 
 
 def _make_points(coordinates: np.ndarray) -> DaskDataFrame:

tests/core/operations/conftest.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import os
+
+# Disable numba JIT compilation for rasterize tests. Datashader (used by rasterize) triggers
+# numba JIT on first call, costing ~1.4s per worker. Python-mode gives identical results for
+# the small test data here — unlike real data, there is no throughput advantage from JIT.
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")

.github/workflows/test.yaml

@@ -20,13 +20,10 @@ jobs:
  matrix:
  include:
  - {os: windows-latest, python: "3.11", dask-version: "2025.12.0", name: "min dask"}
- - {os: windows-latest, python: "3.13", dask-version: "latest"}
  - {os: windows-latest, python: "3.14", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.11", dask-version: "latest"}
- - {os: ubuntu-latest, python: "3.13", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.14", dask-version: "latest"}
  - {os: macos-latest, python: "3.11", dask-version: "latest"}
- - {os: macos-latest, python: "3.13", prerelease: "allow", name: "prerelease"}
  - {os: macos-latest, python: "3.14", prerelease: "allow", name: "prerelease"}
  env:
  OS: ${{ matrix.os }}
@@ -62,7 +59,7 @@ jobs:
  PLATFORM: ${{ matrix.os }}
  DISPLAY: :42
  run: |
- uv run pytest --cov --color=yes --cov-report=xml
+ uv run pytest --cov --color=yes --cov-report=xml -n auto --dist worksteal
  - name: Upload coverage to Codecov
  uses: codecov/codecov-action@v5
  with:

pyproject.toml

@@ -69,6 +69,7 @@ test = [
  "pytest",
  "pytest-cov",
  "pytest-mock",
+ "pytest-xdist",
  "torch",
 ]
 docs = [
@@ -102,6 +103,7 @@ strict = true
 addopts = [
  "--import-mode=importlib", # allow using test files with same name
  "-s", # print output from tests
+ "-p no:napari", # napari registers a pytest plugin via its entry point; disable it here since spatialdata tests don't need it
 ]
 # These are all markers coming from xarray, dask or anndata. Added here to silence warnings.
 markers = [

src/spatialdata/_core/_elements.py

@@ -21,7 +21,6 @@
  PointsModel,
  ShapesModel,
  TableModel,
- get_axes_names,
  get_model,
 )
 
@@ -50,6 +49,7 @@ def _check_key(key: str, element_keys: Iterable[str], shared_keys: set[str | Non
  raise KeyError(*e.args) from e
 
  def __setitem__(self, key: str, value: T) -> None:
+ # note that each __setitem__ in the subclasses calls get_model(), which performs data validation
  self._add_shared_key(key)
  super().__setitem__(key, value)
 
@@ -72,15 +72,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Image2DModel, Image3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 3:
- Image2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 4:
- Image3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 4.")
+ super().__setitem__(key, value)
 
 
 class Labels(Elements[DataArray | DataTree]):
@@ -89,15 +81,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Labels2DModel, Labels3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 2:
- Labels2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 3:
- Labels3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 3.")
+ super().__setitem__(key, value)
 
 
 class Shapes(Elements[GeoDataFrame]):
@@ -106,7 +90,6 @@ def __setitem__(self, key: str, value: GeoDataFrame) -> None:
  schema = get_model(value)
  if schema != ShapesModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ShapesModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -116,7 +99,6 @@ def __setitem__(self, key: str, value: DaskDataFrame) -> None:
  schema = get_model(value)
  if schema != PointsModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- PointsModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -126,5 +108,4 @@ def __setitem__(self, key: str, value: AnnData) -> None:
  schema = get_model(value)
  if schema != TableModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- TableModel.validate(value)
  super().__setitem__(key, value)

src/spatialdata/_core/centroids.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from functools import singledispatch
 
-import dask.array as da
+import numpy as np
 import pandas as pd
 import xarray as xr
 from dask.dataframe import DataFrame as DaskDataFrame
@@ -56,40 +55,29 @@ def get_centroids(
  raise ValueError(f"The object type {type(e)} is not supported.")
 
 
-def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
+def _get_centroids_for_labels(xdata: xr.DataArray) -> pd.DataFrame:
  """
- Compute the component "axis" of the centroid of each label as a weighted average of the xarray coordinates.
+ Compute centroids for all labels in a DataArray in a single O(n_voxels) pass.
 
- Parameters
- ----------
- xdata
- The xarray DataArray containing the labels.
- axis
- The axis for which the centroids are computed.
-
- Returns
- -------
- pd.DataFrame
- A DataFrame containing one column, named after "axis", with the centroids of the labels along that axis.
- The index of the DataFrame is the collection of label values, sorted in ascending order.
+ Works for any number of spatial dimensions (2D and 3D labels).
  """
-centroids: dict[int, float] = defaultdict(float)
-for i in xdata[axis]:
- portion = xdata.sel(**{axis: i}).data
- u = da.unique(portion, return_counts=True)
- labels_values = u[0].compute()
- counts = u[1].compute()
- for j in range(len(labels_values)):
- label_value = labels_values[j]
- count = counts[j]
- centroids[label_value] += count * i.values.item()
-
-all_labels_values, all_labels_counts = da.unique(xdata.data, return_counts=True)
-all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute(), strict=True))
-for label_value in centroids:
-centroids[label_value] /= all_labels[label_value]
- centroids = dict(sorted(centroids.items(), key=lambda x: x[0]))
- return pd.DataFrame({axis: centroids.values()}, index=list(centroids.keys()))
+arr = xdata.data.compute()
+axes = list(xdata.dims)
+
+# Map label values to a contiguous range for bincount efficiency.
+label_ids, inverse = np.unique(arr, return_inverse=True)
+flat_inverse = inverse.ravel()
+counts = np.bincount(flat_inverse) # per-label pixel counts
+
+# indexing="ij" (matrix convention) ensures the i-th grid varies along the i-th
+# dimension of the output, correctly aligning with xdata.dims for any number of axes.
+ coord_grids = np.meshgrid(*[xdata[ax].values for ax in axes], indexing="ij")
+data: dict[str, np.ndarray] = {}
+for ax, grid in zip(axes, coord_grids, strict=True):
+ coord_sums = np.bincount(flat_inverse, weights=grid.ravel().astype(float))
+data[ax] = coord_sums / counts # counts > 0 by construction (unique guarantees this)
+
+ return pd.DataFrame(data, index=label_ids)
 
 
 @get_centroids.register(DataArray)
@@ -109,10 +97,7 @@ def _(
  assert len(e["scale0"]) == 1
  e = next(iter(e["scale0"].values()))
 
- dfs = []
- for axis in get_axes_names(e):
- dfs.append(_get_centroids_for_axis(e, axis))
- df = pd.concat(dfs, axis=1)
+ df = _get_centroids_for_labels(e)
  if not return_background and 0 in df.index:
  df = df.drop(index=0) # drop the background label
  t = get_transformation(e, coordinate_system)

src/spatialdata/models/models.py

@@ -1250,14 +1250,17 @@ def parse(
 
 def get_model(
  e: SpatialElement,
+ validate: bool = True,
 ) -> Schema_t:
  """
- Get the model for the given element.
+ Get the model for the given element. Validate using the model if `validate` is `True`.
 
  Parameters
  ----------
  e
  The element.
+ validate
+ Whether to validate the element using the model.
 
  Returns
  -------
@@ -1268,7 +1271,8 @@ def _validate_and_return(
  schema: Schema_t,
  e: SpatialElement,
  ) -> Schema_t:
- schema.validate(e)
+ if validate:
+ schema.validate(e)
  return schema
 
  if isinstance(e, DataArray | DataTree):

tests/conftest.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+
+# Disable numba JIT for the test suite (the test data is small so initializing the JIT is slower than using plain
+# Python)
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
+
+import copy as _copy
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any
 
@@ -33,6 +40,25 @@
  TableModel,
 )
 
+
+def _fast_deepcopy_sdata(sd: SpatialData) -> SpatialData:
+ """
+ Fast deepcopy for SpatialData objects in tests.
+
+ Uses copy.deepcopy (which skips model re-validation) and manually restores
+ the attrs that copy.deepcopy loses for DaskDataFrame (issue #503) and
+ GeoDataFrame (issue #286).
+ """
+ points_attrs = {k: _copy.deepcopy(v._attrs) for k, v in sd.points.items()}
+ shapes_attrs = {k: _copy.deepcopy(v.attrs) for k, v in sd.shapes.items()}
+ sd_copy = _copy.deepcopy(sd)
+ for k, attrs in points_attrs.items():
+ sd_copy.points[k]._attrs = attrs
+ for k, attrs in shapes_attrs.items():
+ sd_copy.shapes[k].attrs = attrs
+ return sd_copy
+
+
 SEED = 0
 RNG = default_rng(seed=SEED)
 
@@ -41,16 +67,26 @@
 POINT_PATH = Path(__file__).parent / "data/points.json"
 
 
-@pytest.fixture()
-def images() -> SpatialData:
+@pytest.fixture(scope="session")
+def _images_session() -> SpatialData:
  return SpatialData(images=_get_images())
 
 
 @pytest.fixture()
-def labels() -> SpatialData:
+def images(_images_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_images_session)
+
+
+@pytest.fixture(scope="session")
+def _labels_session() -> SpatialData:
  return SpatialData(labels=_get_labels())
 
 
+@pytest.fixture()
+def labels(_labels_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_labels_session)
+
+
 @pytest.fixture()
 def shapes() -> SpatialData:
  return SpatialData(shapes=_get_shapes())
@@ -87,8 +123,8 @@ def tables() -> list[AnnData]:
  return _tables
 
 
-@pytest.fixture()
-def full_sdata() -> SpatialData:
+@pytest.fixture(scope="session")
+def _full_sdata_session() -> SpatialData:
  return SpatialData(
  images=_get_images(),
  labels=_get_labels(),
@@ -98,6 +134,22 @@ def full_sdata() -> SpatialData:
  )
 
 
+@pytest.fixture()
+def full_sdata(_full_sdata_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_full_sdata_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_full_session() -> SpatialData:
+ return SpatialData(
+ images=_get_images(),
+ labels=_get_labels(),
+ shapes=_get_shapes(),
+ points=_get_points(),
+ tables=_get_tables(region="labels2d"),
+ )
+
+
 @pytest.fixture(
  # params=["labels"]
  params=["full", "empty"]
@@ -110,15 +162,9 @@ def full_sdata() -> SpatialData:
  ]
  # + ["empty_" + x for x in ["table"]] # TODO: empty table not supported yet
 )
-def sdata(request) -> SpatialData:
+def sdata(request, _sdata_full_session: SpatialData) -> SpatialData:
  if request.param == "full":
- return SpatialData(
- images=_get_images(),
- labels=_get_labels(),
- shapes=_get_shapes(),
- points=_get_points(),
- tables=_get_tables(region="labels2d"),
- )
+ return _fast_deepcopy_sdata(_sdata_full_session)
  if request.param == "empty":
  return SpatialData()
  return request.getfixturevalue(request.param)
@@ -304,18 +350,38 @@ def _get_new_table(spatial_element: None | str | Sequence[str], instance_id: Non
  return TableModel.parse(adata=adata, spatial_element=spatial_element, instance_id=instance_id)
 
 
-@pytest.fixture()
-def labels_blobs() -> ArrayLike:
- """Create a 2D labels."""
+@pytest.fixture(scope="session")
+def _labels_blobs_session() -> ArrayLike:
  return BlobsDataset()._labels_blobs()
 
 
 @pytest.fixture()
-def sdata_blobs() -> SpatialData:
+def labels_blobs(_labels_blobs_session: ArrayLike) -> ArrayLike:
  """Create a 2D labels."""
+ return deepcopy(_labels_blobs_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_blobs_session() -> SpatialData:
  from spatialdata.datasets import blobs
 
- return deepcopy(blobs(256, 300, 3))
+ return blobs(256, 300, 3)
+
+
+@pytest.fixture()
+def sdata_blobs(_sdata_blobs_session: SpatialData) -> SpatialData:
+ """Create a 2D labels."""
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+
+@pytest.fixture()
+def blobs_factory(_sdata_blobs_session: SpatialData) -> Callable[[], SpatialData]:
+ """Return a factory that creates cheap fresh copies of the session-scoped blobs dataset."""
+
+ def _make() -> SpatialData:
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+ return _make
 
 
 def _make_points(coordinates: np.ndarray) -> DaskDataFrame:

tests/core/operations/conftest.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import os
+
+# Disable numba JIT compilation for rasterize tests. Datashader (used by rasterize) triggers
+# numba JIT on first call, costing ~1.4s per worker. Python-mode gives identical results for
+# the small test data here — unlike real data, there is no throughput advantage from JIT.
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")

.github/workflows/test.yaml

@@ -20,13 +20,10 @@ jobs:
  matrix:
  include:
  - {os: windows-latest, python: "3.11", dask-version: "2025.12.0", name: "min dask"}
- - {os: windows-latest, python: "3.13", dask-version: "latest"}
  - {os: windows-latest, python: "3.14", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.11", dask-version: "latest"}
- - {os: ubuntu-latest, python: "3.13", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.14", dask-version: "latest"}
  - {os: macos-latest, python: "3.11", dask-version: "latest"}
- - {os: macos-latest, python: "3.13", prerelease: "allow", name: "prerelease"}
  - {os: macos-latest, python: "3.14", prerelease: "allow", name: "prerelease"}
  env:
  OS: ${{ matrix.os }}
@@ -62,7 +59,7 @@ jobs:
  PLATFORM: ${{ matrix.os }}
  DISPLAY: :42
  run: |
- uv run pytest --cov --color=yes --cov-report=xml
+ uv run pytest --cov --color=yes --cov-report=xml -n auto --dist worksteal
  - name: Upload coverage to Codecov
  uses: codecov/codecov-action@v5
  with:

pyproject.toml

@@ -69,6 +69,7 @@ test = [
  "pytest",
  "pytest-cov",
  "pytest-mock",
+ "pytest-xdist",
  "torch",
 ]
 docs = [
@@ -102,6 +103,7 @@ strict = true
 addopts = [
  "--import-mode=importlib", # allow using test files with same name
  "-s", # print output from tests
+ "-p no:napari", # napari registers a pytest plugin via its entry point; disable it here since spatialdata tests don't need it
 ]
 # These are all markers coming from xarray, dask or anndata. Added here to silence warnings.
 markers = [

src/spatialdata/_core/_elements.py

@@ -21,7 +21,6 @@
  PointsModel,
  ShapesModel,
  TableModel,
- get_axes_names,
  get_model,
 )
 
@@ -50,6 +49,7 @@ def _check_key(key: str, element_keys: Iterable[str], shared_keys: set[str | Non
  raise KeyError(*e.args) from e
 
  def __setitem__(self, key: str, value: T) -> None:
+ # note that each __setitem__ in the subclasses calls get_model(), which performs data validation
  self._add_shared_key(key)
  super().__setitem__(key, value)
 
@@ -72,15 +72,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Image2DModel, Image3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 3:
- Image2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 4:
- Image3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 4.")
+ super().__setitem__(key, value)
 
 
 class Labels(Elements[DataArray | DataTree]):
@@ -89,15 +81,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Labels2DModel, Labels3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 2:
- Labels2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 3:
- Labels3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 3.")
+ super().__setitem__(key, value)
 
 
 class Shapes(Elements[GeoDataFrame]):
@@ -106,7 +90,6 @@ def __setitem__(self, key: str, value: GeoDataFrame) -> None:
  schema = get_model(value)
  if schema != ShapesModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ShapesModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -116,7 +99,6 @@ def __setitem__(self, key: str, value: DaskDataFrame) -> None:
  schema = get_model(value)
  if schema != PointsModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- PointsModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -126,5 +108,4 @@ def __setitem__(self, key: str, value: AnnData) -> None:
  schema = get_model(value)
  if schema != TableModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- TableModel.validate(value)
  super().__setitem__(key, value)

src/spatialdata/_core/centroids.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from functools import singledispatch
 
-import dask.array as da
+import numpy as np
 import pandas as pd
 import xarray as xr
 from dask.dataframe import DataFrame as DaskDataFrame
@@ -56,40 +55,29 @@ def get_centroids(
  raise ValueError(f"The object type {type(e)} is not supported.")
 
 
-def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
+def _get_centroids_for_labels(xdata: xr.DataArray) -> pd.DataFrame:
  """
- Compute the component "axis" of the centroid of each label as a weighted average of the xarray coordinates.
+ Compute centroids for all labels in a DataArray in a single O(n_voxels) pass.
 
- Parameters
- ----------
- xdata
- The xarray DataArray containing the labels.
- axis
- The axis for which the centroids are computed.
-
- Returns
- -------
- pd.DataFrame
- A DataFrame containing one column, named after "axis", with the centroids of the labels along that axis.
- The index of the DataFrame is the collection of label values, sorted in ascending order.
+ Works for any number of spatial dimensions (2D and 3D labels).
  """
-centroids: dict[int, float] = defaultdict(float)
-for i in xdata[axis]:
- portion = xdata.sel(**{axis: i}).data
- u = da.unique(portion, return_counts=True)
- labels_values = u[0].compute()
- counts = u[1].compute()
- for j in range(len(labels_values)):
- label_value = labels_values[j]
- count = counts[j]
- centroids[label_value] += count * i.values.item()
-
-all_labels_values, all_labels_counts = da.unique(xdata.data, return_counts=True)
-all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute(), strict=True))
-for label_value in centroids:
-centroids[label_value] /= all_labels[label_value]
- centroids = dict(sorted(centroids.items(), key=lambda x: x[0]))
- return pd.DataFrame({axis: centroids.values()}, index=list(centroids.keys()))
+arr = xdata.data.compute()
+axes = list(xdata.dims)
+
+# Map label values to a contiguous range for bincount efficiency.
+label_ids, inverse = np.unique(arr, return_inverse=True)
+flat_inverse = inverse.ravel()
+counts = np.bincount(flat_inverse) # per-label pixel counts
+
+# indexing="ij" (matrix convention) ensures the i-th grid varies along the i-th
+# dimension of the output, correctly aligning with xdata.dims for any number of axes.
+ coord_grids = np.meshgrid(*[xdata[ax].values for ax in axes], indexing="ij")
+data: dict[str, np.ndarray] = {}
+for ax, grid in zip(axes, coord_grids, strict=True):
+ coord_sums = np.bincount(flat_inverse, weights=grid.ravel().astype(float))
+data[ax] = coord_sums / counts # counts > 0 by construction (unique guarantees this)
+
+ return pd.DataFrame(data, index=label_ids)
 
 
 @get_centroids.register(DataArray)
@@ -109,10 +97,7 @@ def _(
  assert len(e["scale0"]) == 1
  e = next(iter(e["scale0"].values()))
 
- dfs = []
- for axis in get_axes_names(e):
- dfs.append(_get_centroids_for_axis(e, axis))
- df = pd.concat(dfs, axis=1)
+ df = _get_centroids_for_labels(e)
  if not return_background and 0 in df.index:
  df = df.drop(index=0) # drop the background label
  t = get_transformation(e, coordinate_system)

src/spatialdata/models/models.py

@@ -1250,14 +1250,17 @@ def parse(
 
 def get_model(
  e: SpatialElement,
+ validate: bool = True,
 ) -> Schema_t:
  """
- Get the model for the given element.
+ Get the model for the given element. Validate using the model if `validate` is `True`.
 
  Parameters
  ----------
  e
  The element.
+ validate
+ Whether to validate the element using the model.
 
  Returns
  -------
@@ -1268,7 +1271,8 @@ def _validate_and_return(
  schema: Schema_t,
  e: SpatialElement,
  ) -> Schema_t:
- schema.validate(e)
+ if validate:
+ schema.validate(e)
  return schema
 
  if isinstance(e, DataArray | DataTree):

tests/conftest.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+
+# Disable numba JIT for the test suite (the test data is small so initializing the JIT is slower than using plain
+# Python)
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
+
+import copy as _copy
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any
 
@@ -33,6 +40,25 @@
  TableModel,
 )
 
+
+def _fast_deepcopy_sdata(sd: SpatialData) -> SpatialData:
+ """
+ Fast deepcopy for SpatialData objects in tests.
+
+ Uses copy.deepcopy (which skips model re-validation) and manually restores
+ the attrs that copy.deepcopy loses for DaskDataFrame (issue #503) and
+ GeoDataFrame (issue #286).
+ """
+ points_attrs = {k: _copy.deepcopy(v._attrs) for k, v in sd.points.items()}
+ shapes_attrs = {k: _copy.deepcopy(v.attrs) for k, v in sd.shapes.items()}
+ sd_copy = _copy.deepcopy(sd)
+ for k, attrs in points_attrs.items():
+ sd_copy.points[k]._attrs = attrs
+ for k, attrs in shapes_attrs.items():
+ sd_copy.shapes[k].attrs = attrs
+ return sd_copy
+
+
 SEED = 0
 RNG = default_rng(seed=SEED)
 
@@ -41,16 +67,26 @@
 POINT_PATH = Path(__file__).parent / "data/points.json"
 
 
-@pytest.fixture()
-def images() -> SpatialData:
+@pytest.fixture(scope="session")
+def _images_session() -> SpatialData:
  return SpatialData(images=_get_images())
 
 
 @pytest.fixture()
-def labels() -> SpatialData:
+def images(_images_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_images_session)
+
+
+@pytest.fixture(scope="session")
+def _labels_session() -> SpatialData:
  return SpatialData(labels=_get_labels())
 
 
+@pytest.fixture()
+def labels(_labels_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_labels_session)
+
+
 @pytest.fixture()
 def shapes() -> SpatialData:
  return SpatialData(shapes=_get_shapes())
@@ -87,8 +123,8 @@ def tables() -> list[AnnData]:
  return _tables
 
 
-@pytest.fixture()
-def full_sdata() -> SpatialData:
+@pytest.fixture(scope="session")
+def _full_sdata_session() -> SpatialData:
  return SpatialData(
  images=_get_images(),
  labels=_get_labels(),
@@ -98,6 +134,22 @@ def full_sdata() -> SpatialData:
  )
 
 
+@pytest.fixture()
+def full_sdata(_full_sdata_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_full_sdata_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_full_session() -> SpatialData:
+ return SpatialData(
+ images=_get_images(),
+ labels=_get_labels(),
+ shapes=_get_shapes(),
+ points=_get_points(),
+ tables=_get_tables(region="labels2d"),
+ )
+
+
 @pytest.fixture(
  # params=["labels"]
  params=["full", "empty"]
@@ -110,15 +162,9 @@ def full_sdata() -> SpatialData:
  ]
  # + ["empty_" + x for x in ["table"]] # TODO: empty table not supported yet
 )
-def sdata(request) -> SpatialData:
+def sdata(request, _sdata_full_session: SpatialData) -> SpatialData:
  if request.param == "full":
- return SpatialData(
- images=_get_images(),
- labels=_get_labels(),
- shapes=_get_shapes(),
- points=_get_points(),
- tables=_get_tables(region="labels2d"),
- )
+ return _fast_deepcopy_sdata(_sdata_full_session)
  if request.param == "empty":
  return SpatialData()
  return request.getfixturevalue(request.param)
@@ -304,18 +350,38 @@ def _get_new_table(spatial_element: None | str | Sequence[str], instance_id: Non
  return TableModel.parse(adata=adata, spatial_element=spatial_element, instance_id=instance_id)
 
 
-@pytest.fixture()
-def labels_blobs() -> ArrayLike:
- """Create a 2D labels."""
+@pytest.fixture(scope="session")
+def _labels_blobs_session() -> ArrayLike:
  return BlobsDataset()._labels_blobs()
 
 
 @pytest.fixture()
-def sdata_blobs() -> SpatialData:
+def labels_blobs(_labels_blobs_session: ArrayLike) -> ArrayLike:
  """Create a 2D labels."""
+ return deepcopy(_labels_blobs_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_blobs_session() -> SpatialData:
  from spatialdata.datasets import blobs
 
- return deepcopy(blobs(256, 300, 3))
+ return blobs(256, 300, 3)
+
+
+@pytest.fixture()
+def sdata_blobs(_sdata_blobs_session: SpatialData) -> SpatialData:
+ """Create a 2D labels."""
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+
+@pytest.fixture()
+def blobs_factory(_sdata_blobs_session: SpatialData) -> Callable[[], SpatialData]:
+ """Return a factory that creates cheap fresh copies of the session-scoped blobs dataset."""
+
+ def _make() -> SpatialData:
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+ return _make
 
 
 def _make_points(coordinates: np.ndarray) -> DaskDataFrame:

tests/core/operations/conftest.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import os
+
+# Disable numba JIT compilation for rasterize tests. Datashader (used by rasterize) triggers
+# numba JIT on first call, costing ~1.4s per worker. Python-mode gives identical results for
+# the small test data here — unlike real data, there is no throughput advantage from JIT.
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")

.github/workflows/test.yaml

@@ -20,13 +20,10 @@ jobs:
  matrix:
  include:
  - {os: windows-latest, python: "3.11", dask-version: "2025.12.0", name: "min dask"}
- - {os: windows-latest, python: "3.13", dask-version: "latest"}
  - {os: windows-latest, python: "3.14", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.11", dask-version: "latest"}
- - {os: ubuntu-latest, python: "3.13", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.14", dask-version: "latest"}
  - {os: macos-latest, python: "3.11", dask-version: "latest"}
- - {os: macos-latest, python: "3.13", prerelease: "allow", name: "prerelease"}
  - {os: macos-latest, python: "3.14", prerelease: "allow", name: "prerelease"}
  env:
  OS: ${{ matrix.os }}
@@ -62,7 +59,7 @@ jobs:
  PLATFORM: ${{ matrix.os }}
  DISPLAY: :42
  run: |
- uv run pytest --cov --color=yes --cov-report=xml
+ uv run pytest --cov --color=yes --cov-report=xml -n auto --dist worksteal
  - name: Upload coverage to Codecov
  uses: codecov/codecov-action@v5
  with:

pyproject.toml

@@ -69,6 +69,7 @@ test = [
  "pytest",
  "pytest-cov",
  "pytest-mock",
+ "pytest-xdist",
  "torch",
 ]
 docs = [
@@ -102,6 +103,7 @@ strict = true
 addopts = [
  "--import-mode=importlib", # allow using test files with same name
  "-s", # print output from tests
+ "-p no:napari", # napari registers a pytest plugin via its entry point; disable it here since spatialdata tests don't need it
 ]
 # These are all markers coming from xarray, dask or anndata. Added here to silence warnings.
 markers = [

src/spatialdata/_core/_elements.py

@@ -21,7 +21,6 @@
  PointsModel,
  ShapesModel,
  TableModel,
- get_axes_names,
  get_model,
 )
 
@@ -50,6 +49,7 @@ def _check_key(key: str, element_keys: Iterable[str], shared_keys: set[str | Non
  raise KeyError(*e.args) from e
 
  def __setitem__(self, key: str, value: T) -> None:
+ # note that each __setitem__ in the subclasses calls get_model(), which performs data validation
  self._add_shared_key(key)
  super().__setitem__(key, value)
 
@@ -72,15 +72,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Image2DModel, Image3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 3:
- Image2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 4:
- Image3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 4.")
+ super().__setitem__(key, value)
 
 
 class Labels(Elements[DataArray | DataTree]):
@@ -89,15 +81,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Labels2DModel, Labels3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 2:
- Labels2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 3:
- Labels3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 3.")
+ super().__setitem__(key, value)
 
 
 class Shapes(Elements[GeoDataFrame]):
@@ -106,7 +90,6 @@ def __setitem__(self, key: str, value: GeoDataFrame) -> None:
  schema = get_model(value)
  if schema != ShapesModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ShapesModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -116,7 +99,6 @@ def __setitem__(self, key: str, value: DaskDataFrame) -> None:
  schema = get_model(value)
  if schema != PointsModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- PointsModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -126,5 +108,4 @@ def __setitem__(self, key: str, value: AnnData) -> None:
  schema = get_model(value)
  if schema != TableModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- TableModel.validate(value)
  super().__setitem__(key, value)

src/spatialdata/_core/centroids.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from functools import singledispatch
 
-import dask.array as da
+import numpy as np
 import pandas as pd
 import xarray as xr
 from dask.dataframe import DataFrame as DaskDataFrame
@@ -56,40 +55,29 @@ def get_centroids(
  raise ValueError(f"The object type {type(e)} is not supported.")
 
 
-def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
+def _get_centroids_for_labels(xdata: xr.DataArray) -> pd.DataFrame:
  """
- Compute the component "axis" of the centroid of each label as a weighted average of the xarray coordinates.
+ Compute centroids for all labels in a DataArray in a single O(n_voxels) pass.
 
- Parameters
- ----------
- xdata
- The xarray DataArray containing the labels.
- axis
- The axis for which the centroids are computed.
-
- Returns
- -------
- pd.DataFrame
- A DataFrame containing one column, named after "axis", with the centroids of the labels along that axis.
- The index of the DataFrame is the collection of label values, sorted in ascending order.
+ Works for any number of spatial dimensions (2D and 3D labels).
  """
-centroids: dict[int, float] = defaultdict(float)
-for i in xdata[axis]:
- portion = xdata.sel(**{axis: i}).data
- u = da.unique(portion, return_counts=True)
- labels_values = u[0].compute()
- counts = u[1].compute()
- for j in range(len(labels_values)):
- label_value = labels_values[j]
- count = counts[j]
- centroids[label_value] += count * i.values.item()
-
-all_labels_values, all_labels_counts = da.unique(xdata.data, return_counts=True)
-all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute(), strict=True))
-for label_value in centroids:
-centroids[label_value] /= all_labels[label_value]
- centroids = dict(sorted(centroids.items(), key=lambda x: x[0]))
- return pd.DataFrame({axis: centroids.values()}, index=list(centroids.keys()))
+arr = xdata.data.compute()
+axes = list(xdata.dims)
+
+# Map label values to a contiguous range for bincount efficiency.
+label_ids, inverse = np.unique(arr, return_inverse=True)
+flat_inverse = inverse.ravel()
+counts = np.bincount(flat_inverse) # per-label pixel counts
+
+# indexing="ij" (matrix convention) ensures the i-th grid varies along the i-th
+# dimension of the output, correctly aligning with xdata.dims for any number of axes.
+ coord_grids = np.meshgrid(*[xdata[ax].values for ax in axes], indexing="ij")
+data: dict[str, np.ndarray] = {}
+for ax, grid in zip(axes, coord_grids, strict=True):
+ coord_sums = np.bincount(flat_inverse, weights=grid.ravel().astype(float))
+data[ax] = coord_sums / counts # counts > 0 by construction (unique guarantees this)
+
+ return pd.DataFrame(data, index=label_ids)
 
 
 @get_centroids.register(DataArray)
@@ -109,10 +97,7 @@ def _(
  assert len(e["scale0"]) == 1
  e = next(iter(e["scale0"].values()))
 
- dfs = []
- for axis in get_axes_names(e):
- dfs.append(_get_centroids_for_axis(e, axis))
- df = pd.concat(dfs, axis=1)
+ df = _get_centroids_for_labels(e)
  if not return_background and 0 in df.index:
  df = df.drop(index=0) # drop the background label
  t = get_transformation(e, coordinate_system)

src/spatialdata/models/models.py

@@ -1250,14 +1250,17 @@ def parse(
 
 def get_model(
  e: SpatialElement,
+ validate: bool = True,
 ) -> Schema_t:
  """
- Get the model for the given element.
+ Get the model for the given element. Validate using the model if `validate` is `True`.
 
  Parameters
  ----------
  e
  The element.
+ validate
+ Whether to validate the element using the model.
 
  Returns
  -------
@@ -1268,7 +1271,8 @@ def _validate_and_return(
  schema: Schema_t,
  e: SpatialElement,
  ) -> Schema_t:
- schema.validate(e)
+ if validate:
+ schema.validate(e)
  return schema
 
  if isinstance(e, DataArray | DataTree):

tests/conftest.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+
+# Disable numba JIT for the test suite (the test data is small so initializing the JIT is slower than using plain
+# Python)
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
+
+import copy as _copy
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any
 
@@ -33,6 +40,25 @@
  TableModel,
 )
 
+
+def _fast_deepcopy_sdata(sd: SpatialData) -> SpatialData:
+ """
+ Fast deepcopy for SpatialData objects in tests.
+
+ Uses copy.deepcopy (which skips model re-validation) and manually restores
+ the attrs that copy.deepcopy loses for DaskDataFrame (issue #503) and
+ GeoDataFrame (issue #286).
+ """
+ points_attrs = {k: _copy.deepcopy(v._attrs) for k, v in sd.points.items()}
+ shapes_attrs = {k: _copy.deepcopy(v.attrs) for k, v in sd.shapes.items()}
+ sd_copy = _copy.deepcopy(sd)
+ for k, attrs in points_attrs.items():
+ sd_copy.points[k]._attrs = attrs
+ for k, attrs in shapes_attrs.items():
+ sd_copy.shapes[k].attrs = attrs
+ return sd_copy
+
+
 SEED = 0
 RNG = default_rng(seed=SEED)
 
@@ -41,16 +67,26 @@
 POINT_PATH = Path(__file__).parent / "data/points.json"
 
 
-@pytest.fixture()
-def images() -> SpatialData:
+@pytest.fixture(scope="session")
+def _images_session() -> SpatialData:
  return SpatialData(images=_get_images())
 
 
 @pytest.fixture()
-def labels() -> SpatialData:
+def images(_images_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_images_session)
+
+
+@pytest.fixture(scope="session")
+def _labels_session() -> SpatialData:
  return SpatialData(labels=_get_labels())
 
 
+@pytest.fixture()
+def labels(_labels_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_labels_session)
+
+
 @pytest.fixture()
 def shapes() -> SpatialData:
  return SpatialData(shapes=_get_shapes())
@@ -87,8 +123,8 @@ def tables() -> list[AnnData]:
  return _tables
 
 
-@pytest.fixture()
-def full_sdata() -> SpatialData:
+@pytest.fixture(scope="session")
+def _full_sdata_session() -> SpatialData:
  return SpatialData(
  images=_get_images(),
  labels=_get_labels(),
@@ -98,6 +134,22 @@ def full_sdata() -> SpatialData:
  )
 
 
+@pytest.fixture()
+def full_sdata(_full_sdata_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_full_sdata_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_full_session() -> SpatialData:
+ return SpatialData(
+ images=_get_images(),
+ labels=_get_labels(),
+ shapes=_get_shapes(),
+ points=_get_points(),
+ tables=_get_tables(region="labels2d"),
+ )
+
+
 @pytest.fixture(
  # params=["labels"]
  params=["full", "empty"]
@@ -110,15 +162,9 @@ def full_sdata() -> SpatialData:
  ]
  # + ["empty_" + x for x in ["table"]] # TODO: empty table not supported yet
 )
-def sdata(request) -> SpatialData:
+def sdata(request, _sdata_full_session: SpatialData) -> SpatialData:
  if request.param == "full":
- return SpatialData(
- images=_get_images(),
- labels=_get_labels(),
- shapes=_get_shapes(),
- points=_get_points(),
- tables=_get_tables(region="labels2d"),
- )
+ return _fast_deepcopy_sdata(_sdata_full_session)
  if request.param == "empty":
  return SpatialData()
  return request.getfixturevalue(request.param)
@@ -304,18 +350,38 @@ def _get_new_table(spatial_element: None | str | Sequence[str], instance_id: Non
  return TableModel.parse(adata=adata, spatial_element=spatial_element, instance_id=instance_id)
 
 
-@pytest.fixture()
-def labels_blobs() -> ArrayLike:
- """Create a 2D labels."""
+@pytest.fixture(scope="session")
+def _labels_blobs_session() -> ArrayLike:
  return BlobsDataset()._labels_blobs()
 
 
 @pytest.fixture()
-def sdata_blobs() -> SpatialData:
+def labels_blobs(_labels_blobs_session: ArrayLike) -> ArrayLike:
  """Create a 2D labels."""
+ return deepcopy(_labels_blobs_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_blobs_session() -> SpatialData:
  from spatialdata.datasets import blobs
 
- return deepcopy(blobs(256, 300, 3))
+ return blobs(256, 300, 3)
+
+
+@pytest.fixture()
+def sdata_blobs(_sdata_blobs_session: SpatialData) -> SpatialData:
+ """Create a 2D labels."""
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+
+@pytest.fixture()
+def blobs_factory(_sdata_blobs_session: SpatialData) -> Callable[[], SpatialData]:
+ """Return a factory that creates cheap fresh copies of the session-scoped blobs dataset."""
+
+ def _make() -> SpatialData:
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+ return _make
 
 
 def _make_points(coordinates: np.ndarray) -> DaskDataFrame:

tests/core/operations/conftest.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import os
+
+# Disable numba JIT compilation for rasterize tests. Datashader (used by rasterize) triggers
+# numba JIT on first call, costing ~1.4s per worker. Python-mode gives identical results for
+# the small test data here — unlike real data, there is no throughput advantage from JIT.
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")

.github/workflows/test.yaml

@@ -20,13 +20,10 @@ jobs:
  matrix:
  include:
  - {os: windows-latest, python: "3.11", dask-version: "2025.12.0", name: "min dask"}
- - {os: windows-latest, python: "3.13", dask-version: "latest"}
  - {os: windows-latest, python: "3.14", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.11", dask-version: "latest"}
- - {os: ubuntu-latest, python: "3.13", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.14", dask-version: "latest"}
  - {os: macos-latest, python: "3.11", dask-version: "latest"}
- - {os: macos-latest, python: "3.13", prerelease: "allow", name: "prerelease"}
  - {os: macos-latest, python: "3.14", prerelease: "allow", name: "prerelease"}
  env:
  OS: ${{ matrix.os }}
@@ -62,7 +59,7 @@ jobs:
  PLATFORM: ${{ matrix.os }}
  DISPLAY: :42
  run: |
- uv run pytest --cov --color=yes --cov-report=xml
+ uv run pytest --cov --color=yes --cov-report=xml -n auto --dist worksteal
  - name: Upload coverage to Codecov
  uses: codecov/codecov-action@v5
  with:

pyproject.toml

@@ -69,6 +69,7 @@ test = [
  "pytest",
  "pytest-cov",
  "pytest-mock",
+ "pytest-xdist",
  "torch",
 ]
 docs = [
@@ -102,6 +103,7 @@ strict = true
 addopts = [
  "--import-mode=importlib", # allow using test files with same name
  "-s", # print output from tests
+ "-p no:napari", # napari registers a pytest plugin via its entry point; disable it here since spatialdata tests don't need it
 ]
 # These are all markers coming from xarray, dask or anndata. Added here to silence warnings.
 markers = [

src/spatialdata/_core/_elements.py

@@ -21,7 +21,6 @@
  PointsModel,
  ShapesModel,
  TableModel,
- get_axes_names,
  get_model,
 )
 
@@ -50,6 +49,7 @@ def _check_key(key: str, element_keys: Iterable[str], shared_keys: set[str | Non
  raise KeyError(*e.args) from e
 
  def __setitem__(self, key: str, value: T) -> None:
+ # note that each __setitem__ in the subclasses calls get_model(), which performs data validation
  self._add_shared_key(key)
  super().__setitem__(key, value)
 
@@ -72,15 +72,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Image2DModel, Image3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 3:
- Image2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 4:
- Image3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 4.")
+ super().__setitem__(key, value)
 
 
 class Labels(Elements[DataArray | DataTree]):
@@ -89,15 +81,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Labels2DModel, Labels3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 2:
- Labels2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 3:
- Labels3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 3.")
+ super().__setitem__(key, value)
 
 
 class Shapes(Elements[GeoDataFrame]):
@@ -106,7 +90,6 @@ def __setitem__(self, key: str, value: GeoDataFrame) -> None:
  schema = get_model(value)
  if schema != ShapesModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ShapesModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -116,7 +99,6 @@ def __setitem__(self, key: str, value: DaskDataFrame) -> None:
  schema = get_model(value)
  if schema != PointsModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- PointsModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -126,5 +108,4 @@ def __setitem__(self, key: str, value: AnnData) -> None:
  schema = get_model(value)
  if schema != TableModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- TableModel.validate(value)
  super().__setitem__(key, value)

src/spatialdata/_core/centroids.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from functools import singledispatch
 
-import dask.array as da
+import numpy as np
 import pandas as pd
 import xarray as xr
 from dask.dataframe import DataFrame as DaskDataFrame
@@ -56,40 +55,29 @@ def get_centroids(
  raise ValueError(f"The object type {type(e)} is not supported.")
 
 
-def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
+def _get_centroids_for_labels(xdata: xr.DataArray) -> pd.DataFrame:
  """
- Compute the component "axis" of the centroid of each label as a weighted average of the xarray coordinates.
+ Compute centroids for all labels in a DataArray in a single O(n_voxels) pass.
 
- Parameters
- ----------
- xdata
- The xarray DataArray containing the labels.
- axis
- The axis for which the centroids are computed.
-
- Returns
- -------
- pd.DataFrame
- A DataFrame containing one column, named after "axis", with the centroids of the labels along that axis.
- The index of the DataFrame is the collection of label values, sorted in ascending order.
+ Works for any number of spatial dimensions (2D and 3D labels).
  """
-centroids: dict[int, float] = defaultdict(float)
-for i in xdata[axis]:
- portion = xdata.sel(**{axis: i}).data
- u = da.unique(portion, return_counts=True)
- labels_values = u[0].compute()
- counts = u[1].compute()
- for j in range(len(labels_values)):
- label_value = labels_values[j]
- count = counts[j]
- centroids[label_value] += count * i.values.item()
-
-all_labels_values, all_labels_counts = da.unique(xdata.data, return_counts=True)
-all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute(), strict=True))
-for label_value in centroids:
-centroids[label_value] /= all_labels[label_value]
- centroids = dict(sorted(centroids.items(), key=lambda x: x[0]))
- return pd.DataFrame({axis: centroids.values()}, index=list(centroids.keys()))
+arr = xdata.data.compute()
+axes = list(xdata.dims)
+
+# Map label values to a contiguous range for bincount efficiency.
+label_ids, inverse = np.unique(arr, return_inverse=True)
+flat_inverse = inverse.ravel()
+counts = np.bincount(flat_inverse) # per-label pixel counts
+
+# indexing="ij" (matrix convention) ensures the i-th grid varies along the i-th
+# dimension of the output, correctly aligning with xdata.dims for any number of axes.
+ coord_grids = np.meshgrid(*[xdata[ax].values for ax in axes], indexing="ij")
+data: dict[str, np.ndarray] = {}
+for ax, grid in zip(axes, coord_grids, strict=True):
+ coord_sums = np.bincount(flat_inverse, weights=grid.ravel().astype(float))
+data[ax] = coord_sums / counts # counts > 0 by construction (unique guarantees this)
+
+ return pd.DataFrame(data, index=label_ids)
 
 
 @get_centroids.register(DataArray)
@@ -109,10 +97,7 @@ def _(
  assert len(e["scale0"]) == 1
  e = next(iter(e["scale0"].values()))
 
- dfs = []
- for axis in get_axes_names(e):
- dfs.append(_get_centroids_for_axis(e, axis))
- df = pd.concat(dfs, axis=1)
+ df = _get_centroids_for_labels(e)
  if not return_background and 0 in df.index:
  df = df.drop(index=0) # drop the background label
  t = get_transformation(e, coordinate_system)

src/spatialdata/models/models.py

@@ -1250,14 +1250,17 @@ def parse(
 
 def get_model(
  e: SpatialElement,
+ validate: bool = True,
 ) -> Schema_t:
  """
- Get the model for the given element.
+ Get the model for the given element. Validate using the model if `validate` is `True`.
 
  Parameters
  ----------
  e
  The element.
+ validate
+ Whether to validate the element using the model.
 
  Returns
  -------
@@ -1268,7 +1271,8 @@ def _validate_and_return(
  schema: Schema_t,
  e: SpatialElement,
  ) -> Schema_t:
- schema.validate(e)
+ if validate:
+ schema.validate(e)
  return schema
 
  if isinstance(e, DataArray | DataTree):

tests/conftest.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+
+# Disable numba JIT for the test suite (the test data is small so initializing the JIT is slower than using plain
+# Python)
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
+
+import copy as _copy
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any
 
@@ -33,6 +40,25 @@
  TableModel,
 )
 
+
+def _fast_deepcopy_sdata(sd: SpatialData) -> SpatialData:
+ """
+ Fast deepcopy for SpatialData objects in tests.
+
+ Uses copy.deepcopy (which skips model re-validation) and manually restores
+ the attrs that copy.deepcopy loses for DaskDataFrame (issue #503) and
+ GeoDataFrame (issue #286).
+ """
+ points_attrs = {k: _copy.deepcopy(v._attrs) for k, v in sd.points.items()}
+ shapes_attrs = {k: _copy.deepcopy(v.attrs) for k, v in sd.shapes.items()}
+ sd_copy = _copy.deepcopy(sd)
+ for k, attrs in points_attrs.items():
+ sd_copy.points[k]._attrs = attrs
+ for k, attrs in shapes_attrs.items():
+ sd_copy.shapes[k].attrs = attrs
+ return sd_copy
+
+
 SEED = 0
 RNG = default_rng(seed=SEED)
 
@@ -41,16 +67,26 @@
 POINT_PATH = Path(__file__).parent / "data/points.json"
 
 
-@pytest.fixture()
-def images() -> SpatialData:
+@pytest.fixture(scope="session")
+def _images_session() -> SpatialData:
  return SpatialData(images=_get_images())
 
 
 @pytest.fixture()
-def labels() -> SpatialData:
+def images(_images_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_images_session)
+
+
+@pytest.fixture(scope="session")
+def _labels_session() -> SpatialData:
  return SpatialData(labels=_get_labels())
 
 
+@pytest.fixture()
+def labels(_labels_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_labels_session)
+
+
 @pytest.fixture()
 def shapes() -> SpatialData:
  return SpatialData(shapes=_get_shapes())
@@ -87,8 +123,8 @@ def tables() -> list[AnnData]:
  return _tables
 
 
-@pytest.fixture()
-def full_sdata() -> SpatialData:
+@pytest.fixture(scope="session")
+def _full_sdata_session() -> SpatialData:
  return SpatialData(
  images=_get_images(),
  labels=_get_labels(),
@@ -98,6 +134,22 @@ def full_sdata() -> SpatialData:
  )
 
 
+@pytest.fixture()
+def full_sdata(_full_sdata_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_full_sdata_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_full_session() -> SpatialData:
+ return SpatialData(
+ images=_get_images(),
+ labels=_get_labels(),
+ shapes=_get_shapes(),
+ points=_get_points(),
+ tables=_get_tables(region="labels2d"),
+ )
+
+
 @pytest.fixture(
  # params=["labels"]
  params=["full", "empty"]
@@ -110,15 +162,9 @@ def full_sdata() -> SpatialData:
  ]
  # + ["empty_" + x for x in ["table"]] # TODO: empty table not supported yet
 )
-def sdata(request) -> SpatialData:
+def sdata(request, _sdata_full_session: SpatialData) -> SpatialData:
  if request.param == "full":
- return SpatialData(
- images=_get_images(),
- labels=_get_labels(),
- shapes=_get_shapes(),
- points=_get_points(),
- tables=_get_tables(region="labels2d"),
- )
+ return _fast_deepcopy_sdata(_sdata_full_session)
  if request.param == "empty":
  return SpatialData()
  return request.getfixturevalue(request.param)
@@ -304,18 +350,38 @@ def _get_new_table(spatial_element: None | str | Sequence[str], instance_id: Non
  return TableModel.parse(adata=adata, spatial_element=spatial_element, instance_id=instance_id)
 
 
-@pytest.fixture()
-def labels_blobs() -> ArrayLike:
- """Create a 2D labels."""
+@pytest.fixture(scope="session")
+def _labels_blobs_session() -> ArrayLike:
  return BlobsDataset()._labels_blobs()
 
 
 @pytest.fixture()
-def sdata_blobs() -> SpatialData:
+def labels_blobs(_labels_blobs_session: ArrayLike) -> ArrayLike:
  """Create a 2D labels."""
+ return deepcopy(_labels_blobs_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_blobs_session() -> SpatialData:
  from spatialdata.datasets import blobs
 
- return deepcopy(blobs(256, 300, 3))
+ return blobs(256, 300, 3)
+
+
+@pytest.fixture()
+def sdata_blobs(_sdata_blobs_session: SpatialData) -> SpatialData:
+ """Create a 2D labels."""
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+
+@pytest.fixture()
+def blobs_factory(_sdata_blobs_session: SpatialData) -> Callable[[], SpatialData]:
+ """Return a factory that creates cheap fresh copies of the session-scoped blobs dataset."""
+
+ def _make() -> SpatialData:
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+ return _make
 
 
 def _make_points(coordinates: np.ndarray) -> DaskDataFrame:

tests/core/operations/conftest.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import os
+
+# Disable numba JIT compilation for rasterize tests. Datashader (used by rasterize) triggers
+# numba JIT on first call, costing ~1.4s per worker. Python-mode gives identical results for
+# the small test data here — unlike real data, there is no throughput advantage from JIT.
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")

.github/workflows/test.yaml

@@ -20,13 +20,10 @@ jobs:
  matrix:
  include:
  - {os: windows-latest, python: "3.11", dask-version: "2025.12.0", name: "min dask"}
- - {os: windows-latest, python: "3.13", dask-version: "latest"}
  - {os: windows-latest, python: "3.14", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.11", dask-version: "latest"}
- - {os: ubuntu-latest, python: "3.13", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.14", dask-version: "latest"}
  - {os: macos-latest, python: "3.11", dask-version: "latest"}
- - {os: macos-latest, python: "3.13", prerelease: "allow", name: "prerelease"}
  - {os: macos-latest, python: "3.14", prerelease: "allow", name: "prerelease"}
  env:
  OS: ${{ matrix.os }}
@@ -62,7 +59,7 @@ jobs:
  PLATFORM: ${{ matrix.os }}
  DISPLAY: :42
  run: |
- uv run pytest --cov --color=yes --cov-report=xml
+ uv run pytest --cov --color=yes --cov-report=xml -n auto --dist worksteal
  - name: Upload coverage to Codecov
  uses: codecov/codecov-action@v5
  with:

pyproject.toml

@@ -69,6 +69,7 @@ test = [
  "pytest",
  "pytest-cov",
  "pytest-mock",
+ "pytest-xdist",
  "torch",
 ]
 docs = [
@@ -102,6 +103,7 @@ strict = true
 addopts = [
  "--import-mode=importlib", # allow using test files with same name
  "-s", # print output from tests
+ "-p no:napari", # napari registers a pytest plugin via its entry point; disable it here since spatialdata tests don't need it
 ]
 # These are all markers coming from xarray, dask or anndata. Added here to silence warnings.
 markers = [

src/spatialdata/_core/_elements.py

@@ -21,7 +21,6 @@
  PointsModel,
  ShapesModel,
  TableModel,
- get_axes_names,
  get_model,
 )
 
@@ -50,6 +49,7 @@ def _check_key(key: str, element_keys: Iterable[str], shared_keys: set[str | Non
  raise KeyError(*e.args) from e
 
  def __setitem__(self, key: str, value: T) -> None:
+ # note that each __setitem__ in the subclasses calls get_model(), which performs data validation
  self._add_shared_key(key)
  super().__setitem__(key, value)
 
@@ -72,15 +72,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Image2DModel, Image3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 3:
- Image2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 4:
- Image3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 4.")
+ super().__setitem__(key, value)
 
 
 class Labels(Elements[DataArray | DataTree]):
@@ -89,15 +81,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Labels2DModel, Labels3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 2:
- Labels2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 3:
- Labels3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 3.")
+ super().__setitem__(key, value)
 
 
 class Shapes(Elements[GeoDataFrame]):
@@ -106,7 +90,6 @@ def __setitem__(self, key: str, value: GeoDataFrame) -> None:
  schema = get_model(value)
  if schema != ShapesModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ShapesModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -116,7 +99,6 @@ def __setitem__(self, key: str, value: DaskDataFrame) -> None:
  schema = get_model(value)
  if schema != PointsModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- PointsModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -126,5 +108,4 @@ def __setitem__(self, key: str, value: AnnData) -> None:
  schema = get_model(value)
  if schema != TableModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- TableModel.validate(value)
  super().__setitem__(key, value)

src/spatialdata/_core/centroids.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from functools import singledispatch
 
-import dask.array as da
+import numpy as np
 import pandas as pd
 import xarray as xr
 from dask.dataframe import DataFrame as DaskDataFrame
@@ -56,40 +55,29 @@ def get_centroids(
  raise ValueError(f"The object type {type(e)} is not supported.")
 
 
-def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
+def _get_centroids_for_labels(xdata: xr.DataArray) -> pd.DataFrame:
  """
- Compute the component "axis" of the centroid of each label as a weighted average of the xarray coordinates.
+ Compute centroids for all labels in a DataArray in a single O(n_voxels) pass.
 
- Parameters
- ----------
- xdata
- The xarray DataArray containing the labels.
- axis
- The axis for which the centroids are computed.
-
- Returns
- -------
- pd.DataFrame
- A DataFrame containing one column, named after "axis", with the centroids of the labels along that axis.
- The index of the DataFrame is the collection of label values, sorted in ascending order.
+ Works for any number of spatial dimensions (2D and 3D labels).
  """
-centroids: dict[int, float] = defaultdict(float)
-for i in xdata[axis]:
- portion = xdata.sel(**{axis: i}).data
- u = da.unique(portion, return_counts=True)
- labels_values = u[0].compute()
- counts = u[1].compute()
- for j in range(len(labels_values)):
- label_value = labels_values[j]
- count = counts[j]
- centroids[label_value] += count * i.values.item()
-
-all_labels_values, all_labels_counts = da.unique(xdata.data, return_counts=True)
-all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute(), strict=True))
-for label_value in centroids:
-centroids[label_value] /= all_labels[label_value]
- centroids = dict(sorted(centroids.items(), key=lambda x: x[0]))
- return pd.DataFrame({axis: centroids.values()}, index=list(centroids.keys()))
+arr = xdata.data.compute()
+axes = list(xdata.dims)
+
+# Map label values to a contiguous range for bincount efficiency.
+label_ids, inverse = np.unique(arr, return_inverse=True)
+flat_inverse = inverse.ravel()
+counts = np.bincount(flat_inverse) # per-label pixel counts
+
+# indexing="ij" (matrix convention) ensures the i-th grid varies along the i-th
+# dimension of the output, correctly aligning with xdata.dims for any number of axes.
+ coord_grids = np.meshgrid(*[xdata[ax].values for ax in axes], indexing="ij")
+data: dict[str, np.ndarray] = {}
+for ax, grid in zip(axes, coord_grids, strict=True):
+ coord_sums = np.bincount(flat_inverse, weights=grid.ravel().astype(float))
+data[ax] = coord_sums / counts # counts > 0 by construction (unique guarantees this)
+
+ return pd.DataFrame(data, index=label_ids)
 
 
 @get_centroids.register(DataArray)
@@ -109,10 +97,7 @@ def _(
  assert len(e["scale0"]) == 1
  e = next(iter(e["scale0"].values()))
 
- dfs = []
- for axis in get_axes_names(e):
- dfs.append(_get_centroids_for_axis(e, axis))
- df = pd.concat(dfs, axis=1)
+ df = _get_centroids_for_labels(e)
  if not return_background and 0 in df.index:
  df = df.drop(index=0) # drop the background label
  t = get_transformation(e, coordinate_system)

src/spatialdata/models/models.py

@@ -1250,14 +1250,17 @@ def parse(
 
 def get_model(
  e: SpatialElement,
+ validate: bool = True,
 ) -> Schema_t:
  """
- Get the model for the given element.
+ Get the model for the given element. Validate using the model if `validate` is `True`.
 
  Parameters
  ----------
  e
  The element.
+ validate
+ Whether to validate the element using the model.
 
  Returns
  -------
@@ -1268,7 +1271,8 @@ def _validate_and_return(
  schema: Schema_t,
  e: SpatialElement,
  ) -> Schema_t:
- schema.validate(e)
+ if validate:
+ schema.validate(e)
  return schema
 
  if isinstance(e, DataArray | DataTree):

tests/conftest.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+
+# Disable numba JIT for the test suite (the test data is small so initializing the JIT is slower than using plain
+# Python)
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
+
+import copy as _copy
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any
 
@@ -33,6 +40,25 @@
  TableModel,
 )
 
+
+def _fast_deepcopy_sdata(sd: SpatialData) -> SpatialData:
+ """
+ Fast deepcopy for SpatialData objects in tests.
+
+ Uses copy.deepcopy (which skips model re-validation) and manually restores
+ the attrs that copy.deepcopy loses for DaskDataFrame (issue #503) and
+ GeoDataFrame (issue #286).
+ """
+ points_attrs = {k: _copy.deepcopy(v._attrs) for k, v in sd.points.items()}
+ shapes_attrs = {k: _copy.deepcopy(v.attrs) for k, v in sd.shapes.items()}
+ sd_copy = _copy.deepcopy(sd)
+ for k, attrs in points_attrs.items():
+ sd_copy.points[k]._attrs = attrs
+ for k, attrs in shapes_attrs.items():
+ sd_copy.shapes[k].attrs = attrs
+ return sd_copy
+
+
 SEED = 0
 RNG = default_rng(seed=SEED)
 
@@ -41,16 +67,26 @@
 POINT_PATH = Path(__file__).parent / "data/points.json"
 
 
-@pytest.fixture()
-def images() -> SpatialData:
+@pytest.fixture(scope="session")
+def _images_session() -> SpatialData:
  return SpatialData(images=_get_images())
 
 
 @pytest.fixture()
-def labels() -> SpatialData:
+def images(_images_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_images_session)
+
+
+@pytest.fixture(scope="session")
+def _labels_session() -> SpatialData:
  return SpatialData(labels=_get_labels())
 
 
+@pytest.fixture()
+def labels(_labels_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_labels_session)
+
+
 @pytest.fixture()
 def shapes() -> SpatialData:
  return SpatialData(shapes=_get_shapes())
@@ -87,8 +123,8 @@ def tables() -> list[AnnData]:
  return _tables
 
 
-@pytest.fixture()
-def full_sdata() -> SpatialData:
+@pytest.fixture(scope="session")
+def _full_sdata_session() -> SpatialData:
  return SpatialData(
  images=_get_images(),
  labels=_get_labels(),
@@ -98,6 +134,22 @@ def full_sdata() -> SpatialData:
  )
 
 
+@pytest.fixture()
+def full_sdata(_full_sdata_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_full_sdata_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_full_session() -> SpatialData:
+ return SpatialData(
+ images=_get_images(),
+ labels=_get_labels(),
+ shapes=_get_shapes(),
+ points=_get_points(),
+ tables=_get_tables(region="labels2d"),
+ )
+
+
 @pytest.fixture(
  # params=["labels"]
  params=["full", "empty"]
@@ -110,15 +162,9 @@ def full_sdata() -> SpatialData:
  ]
  # + ["empty_" + x for x in ["table"]] # TODO: empty table not supported yet
 )
-def sdata(request) -> SpatialData:
+def sdata(request, _sdata_full_session: SpatialData) -> SpatialData:
  if request.param == "full":
- return SpatialData(
- images=_get_images(),
- labels=_get_labels(),
- shapes=_get_shapes(),
- points=_get_points(),
- tables=_get_tables(region="labels2d"),
- )
+ return _fast_deepcopy_sdata(_sdata_full_session)
  if request.param == "empty":
  return SpatialData()
  return request.getfixturevalue(request.param)
@@ -304,18 +350,38 @@ def _get_new_table(spatial_element: None | str | Sequence[str], instance_id: Non
  return TableModel.parse(adata=adata, spatial_element=spatial_element, instance_id=instance_id)
 
 
-@pytest.fixture()
-def labels_blobs() -> ArrayLike:
- """Create a 2D labels."""
+@pytest.fixture(scope="session")
+def _labels_blobs_session() -> ArrayLike:
  return BlobsDataset()._labels_blobs()
 
 
 @pytest.fixture()
-def sdata_blobs() -> SpatialData:
+def labels_blobs(_labels_blobs_session: ArrayLike) -> ArrayLike:
  """Create a 2D labels."""
+ return deepcopy(_labels_blobs_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_blobs_session() -> SpatialData:
  from spatialdata.datasets import blobs
 
- return deepcopy(blobs(256, 300, 3))
+ return blobs(256, 300, 3)
+
+
+@pytest.fixture()
+def sdata_blobs(_sdata_blobs_session: SpatialData) -> SpatialData:
+ """Create a 2D labels."""
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+
+@pytest.fixture()
+def blobs_factory(_sdata_blobs_session: SpatialData) -> Callable[[], SpatialData]:
+ """Return a factory that creates cheap fresh copies of the session-scoped blobs dataset."""
+
+ def _make() -> SpatialData:
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+ return _make
 
 
 def _make_points(coordinates: np.ndarray) -> DaskDataFrame:

tests/core/operations/conftest.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import os
+
+# Disable numba JIT compilation for rasterize tests. Datashader (used by rasterize) triggers
+# numba JIT on first call, costing ~1.4s per worker. Python-mode gives identical results for
+# the small test data here — unlike real data, there is no throughput advantage from JIT.
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")

.github/workflows/test.yaml

@@ -20,13 +20,10 @@ jobs:
  matrix:
  include:
  - {os: windows-latest, python: "3.11", dask-version: "2025.12.0", name: "min dask"}
- - {os: windows-latest, python: "3.13", dask-version: "latest"}
  - {os: windows-latest, python: "3.14", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.11", dask-version: "latest"}
- - {os: ubuntu-latest, python: "3.13", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.14", dask-version: "latest"}
  - {os: macos-latest, python: "3.11", dask-version: "latest"}
- - {os: macos-latest, python: "3.13", prerelease: "allow", name: "prerelease"}
  - {os: macos-latest, python: "3.14", prerelease: "allow", name: "prerelease"}
  env:
  OS: ${{ matrix.os }}
@@ -62,7 +59,7 @@ jobs:
  PLATFORM: ${{ matrix.os }}
  DISPLAY: :42
  run: |
- uv run pytest --cov --color=yes --cov-report=xml
+ uv run pytest --cov --color=yes --cov-report=xml -n auto --dist worksteal
  - name: Upload coverage to Codecov
  uses: codecov/codecov-action@v5
  with:

pyproject.toml

@@ -69,6 +69,7 @@ test = [
  "pytest",
  "pytest-cov",
  "pytest-mock",
+ "pytest-xdist",
  "torch",
 ]
 docs = [
@@ -102,6 +103,7 @@ strict = true
 addopts = [
  "--import-mode=importlib", # allow using test files with same name
  "-s", # print output from tests
+ "-p no:napari", # napari registers a pytest plugin via its entry point; disable it here since spatialdata tests don't need it
 ]
 # These are all markers coming from xarray, dask or anndata. Added here to silence warnings.
 markers = [

src/spatialdata/_core/_elements.py

@@ -21,7 +21,6 @@
  PointsModel,
  ShapesModel,
  TableModel,
- get_axes_names,
  get_model,
 )
 
@@ -50,6 +49,7 @@ def _check_key(key: str, element_keys: Iterable[str], shared_keys: set[str | Non
  raise KeyError(*e.args) from e
 
  def __setitem__(self, key: str, value: T) -> None:
+ # note that each __setitem__ in the subclasses calls get_model(), which performs data validation
  self._add_shared_key(key)
  super().__setitem__(key, value)
 
@@ -72,15 +72,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Image2DModel, Image3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 3:
- Image2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 4:
- Image3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 4.")
+ super().__setitem__(key, value)
 
 
 class Labels(Elements[DataArray | DataTree]):
@@ -89,15 +81,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Labels2DModel, Labels3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 2:
- Labels2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 3:
- Labels3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 3.")
+ super().__setitem__(key, value)
 
 
 class Shapes(Elements[GeoDataFrame]):
@@ -106,7 +90,6 @@ def __setitem__(self, key: str, value: GeoDataFrame) -> None:
  schema = get_model(value)
  if schema != ShapesModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ShapesModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -116,7 +99,6 @@ def __setitem__(self, key: str, value: DaskDataFrame) -> None:
  schema = get_model(value)
  if schema != PointsModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- PointsModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -126,5 +108,4 @@ def __setitem__(self, key: str, value: AnnData) -> None:
  schema = get_model(value)
  if schema != TableModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- TableModel.validate(value)
  super().__setitem__(key, value)

src/spatialdata/_core/centroids.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from functools import singledispatch
 
-import dask.array as da
+import numpy as np
 import pandas as pd
 import xarray as xr
 from dask.dataframe import DataFrame as DaskDataFrame
@@ -56,40 +55,29 @@ def get_centroids(
  raise ValueError(f"The object type {type(e)} is not supported.")
 
 
-def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
+def _get_centroids_for_labels(xdata: xr.DataArray) -> pd.DataFrame:
  """
- Compute the component "axis" of the centroid of each label as a weighted average of the xarray coordinates.
+ Compute centroids for all labels in a DataArray in a single O(n_voxels) pass.
 
- Parameters
- ----------
- xdata
- The xarray DataArray containing the labels.
- axis
- The axis for which the centroids are computed.
-
- Returns
- -------
- pd.DataFrame
- A DataFrame containing one column, named after "axis", with the centroids of the labels along that axis.
- The index of the DataFrame is the collection of label values, sorted in ascending order.
+ Works for any number of spatial dimensions (2D and 3D labels).
  """
-centroids: dict[int, float] = defaultdict(float)
-for i in xdata[axis]:
- portion = xdata.sel(**{axis: i}).data
- u = da.unique(portion, return_counts=True)
- labels_values = u[0].compute()
- counts = u[1].compute()
- for j in range(len(labels_values)):
- label_value = labels_values[j]
- count = counts[j]
- centroids[label_value] += count * i.values.item()
-
-all_labels_values, all_labels_counts = da.unique(xdata.data, return_counts=True)
-all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute(), strict=True))
-for label_value in centroids:
-centroids[label_value] /= all_labels[label_value]
- centroids = dict(sorted(centroids.items(), key=lambda x: x[0]))
- return pd.DataFrame({axis: centroids.values()}, index=list(centroids.keys()))
+arr = xdata.data.compute()
+axes = list(xdata.dims)
+
+# Map label values to a contiguous range for bincount efficiency.
+label_ids, inverse = np.unique(arr, return_inverse=True)
+flat_inverse = inverse.ravel()
+counts = np.bincount(flat_inverse) # per-label pixel counts
+
+# indexing="ij" (matrix convention) ensures the i-th grid varies along the i-th
+# dimension of the output, correctly aligning with xdata.dims for any number of axes.
+ coord_grids = np.meshgrid(*[xdata[ax].values for ax in axes], indexing="ij")
+data: dict[str, np.ndarray] = {}
+for ax, grid in zip(axes, coord_grids, strict=True):
+ coord_sums = np.bincount(flat_inverse, weights=grid.ravel().astype(float))
+data[ax] = coord_sums / counts # counts > 0 by construction (unique guarantees this)
+
+ return pd.DataFrame(data, index=label_ids)
 
 
 @get_centroids.register(DataArray)
@@ -109,10 +97,7 @@ def _(
  assert len(e["scale0"]) == 1
  e = next(iter(e["scale0"].values()))
 
- dfs = []
- for axis in get_axes_names(e):
- dfs.append(_get_centroids_for_axis(e, axis))
- df = pd.concat(dfs, axis=1)
+ df = _get_centroids_for_labels(e)
  if not return_background and 0 in df.index:
  df = df.drop(index=0) # drop the background label
  t = get_transformation(e, coordinate_system)

src/spatialdata/models/models.py

@@ -1250,14 +1250,17 @@ def parse(
 
 def get_model(
  e: SpatialElement,
+ validate: bool = True,
 ) -> Schema_t:
  """
- Get the model for the given element.
+ Get the model for the given element. Validate using the model if `validate` is `True`.
 
  Parameters
  ----------
  e
  The element.
+ validate
+ Whether to validate the element using the model.
 
  Returns
  -------
@@ -1268,7 +1271,8 @@ def _validate_and_return(
  schema: Schema_t,
  e: SpatialElement,
  ) -> Schema_t:
- schema.validate(e)
+ if validate:
+ schema.validate(e)
  return schema
 
  if isinstance(e, DataArray | DataTree):

tests/conftest.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+
+# Disable numba JIT for the test suite (the test data is small so initializing the JIT is slower than using plain
+# Python)
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
+
+import copy as _copy
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any
 
@@ -33,6 +40,25 @@
  TableModel,
 )
 
+
+def _fast_deepcopy_sdata(sd: SpatialData) -> SpatialData:
+ """
+ Fast deepcopy for SpatialData objects in tests.
+
+ Uses copy.deepcopy (which skips model re-validation) and manually restores
+ the attrs that copy.deepcopy loses for DaskDataFrame (issue #503) and
+ GeoDataFrame (issue #286).
+ """
+ points_attrs = {k: _copy.deepcopy(v._attrs) for k, v in sd.points.items()}
+ shapes_attrs = {k: _copy.deepcopy(v.attrs) for k, v in sd.shapes.items()}
+ sd_copy = _copy.deepcopy(sd)
+ for k, attrs in points_attrs.items():
+ sd_copy.points[k]._attrs = attrs
+ for k, attrs in shapes_attrs.items():
+ sd_copy.shapes[k].attrs = attrs
+ return sd_copy
+
+
 SEED = 0
 RNG = default_rng(seed=SEED)
 
@@ -41,16 +67,26 @@
 POINT_PATH = Path(__file__).parent / "data/points.json"
 
 
-@pytest.fixture()
-def images() -> SpatialData:
+@pytest.fixture(scope="session")
+def _images_session() -> SpatialData:
  return SpatialData(images=_get_images())
 
 
 @pytest.fixture()
-def labels() -> SpatialData:
+def images(_images_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_images_session)
+
+
+@pytest.fixture(scope="session")
+def _labels_session() -> SpatialData:
  return SpatialData(labels=_get_labels())
 
 
+@pytest.fixture()
+def labels(_labels_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_labels_session)
+
+
 @pytest.fixture()
 def shapes() -> SpatialData:
  return SpatialData(shapes=_get_shapes())
@@ -87,8 +123,8 @@ def tables() -> list[AnnData]:
  return _tables
 
 
-@pytest.fixture()
-def full_sdata() -> SpatialData:
+@pytest.fixture(scope="session")
+def _full_sdata_session() -> SpatialData:
  return SpatialData(
  images=_get_images(),
  labels=_get_labels(),
@@ -98,6 +134,22 @@ def full_sdata() -> SpatialData:
  )
 
 
+@pytest.fixture()
+def full_sdata(_full_sdata_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_full_sdata_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_full_session() -> SpatialData:
+ return SpatialData(
+ images=_get_images(),
+ labels=_get_labels(),
+ shapes=_get_shapes(),
+ points=_get_points(),
+ tables=_get_tables(region="labels2d"),
+ )
+
+
 @pytest.fixture(
  # params=["labels"]
  params=["full", "empty"]
@@ -110,15 +162,9 @@ def full_sdata() -> SpatialData:
  ]
  # + ["empty_" + x for x in ["table"]] # TODO: empty table not supported yet
 )
-def sdata(request) -> SpatialData:
+def sdata(request, _sdata_full_session: SpatialData) -> SpatialData:
  if request.param == "full":
- return SpatialData(
- images=_get_images(),
- labels=_get_labels(),
- shapes=_get_shapes(),
- points=_get_points(),
- tables=_get_tables(region="labels2d"),
- )
+ return _fast_deepcopy_sdata(_sdata_full_session)
  if request.param == "empty":
  return SpatialData()
  return request.getfixturevalue(request.param)
@@ -304,18 +350,38 @@ def _get_new_table(spatial_element: None | str | Sequence[str], instance_id: Non
  return TableModel.parse(adata=adata, spatial_element=spatial_element, instance_id=instance_id)
 
 
-@pytest.fixture()
-def labels_blobs() -> ArrayLike:
- """Create a 2D labels."""
+@pytest.fixture(scope="session")
+def _labels_blobs_session() -> ArrayLike:
  return BlobsDataset()._labels_blobs()
 
 
 @pytest.fixture()
-def sdata_blobs() -> SpatialData:
+def labels_blobs(_labels_blobs_session: ArrayLike) -> ArrayLike:
  """Create a 2D labels."""
+ return deepcopy(_labels_blobs_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_blobs_session() -> SpatialData:
  from spatialdata.datasets import blobs
 
- return deepcopy(blobs(256, 300, 3))
+ return blobs(256, 300, 3)
+
+
+@pytest.fixture()
+def sdata_blobs(_sdata_blobs_session: SpatialData) -> SpatialData:
+ """Create a 2D labels."""
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+
+@pytest.fixture()
+def blobs_factory(_sdata_blobs_session: SpatialData) -> Callable[[], SpatialData]:
+ """Return a factory that creates cheap fresh copies of the session-scoped blobs dataset."""
+
+ def _make() -> SpatialData:
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+ return _make
 
 
 def _make_points(coordinates: np.ndarray) -> DaskDataFrame:

tests/core/operations/conftest.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import os
+
+# Disable numba JIT compilation for rasterize tests. Datashader (used by rasterize) triggers
+# numba JIT on first call, costing ~1.4s per worker. Python-mode gives identical results for
+# the small test data here — unlike real data, there is no throughput advantage from JIT.
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")

.github/workflows/test.yaml

@@ -20,13 +20,10 @@ jobs:
  matrix:
  include:
  - {os: windows-latest, python: "3.11", dask-version: "2025.12.0", name: "min dask"}
- - {os: windows-latest, python: "3.13", dask-version: "latest"}
  - {os: windows-latest, python: "3.14", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.11", dask-version: "latest"}
- - {os: ubuntu-latest, python: "3.13", dask-version: "latest"}
  - {os: ubuntu-latest, python: "3.14", dask-version: "latest"}
  - {os: macos-latest, python: "3.11", dask-version: "latest"}
- - {os: macos-latest, python: "3.13", prerelease: "allow", name: "prerelease"}
  - {os: macos-latest, python: "3.14", prerelease: "allow", name: "prerelease"}
  env:
  OS: ${{ matrix.os }}
@@ -62,7 +59,7 @@ jobs:
  PLATFORM: ${{ matrix.os }}
  DISPLAY: :42
  run: |
- uv run pytest --cov --color=yes --cov-report=xml
+ uv run pytest --cov --color=yes --cov-report=xml -n auto --dist worksteal
  - name: Upload coverage to Codecov
  uses: codecov/codecov-action@v5
  with:

pyproject.toml

@@ -69,6 +69,7 @@ test = [
  "pytest",
  "pytest-cov",
  "pytest-mock",
+ "pytest-xdist",
  "torch",
 ]
 docs = [
@@ -102,6 +103,7 @@ strict = true
 addopts = [
  "--import-mode=importlib", # allow using test files with same name
  "-s", # print output from tests
+ "-p no:napari", # napari registers a pytest plugin via its entry point; disable it here since spatialdata tests don't need it
 ]
 # These are all markers coming from xarray, dask or anndata. Added here to silence warnings.
 markers = [

src/spatialdata/_core/_elements.py

@@ -21,7 +21,6 @@
  PointsModel,
  ShapesModel,
  TableModel,
- get_axes_names,
  get_model,
 )
 
@@ -50,6 +49,7 @@ def _check_key(key: str, element_keys: Iterable[str], shared_keys: set[str | Non
  raise KeyError(*e.args) from e
 
  def __setitem__(self, key: str, value: T) -> None:
+ # note that each __setitem__ in the subclasses calls get_model(), which performs data validation
  self._add_shared_key(key)
  super().__setitem__(key, value)
 
@@ -72,15 +72,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Image2DModel, Image3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 3:
- Image2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 4:
- Image3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 4.")
+ super().__setitem__(key, value)
 
 
 class Labels(Elements[DataArray | DataTree]):
@@ -89,15 +81,7 @@ def __setitem__(self, key: str, value: Raster_T) -> None:
  schema = get_model(value)
  if schema not in (Labels2DModel, Labels3DModel):
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ndim = len(get_axes_names(value))
- if ndim == 2:
- Labels2DModel.validate(value)
- super().__setitem__(key, value)
- elif ndim == 3:
- Labels3DModel.validate(value)
- super().__setitem__(key, value)
- else:
- NotImplementedError("TODO: implement for ndim > 3.")
+ super().__setitem__(key, value)
 
 
 class Shapes(Elements[GeoDataFrame]):
@@ -106,7 +90,6 @@ def __setitem__(self, key: str, value: GeoDataFrame) -> None:
  schema = get_model(value)
  if schema != ShapesModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- ShapesModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -116,7 +99,6 @@ def __setitem__(self, key: str, value: DaskDataFrame) -> None:
  schema = get_model(value)
  if schema != PointsModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- PointsModel.validate(value)
  super().__setitem__(key, value)
 
 
@@ -126,5 +108,4 @@ def __setitem__(self, key: str, value: AnnData) -> None:
  schema = get_model(value)
  if schema != TableModel:
  raise TypeError(f"Unknown element type with schema: {schema!r}.")
- TableModel.validate(value)
  super().__setitem__(key, value)

src/spatialdata/_core/centroids.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from functools import singledispatch
 
-import dask.array as da
+import numpy as np
 import pandas as pd
 import xarray as xr
 from dask.dataframe import DataFrame as DaskDataFrame
@@ -56,40 +55,29 @@ def get_centroids(
  raise ValueError(f"The object type {type(e)} is not supported.")
 
 
-def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
+def _get_centroids_for_labels(xdata: xr.DataArray) -> pd.DataFrame:
  """
- Compute the component "axis" of the centroid of each label as a weighted average of the xarray coordinates.
+ Compute centroids for all labels in a DataArray in a single O(n_voxels) pass.
 
- Parameters
- ----------
- xdata
- The xarray DataArray containing the labels.
- axis
- The axis for which the centroids are computed.
-
- Returns
- -------
- pd.DataFrame
- A DataFrame containing one column, named after "axis", with the centroids of the labels along that axis.
- The index of the DataFrame is the collection of label values, sorted in ascending order.
+ Works for any number of spatial dimensions (2D and 3D labels).
  """
-centroids: dict[int, float] = defaultdict(float)
-for i in xdata[axis]:
- portion = xdata.sel(**{axis: i}).data
- u = da.unique(portion, return_counts=True)
- labels_values = u[0].compute()
- counts = u[1].compute()
- for j in range(len(labels_values)):
- label_value = labels_values[j]
- count = counts[j]
- centroids[label_value] += count * i.values.item()
-
-all_labels_values, all_labels_counts = da.unique(xdata.data, return_counts=True)
-all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute(), strict=True))
-for label_value in centroids:
-centroids[label_value] /= all_labels[label_value]
- centroids = dict(sorted(centroids.items(), key=lambda x: x[0]))
- return pd.DataFrame({axis: centroids.values()}, index=list(centroids.keys()))
+arr = xdata.data.compute()
+axes = list(xdata.dims)
+
+# Map label values to a contiguous range for bincount efficiency.
+label_ids, inverse = np.unique(arr, return_inverse=True)
+flat_inverse = inverse.ravel()
+counts = np.bincount(flat_inverse) # per-label pixel counts
+
+# indexing="ij" (matrix convention) ensures the i-th grid varies along the i-th
+# dimension of the output, correctly aligning with xdata.dims for any number of axes.
+ coord_grids = np.meshgrid(*[xdata[ax].values for ax in axes], indexing="ij")
+data: dict[str, np.ndarray] = {}
+for ax, grid in zip(axes, coord_grids, strict=True):
+ coord_sums = np.bincount(flat_inverse, weights=grid.ravel().astype(float))
+data[ax] = coord_sums / counts # counts > 0 by construction (unique guarantees this)
+
+ return pd.DataFrame(data, index=label_ids)
 
 
 @get_centroids.register(DataArray)
@@ -109,10 +97,7 @@ def _(
  assert len(e["scale0"]) == 1
  e = next(iter(e["scale0"].values()))
 
- dfs = []
- for axis in get_axes_names(e):
- dfs.append(_get_centroids_for_axis(e, axis))
- df = pd.concat(dfs, axis=1)
+ df = _get_centroids_for_labels(e)
  if not return_background and 0 in df.index:
  df = df.drop(index=0) # drop the background label
  t = get_transformation(e, coordinate_system)

src/spatialdata/models/models.py

@@ -1250,14 +1250,17 @@ def parse(
 
 def get_model(
  e: SpatialElement,
+ validate: bool = True,
 ) -> Schema_t:
  """
- Get the model for the given element.
+ Get the model for the given element. Validate using the model if `validate` is `True`.
 
  Parameters
  ----------
  e
  The element.
+ validate
+ Whether to validate the element using the model.
 
  Returns
  -------
@@ -1268,7 +1271,8 @@ def _validate_and_return(
  schema: Schema_t,
  e: SpatialElement,
  ) -> Schema_t:
- schema.validate(e)
+ if validate:
+ schema.validate(e)
  return schema
 
  if isinstance(e, DataArray | DataTree):

tests/conftest.py

@@ -1,6 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+import os
+
+# Disable numba JIT for the test suite (the test data is small so initializing the JIT is slower than using plain
+# Python)
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")
+
+import copy as _copy
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any
 
@@ -33,6 +40,25 @@
  TableModel,
 )
 
+
+def _fast_deepcopy_sdata(sd: SpatialData) -> SpatialData:
+ """
+ Fast deepcopy for SpatialData objects in tests.
+
+ Uses copy.deepcopy (which skips model re-validation) and manually restores
+ the attrs that copy.deepcopy loses for DaskDataFrame (issue #503) and
+ GeoDataFrame (issue #286).
+ """
+ points_attrs = {k: _copy.deepcopy(v._attrs) for k, v in sd.points.items()}
+ shapes_attrs = {k: _copy.deepcopy(v.attrs) for k, v in sd.shapes.items()}
+ sd_copy = _copy.deepcopy(sd)
+ for k, attrs in points_attrs.items():
+ sd_copy.points[k]._attrs = attrs
+ for k, attrs in shapes_attrs.items():
+ sd_copy.shapes[k].attrs = attrs
+ return sd_copy
+
+
 SEED = 0
 RNG = default_rng(seed=SEED)
 
@@ -41,16 +67,26 @@
 POINT_PATH = Path(__file__).parent / "data/points.json"
 
 
-@pytest.fixture()
-def images() -> SpatialData:
+@pytest.fixture(scope="session")
+def _images_session() -> SpatialData:
  return SpatialData(images=_get_images())
 
 
 @pytest.fixture()
-def labels() -> SpatialData:
+def images(_images_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_images_session)
+
+
+@pytest.fixture(scope="session")
+def _labels_session() -> SpatialData:
  return SpatialData(labels=_get_labels())
 
 
+@pytest.fixture()
+def labels(_labels_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_labels_session)
+
+
 @pytest.fixture()
 def shapes() -> SpatialData:
  return SpatialData(shapes=_get_shapes())
@@ -87,8 +123,8 @@ def tables() -> list[AnnData]:
  return _tables
 
 
-@pytest.fixture()
-def full_sdata() -> SpatialData:
+@pytest.fixture(scope="session")
+def _full_sdata_session() -> SpatialData:
  return SpatialData(
  images=_get_images(),
  labels=_get_labels(),
@@ -98,6 +134,22 @@ def full_sdata() -> SpatialData:
  )
 
 
+@pytest.fixture()
+def full_sdata(_full_sdata_session: SpatialData) -> SpatialData:
+ return _fast_deepcopy_sdata(_full_sdata_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_full_session() -> SpatialData:
+ return SpatialData(
+ images=_get_images(),
+ labels=_get_labels(),
+ shapes=_get_shapes(),
+ points=_get_points(),
+ tables=_get_tables(region="labels2d"),
+ )
+
+
 @pytest.fixture(
  # params=["labels"]
  params=["full", "empty"]
@@ -110,15 +162,9 @@ def full_sdata() -> SpatialData:
  ]
  # + ["empty_" + x for x in ["table"]] # TODO: empty table not supported yet
 )
-def sdata(request) -> SpatialData:
+def sdata(request, _sdata_full_session: SpatialData) -> SpatialData:
  if request.param == "full":
- return SpatialData(
- images=_get_images(),
- labels=_get_labels(),
- shapes=_get_shapes(),
- points=_get_points(),
- tables=_get_tables(region="labels2d"),
- )
+ return _fast_deepcopy_sdata(_sdata_full_session)
  if request.param == "empty":
  return SpatialData()
  return request.getfixturevalue(request.param)
@@ -304,18 +350,38 @@ def _get_new_table(spatial_element: None | str | Sequence[str], instance_id: Non
  return TableModel.parse(adata=adata, spatial_element=spatial_element, instance_id=instance_id)
 
 
-@pytest.fixture()
-def labels_blobs() -> ArrayLike:
- """Create a 2D labels."""
+@pytest.fixture(scope="session")
+def _labels_blobs_session() -> ArrayLike:
  return BlobsDataset()._labels_blobs()
 
 
 @pytest.fixture()
-def sdata_blobs() -> SpatialData:
+def labels_blobs(_labels_blobs_session: ArrayLike) -> ArrayLike:
  """Create a 2D labels."""
+ return deepcopy(_labels_blobs_session)
+
+
+@pytest.fixture(scope="session")
+def _sdata_blobs_session() -> SpatialData:
  from spatialdata.datasets import blobs
 
- return deepcopy(blobs(256, 300, 3))
+ return blobs(256, 300, 3)
+
+
+@pytest.fixture()
+def sdata_blobs(_sdata_blobs_session: SpatialData) -> SpatialData:
+ """Create a 2D labels."""
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+
+@pytest.fixture()
+def blobs_factory(_sdata_blobs_session: SpatialData) -> Callable[[], SpatialData]:
+ """Return a factory that creates cheap fresh copies of the session-scoped blobs dataset."""
+
+ def _make() -> SpatialData:
+ return _fast_deepcopy_sdata(_sdata_blobs_session)
+
+ return _make
 
 
 def _make_points(coordinates: np.ndarray) -> DaskDataFrame:

tests/core/operations/conftest.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import os
+
+# Disable numba JIT compilation for rasterize tests. Datashader (used by rasterize) triggers
+# numba JIT on first call, costing ~1.4s per worker. Python-mode gives identical results for
+# the small test data here — unlike real data, there is no throughput advantage from JIT.
+os.environ.setdefault("NUMBA_DISABLE_JIT", "1")