From 40f6477c3847b12455008978cb9607f58f75f2df Mon Sep 17 00:00:00 2001 From: Tomatokeftes <129113023+Tomatokeftes@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:24:57 +0200 Subject: [PATCH 1/5] feat: add lazy table loading via anndata.experimental.read_lazy Add a `lazy: bool = False` parameter to SpatialData.read(), read_zarr() and _read_table(), backing tables with anndata.experimental.read_lazy() instead of loading them into memory. Large tables -- Mass Spectrometry Imaging data can reach millions of pixels by hundreds of thousands of m/z bins -- otherwise have to be read eagerly even though every other element type is already lazy. Adds a _is_lazy_anndata() helper and skips the eager dtype/null checks in TableModel validation for lazy tables, since those checks would materialize the data and defeat the purpose. If lazy=True and read_lazy() fails the error propagates rather than silently falling back to an eager read, which could exhaust memory on large stores. --- src/spatialdata/_core/spatialdata.py | 61 +++++++++++++--------------- src/spatialdata/_io/io_table.py | 38 +++++++++++++---- src/spatialdata/_io/io_zarr.py | 21 +++++----- src/spatialdata/models/models.py | 36 +++++++++++++++- 4 files changed, 102 insertions(+), 54 deletions(-) diff --git a/src/spatialdata/_core/spatialdata.py b/src/spatialdata/_core/spatialdata.py index fb55ab086..b1e9a485c 100644 --- a/src/spatialdata/_core/spatialdata.py +++ b/src/spatialdata/_core/spatialdata.py @@ -52,10 +52,7 @@ if TYPE_CHECKING: from spatialdata._core.query.spatial_query import BaseSpatialRequest - from spatialdata._io.format import ( - SpatialDataContainerFormatType, - SpatialDataFormatType, - ) + from spatialdata._io.format import SpatialDataContainerFormatType, SpatialDataFormatType class SpatialData: @@ -232,9 +229,7 @@ def get_annotated_regions(table: AnnData) -> list[str]: ------- The annotated regions. """ - from spatialdata.models.models import ( - _get_region_metadata_from_region_key_column, - ) + from spatialdata.models.models import _get_region_metadata_from_region_key_column return _get_region_metadata_from_region_key_column(table) @@ -705,9 +700,7 @@ def _filter_tables( if table is not None and len(table) != 0: tables[table_name] = table elif by == "elements": - from spatialdata._core.query.relational_query import ( - _filter_table_by_elements, - ) + from spatialdata._core.query.relational_query import _filter_table_by_elements assert elements_dict is not None table = _filter_table_by_elements(table, elements_dict=elements_dict) @@ -732,10 +725,7 @@ def rename_coordinate_systems(self, rename_dict: dict[str, str]) -> None: The method does not allow to rename a coordinate system into an existing one, unless the existing one is also renamed in the same call. """ - from spatialdata.transformations.operations import ( - get_transformation, - set_transformation, - ) + from spatialdata.transformations.operations import get_transformation, set_transformation # check that the rename_dict is valid old_names = self.coordinate_systems @@ -1111,7 +1101,7 @@ def write( overwrite: bool = False, consolidate_metadata: bool = True, update_sdata_path: bool = True, - sdata_formats: SpatialDataFormatType | list[SpatialDataFormatType] | None = None, + sdata_formats: (SpatialDataFormatType | list[SpatialDataFormatType] | None) = None, shapes_geometry_encoding: Literal["WKB", "geoarrow"] | None = None, raster_compressor: dict[Literal["lz4", "zstd"], int] | None = None, ) -> None: @@ -1225,15 +1215,12 @@ def _write_element( ) root_group, element_type_group, element_group = _get_groups_for_element( - zarr_path=zarr_container_path, element_type=element_type, element_name=element_name, use_consolidated=False - ) - from spatialdata._io import ( - write_image, - write_labels, - write_points, - write_shapes, - write_table, + zarr_path=zarr_container_path, + element_type=element_type, + element_name=element_name, + use_consolidated=False, ) + from spatialdata._io import write_image, write_labels, write_points, write_shapes, write_table from spatialdata._io.format import _parse_formats if parsed_formats is None: @@ -1287,7 +1274,7 @@ def write_element( self, element_name: str | list[str], overwrite: bool = False, - sdata_formats: SpatialDataFormatType | list[SpatialDataFormatType] | None = None, + sdata_formats: (SpatialDataFormatType | list[SpatialDataFormatType] | None) = None, shapes_geometry_encoding: Literal["WKB", "geoarrow"] | None = None, raster_compressor: dict[Literal["lz4", "zstd"], int] | None = None, ) -> None: @@ -1573,7 +1560,10 @@ def write_channel_names(self, element_name: str | None = None) -> None: # Mypy does not understand that path is not None so we have the check in the conditional if element_type == "images" and self.path is not None: _, _, element_group = _get_groups_for_element( - zarr_path=Path(self.path), element_type=element_type, element_name=element_name, use_consolidated=False + zarr_path=Path(self.path), + element_type=element_type, + element_name=element_name, + use_consolidated=False, ) from spatialdata._io._utils import overwrite_channel_names @@ -1624,19 +1614,18 @@ def write_transformations(self, element_name: str | None = None) -> None: ) axes = get_axes_names(element) if isinstance(element, DataArray | DataTree): - from spatialdata._io._utils import ( - overwrite_coordinate_transformations_raster, - ) + from spatialdata._io._utils import overwrite_coordinate_transformations_raster from spatialdata._io.format import RasterFormats raster_format = RasterFormats[element_group.metadata.attributes["spatialdata_attrs"]["version"]] overwrite_coordinate_transformations_raster( - group=element_group, axes=axes, transformations=transformations, raster_format=raster_format + group=element_group, + axes=axes, + transformations=transformations, + raster_format=raster_format, ) elif isinstance(element, DaskDataFrame | GeoDataFrame | AnnData): - from spatialdata._io._utils import ( - overwrite_coordinate_transformations_non_raster, - ) + from spatialdata._io._utils import overwrite_coordinate_transformations_non_raster overwrite_coordinate_transformations_non_raster( group=element_group, @@ -1855,6 +1844,7 @@ def read( file_path: str | Path | UPath | zarr.Group, selection: tuple[str] | None = None, reconsolidate_metadata: bool = False, + lazy: bool = False, ) -> SpatialData: """ Read a SpatialData object from a Zarr storage (on-disk or remote). @@ -1867,6 +1857,11 @@ def read( The elements to read (images, labels, points, shapes, table). If None, all elements are read. reconsolidate_metadata If the consolidated metadata store got corrupted this can lead to errors when trying to read the data. + lazy + If True, read tables lazily using anndata.experimental.read_lazy. + This keeps large tables out of memory until needed. Requires anndata >= 0.12. + Note: Images, labels, and points are always read lazily (using Dask). + This parameter only affects tables, which are normally loaded into memory. Returns ------- @@ -1879,7 +1874,7 @@ def read( _write_consolidated_metadata(file_path) - return read_zarr(file_path, selection=selection) + return read_zarr(file_path, selection=selection, lazy=lazy) @property def images(self) -> Images: diff --git a/src/spatialdata/_io/io_table.py b/src/spatialdata/_io/io_table.py index 3eb4b0927..f75c91b7a 100644 --- a/src/spatialdata/_io/io_table.py +++ b/src/spatialdata/_io/io_table.py @@ -9,18 +9,38 @@ from anndata._io.specs import write_elem as write_adata from ome_zarr.format import Format -from spatialdata._io.format import ( - CurrentTablesFormat, - TablesFormats, - TablesFormatV01, - TablesFormatV02, - _parse_version, -) +from spatialdata._io.format import CurrentTablesFormat, TablesFormats, TablesFormatV01, TablesFormatV02, _parse_version from spatialdata.models import TableModel, get_table_keys -def _read_table(store: str | Path) -> AnnData: - table = read_anndata_zarr(str(store)) +def _read_table(store: str | Path, lazy: bool = False) -> AnnData: + """ + Read a table from a zarr store. + + Parameters + ---------- + store + Path to the zarr store containing the table. + lazy + If True, read the table lazily using ``anndata.experimental.read_lazy``. + This keeps large matrices (X, layers) as dask arrays backed by zarr, + so they are only loaded into memory on demand. Requires anndata >= 0.12. + + Returns + ------- + The AnnData table, either lazily loaded or in-memory. + + Raises + ------ + ImportError + If ``lazy=True`` but anndata >= 0.12 is not installed. + """ + if lazy: + from anndata.experimental import read_lazy + + table = read_lazy(str(store)) + else: + table = read_anndata_zarr(str(store)) f = zarr.open(Path(store), mode="r") # Path avoids zarr v3 URL-parsing special chars (e.g. #) in names version = _parse_version(f, expect_attrs_key=False) diff --git a/src/spatialdata/_io/io_zarr.py b/src/spatialdata/_io/io_zarr.py index 4c410fab0..0e36b2a04 100644 --- a/src/spatialdata/_io/io_zarr.py +++ b/src/spatialdata/_io/io_zarr.py @@ -3,6 +3,7 @@ import os import warnings from collections.abc import Callable +from functools import partial from json import JSONDecodeError from pathlib import Path from typing import Any, Literal, cast @@ -17,11 +18,7 @@ from zarr.errors import ArrayNotFoundError from spatialdata._core.spatialdata import SpatialData -from spatialdata._io._utils import ( - BadFileHandleMethod, - _resolve_zarr_store, - handle_read_errors, -) +from spatialdata._io._utils import BadFileHandleMethod, _resolve_zarr_store, handle_read_errors from spatialdata._io.io_points import _read_points from spatialdata._io.io_raster import _read_multiscale from spatialdata._io.io_shapes import _read_shapes @@ -106,10 +103,7 @@ def get_raster_format_for_read( ------- The ome-zarr format to use for reading the raster element. """ - from spatialdata._io.format import ( - sdata_zarr_version_to_ome_zarr_format, - sdata_zarr_version_to_raster_format, - ) + from spatialdata._io.format import sdata_zarr_version_to_ome_zarr_format, sdata_zarr_version_to_raster_format if sdata_version == "0.1": group_version = group.metadata.attributes["multiscales"][0]["version"] @@ -126,6 +120,7 @@ def read_zarr( store: str | Path | UPath | zarr.Group, selection: None | tuple[str] = None, on_bad_files: Literal[BadFileHandleMethod.ERROR, BadFileHandleMethod.WARN] = BadFileHandleMethod.ERROR, + lazy: bool = False, ) -> SpatialData: """ Read a SpatialData dataset from a zarr store (on-disk or remote). @@ -149,6 +144,12 @@ def read_zarr( object is returned containing only elements that could be read. Failures can only be determined from the warnings. + lazy + If True, read tables lazily using anndata.experimental.read_lazy. + This keeps large tables out of memory until needed. Requires anndata >= 0.12. + Note: Images, labels, and points are always read lazily (using Dask). + This parameter only affects tables, which are normally loaded into memory. + Returns ------- A SpatialData object. @@ -195,7 +196,7 @@ def read_zarr( "labels": (_read_multiscale, "labels", labels), "points": (_read_points, "points", points), "shapes": (_read_shapes, "shapes", shapes), - "tables": (_read_table, "tables", tables), + "tables": (partial(_read_table, lazy=lazy), "tables", tables), } for group_name, ( read_func, diff --git a/src/spatialdata/models/models.py b/src/spatialdata/models/models.py index a818bdadc..ad9231016 100644 --- a/src/spatialdata/models/models.py +++ b/src/spatialdata/models/models.py @@ -55,6 +55,25 @@ ATTRS_KEY = "spatialdata_attrs" +def _is_lazy_anndata(adata: AnnData) -> bool: + """Check if an AnnData object is lazily loaded. + + Lazy AnnData objects (from anndata.experimental.read_lazy) have obs/var + stored as xarray Dataset2D instead of pandas DataFrame. + + Parameters + ---------- + adata + The AnnData object to check. + + Returns + ------- + True if the AnnData is lazily loaded, False otherwise. + """ + # Check if obs is not a pandas DataFrame (lazy AnnData uses xarray Dataset2D) + return not isinstance(adata.obs, pd.DataFrame) + + def _parse_transformations(element: SpatialElement, transformations: MappingToCoordinateSystem_t | None = None) -> None: _validate_mapping_to_coordinate_system_type(transformations) transformations_in_element = _get_transformations(element) @@ -1085,6 +1104,13 @@ def _validate_table_annotation_metadata(cls, data: AnnData) -> None: raise ValueError(f"`{attr[cls.REGION_KEY_KEY]}` not found in `adata.obs`. Please create the column.") if attr[cls.INSTANCE_KEY] not in data.obs: raise ValueError(f"`{attr[cls.INSTANCE_KEY]}` not found in `adata.obs`. Please create the column.") + + # Skip detailed dtype/value validation for lazy-loaded AnnData + # These checks would trigger data loading, defeating the purpose of lazy loading + # Validation will occur when data is actually computed/accessed + if _is_lazy_anndata(data): + return + instance_col = data.obs[attr[cls.INSTANCE_KEY]] dtype = instance_col.dtype @@ -1154,6 +1180,10 @@ def validate( if ATTRS_KEY not in data.uns: return data + # Check if this is a lazy-loaded AnnData (from anndata.experimental.read_lazy) + # Lazy AnnData has xarray-based obs/var, which requires different validation + is_lazy = _is_lazy_anndata(data) + _, region_key, instance_key = get_table_keys(data) if region_key is not None: if region_key not in data.obs: @@ -1161,7 +1191,8 @@ def validate( f"Region key `{region_key}` not in `adata.obs`. Please create the column and parse " f"using TableModel.parse(adata)." ) - if not isinstance(data.obs[region_key].dtype, CategoricalDtype): + # Skip dtype validation for lazy tables (would require loading data) + if not is_lazy and not isinstance(data.obs[region_key].dtype, CategoricalDtype): raise ValueError( f"`table.obs[{region_key}]` must be of type `categorical`, not `{type(data.obs[region_key])}`." ) @@ -1171,7 +1202,8 @@ def validate( f"Instance key `{instance_key}` not in `adata.obs`. Please create the column and parse" f" using TableModel.parse(adata)." ) - if data.obs[instance_key].isnull().values.any(): + # Skip null check for lazy tables (would require loading data) + if not is_lazy and data.obs[instance_key].isnull().values.any(): raise ValueError("`table.obs[instance_key]` must not contain null values, but it does.") cls._validate_table_annotation_metadata(data) From 966a4fe13dadba6d20ca523be7b3770817dae17d Mon Sep 17 00:00:00 2001 From: Tomatokeftes <129113023+Tomatokeftes@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:25:07 +0200 Subject: [PATCH 2/5] fix: handle lazy Dataset2D obs in query and categorical-subset paths Tables read via anndata.experimental.read_lazy() expose obs as an xarray Dataset2D rather than a pandas DataFrame. pd.DataFrame(obs) accepts it but returns a malformed frame instead of raising, so the failure surfaces later as wrong values rather than an error. Materialize with obs.to_memory() on that branch in get_values() and _inplace_fix_subset_categorical_obs(), which is what bounding_box_query() and aggregate() go through. --- src/spatialdata/_core/query/relational_query.py | 7 ++++++- src/spatialdata/_utils.py | 10 ++++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/spatialdata/_core/query/relational_query.py b/src/spatialdata/_core/query/relational_query.py index 7ef7c1a07..f82221fc6 100644 --- a/src/spatialdata/_core/query/relational_query.py +++ b/src/spatialdata/_core/query/relational_query.py @@ -1076,7 +1076,12 @@ def get_values( if origin == "obs": df = obs[value_key_values].copy() if origin == "var": - matched_table.obs = pd.DataFrame(obs) + # When the table came from anndata.experimental.read_lazy, obs is a Dataset2D, not a + # DataFrame, and pd.DataFrame(obs) returns a malformed frame. Materialize via to_memory(). + if isinstance(obs, pd.DataFrame): + matched_table.obs = pd.DataFrame(obs) + else: + matched_table.obs = obs.to_memory() if table_layer is None: x = matched_table[:, value_key_values].X else: diff --git a/src/spatialdata/_utils.py b/src/spatialdata/_utils.py index 609cd0403..7473df026 100644 --- a/src/spatialdata/_utils.py +++ b/src/spatialdata/_utils.py @@ -225,11 +225,17 @@ def _inplace_fix_subset_categorical_obs(subset_adata: AnnData, original_adata: A """ if not hasattr(subset_adata, "obs") or not hasattr(original_adata, "obs"): return - obs = pd.DataFrame(subset_adata.obs) + # Tables read via anndata.experimental.read_lazy have a Dataset2D obs instead of a DataFrame; + # pd.DataFrame() would silently malform it, so materialize with to_memory() in that case. + obs = pd.DataFrame(subset_adata.obs) if isinstance(subset_adata.obs, pd.DataFrame) else subset_adata.obs.to_memory() + original_obs = ( + original_adata.obs if isinstance(original_adata.obs, pd.DataFrame) else original_adata.obs.to_memory() + ) + for column in obs.columns: is_categorical = isinstance(obs[column].dtype, pd.CategoricalDtype) if is_categorical: - c = obs[column].cat.set_categories(original_adata.obs[column].cat.categories) + c = obs[column].cat.set_categories(original_obs[column].cat.categories) obs[column] = c subset_adata.obs = obs From 1f6e8ca9a65748ddadafbb4f730acf026da50b61 Mon Sep 17 00:00:00 2001 From: Tomatokeftes <129113023+Tomatokeftes@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:25:07 +0200 Subject: [PATCH 3/5] test: cover lazy table reading Adds coverage that lazy=True builds a SpatialData object, that lazy=False is unchanged, and that the flag is threaded through read_zarr(). --- tests/io/test_readwrite.py | 90 +++++++++++++++++++++++++++++++++++++- 1 file changed, 89 insertions(+), 1 deletion(-) diff --git a/tests/io/test_readwrite.py b/tests/io/test_readwrite.py index 034c01d37..b14669575 100644 --- a/tests/io/test_readwrite.py +++ b/tests/io/test_readwrite.py @@ -101,7 +101,11 @@ def test_shapes( # add a mixed Polygon + MultiPolygon element shapes["mixed"] = pd.concat([shapes["poly"], shapes["multipoly"]]) - shapes.write(tmpdir, sdata_formats=sdata_container_format, shapes_geometry_encoding=geometry_encoding) + shapes.write( + tmpdir, + sdata_formats=sdata_container_format, + shapes_geometry_encoding=geometry_encoding, + ) sdata = SpatialData.read(tmpdir) if geometry_encoding == "WKB": @@ -1334,3 +1338,87 @@ def test_sdata_with_nan_in_obs(tmp_path: Path) -> None: else: # After round-trip, NaN in object-dtype column becomes string "nan" on pandas 2 assert r1.iloc[1] == "nan" assert np.isnan(r2.iloc[0]) + + +class TestLazyTableLoading: + """Tests for lazy table loading functionality. + + Lazy loading uses anndata.experimental.read_lazy() to keep large tables + out of memory until needed. This is particularly useful for MSI data + where tables can contain millions of pixels. + """ + + @pytest.fixture + def sdata_with_table(self) -> SpatialData: + """Create a SpatialData object with a simple table for testing.""" + from spatialdata.models import TableModel + + rng = default_rng(42) + table = TableModel.parse( + AnnData( + X=rng.random((100, 50)), + obs=pd.DataFrame( + { + "region": pd.Categorical(["region1"] * 100), + "instance": np.arange(100), + } + ), + ), + region_key="region", + instance_key="instance", + region="region1", + ) + return SpatialData(tables={"test_table": table}) + + def test_lazy_read_basic(self, sdata_with_table: SpatialData) -> None: + """Test that lazy=True reads tables without loading into memory.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "data.zarr") + sdata_with_table.write(path) + + # Read with lazy=True + try: + sdata_lazy = SpatialData.read(path, lazy=True) + + # Table should be present + assert "test_table" in sdata_lazy.tables + + # Check that X is a lazy array (dask or similar) + # Lazy AnnData from read_lazy uses dask arrays + table = sdata_lazy.tables["test_table"] + assert hasattr(table, "X") + + except ImportError: + # If anndata.experimental.read_lazy is not available, skip + pytest.skip("anndata.experimental.read_lazy not available") + + def test_lazy_false_loads_normally(self, sdata_with_table: SpatialData) -> None: + """Test that lazy=False (default) loads tables into memory normally.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "data.zarr") + sdata_with_table.write(path) + + # Read with lazy=False (default) + sdata_normal = SpatialData.read(path, lazy=False) + + # Table should be present and loaded normally + assert "test_table" in sdata_normal.tables + table = sdata_normal.tables["test_table"] + + # X should be a numpy array or scipy sparse matrix (in-memory) + import scipy.sparse as sp + + assert isinstance(table.X, np.ndarray | sp.spmatrix) + + def test_read_zarr_lazy_parameter(self, sdata_with_table: SpatialData) -> None: + """Test that read_zarr function accepts lazy parameter.""" + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "data.zarr") + sdata_with_table.write(path) + + # Test read_zarr directly with lazy parameter + try: + sdata = read_zarr(path, lazy=True) + assert "test_table" in sdata.tables + except ImportError: + pytest.skip("anndata.experimental.read_lazy not available") From 9e2a8a1dac459cf23538c0a9319491f12dfa4501 Mon Sep 17 00:00:00 2001 From: Tomatokeftes <129113023+Tomatokeftes@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:54:01 +0200 Subject: [PATCH 4/5] fix: keep relational queries working on lazily-read tables Upstream #1131 rewrote the join helpers to call pandas methods on table.obs directly (reset_index, groupby). A table read with anndata.experimental.read_lazy exposes obs as an xarray Dataset2D, which implements neither, so join_spatialelement_table, bounding_box_query and get_values all raised AttributeError on lazy tables. Route those five call sites through _obs_as_dataframe(), which materializes obs only when it is not already a DataFrame. obs is the small axis, so X stays lazy. get_values additionally needs X computed before it reaches pd.DataFrame(), which cannot consume a dask array; that is a selection of the requested columns only. Adds a regression test covering all five join modes, bounding_box_query and get_values against a lazy table, asserting the results equal the eager read and that X is still a dask array afterwards. The test fails without this change. --- .../_core/query/relational_query.py | 27 ++++-- tests/io/test_readwrite.py | 85 +++++++++++++++++++ 2 files changed, 107 insertions(+), 5 deletions(-) diff --git a/src/spatialdata/_core/query/relational_query.py b/src/spatialdata/_core/query/relational_query.py index f82221fc6..e0ec9d9b5 100644 --- a/src/spatialdata/_core/query/relational_query.py +++ b/src/spatialdata/_core/query/relational_query.py @@ -246,6 +246,18 @@ def _region_as_str_if_list_of_len_one(region: list[str]) -> str | list[str]: return region if len(region) > 1 else region[0] +def _obs_as_dataframe(table: AnnData) -> pd.DataFrame: + """Return ``table.obs`` as a pandas DataFrame. + + Tables read with ``anndata.experimental.read_lazy`` expose ``obs`` as an xarray + ``Dataset2D``, which does not implement the pandas methods (``reset_index``, + ``groupby``) the join helpers rely on. Materialize it in that case; ``obs`` is the + small axis of the table, so this does not pull in ``X``, which stays lazy. + """ + obs = table.obs + return obs if isinstance(obs, pd.DataFrame) else obs.to_memory() + + def _right_exclusive_join_spatialelement_table( element_dict: dict[str, dict[str, Any]], table: AnnData, @@ -256,7 +268,7 @@ def _right_exclusive_join_spatialelement_table( if isinstance(regions, str): regions = [regions] # reset_index so group_df.index gives integer positions — safe with duplicate obs names - obs = table.obs.reset_index() + obs = _obs_as_dataframe(table).reset_index() groups_df = obs.groupby(by=region_column_name, observed=False) keep = np.zeros(len(table), dtype=bool) has_match = False @@ -301,7 +313,7 @@ def _right_join_spatialelement_table( regions, region_column_name, instance_key = get_table_keys(table) if isinstance(regions, str): regions = [regions] - groups_df = table.obs.groupby(by=region_column_name, observed=False) + groups_df = _obs_as_dataframe(table).groupby(by=region_column_name, observed=False) for element_type, name_element in element_dict.items(): for name, element in name_element.items(): if name in regions: @@ -343,7 +355,7 @@ def _inner_join_spatialelement_table( regions, region_column_name, instance_key = get_table_keys(table) if isinstance(regions, str): regions = [regions] - obs = table.obs.reset_index() + obs = _obs_as_dataframe(table).reset_index() groups_df = obs.groupby(by=region_column_name, observed=False) joined_indices = None for element_type, name_element in element_dict.items(): @@ -404,7 +416,7 @@ def _left_exclusive_join_spatialelement_table( regions, region_column_name, instance_key = get_table_keys(table) if isinstance(regions, str): regions = [regions] - groups_df = table.obs.groupby(by=region_column_name, observed=False) + groups_df = _obs_as_dataframe(table).groupby(by=region_column_name, observed=False) for element_type, name_element in element_dict.items(): for name, element in name_element.items(): if name in regions: @@ -442,7 +454,7 @@ def _left_join_spatialelement_table( regions, region_column_name, instance_key = get_table_keys(table) if isinstance(regions, str): regions = [regions] - obs = table.obs.reset_index() + obs = _obs_as_dataframe(table).reset_index() groups_df = obs.groupby(by=region_column_name, observed=False) joined_indices = None for element_type, name_element in element_dict.items(): @@ -1090,6 +1102,11 @@ def get_values( x = matched_table[:, value_key_values].layers[table_layer] import scipy + if isinstance(x, da.Array): + # A lazy table backs X with dask, and pd.DataFrame() cannot consume that. + # get_values returns in-memory values by contract, and this is a selection of + # the requested columns only, so the materialization is bounded by those. + x = x.compute() if isinstance(x, scipy.sparse.csr_matrix | scipy.sparse.csc_matrix | scipy.sparse.coo_matrix): x = x.todense() df = pd.DataFrame(x, columns=value_key_values) diff --git a/tests/io/test_readwrite.py b/tests/io/test_readwrite.py index b14669575..ba9c6f381 100644 --- a/tests/io/test_readwrite.py +++ b/tests/io/test_readwrite.py @@ -1422,3 +1422,88 @@ def test_read_zarr_lazy_parameter(self, sdata_with_table: SpatialData) -> None: assert "test_table" in sdata.tables except ImportError: pytest.skip("anndata.experimental.read_lazy not available") + + @pytest.fixture + def sdata_with_shapes_and_table(self) -> SpatialData: + """A table annotating a shapes element, so the relational queries have something to join.""" + from geopandas import GeoDataFrame + from scipy.sparse import csr_matrix + from shapely.geometry import Point + + from spatialdata.models import ShapesModel, TableModel + + n = 20 + rng = default_rng(42) + # Circles on a line, so a bounding box selects a known prefix. + circles = ShapesModel.parse( + GeoDataFrame( + {"geometry": [Point(float(i), 0.0) for i in range(n)], "radius": np.full(n, 0.1)}, + index=np.arange(n), + ) + ) + table = TableModel.parse( + AnnData( + X=csr_matrix(rng.random((n, 5)) * (rng.random((n, 5)) > 0.5)), + obs=pd.DataFrame({"region": pd.Categorical(["circles"] * n), "instance": np.arange(n)}), + var=pd.DataFrame(index=[f"g{i}" for i in range(5)]), + ), + region_key="region", + instance_key="instance", + region="circles", + ) + return SpatialData(shapes={"circles": circles}, tables={"table": table}) + + def test_lazy_table_relational_queries_match_eager(self, sdata_with_shapes_and_table: SpatialData) -> None: + """Relational queries must work on a lazy table and agree with the eager read. + + A lazy table's obs is an xarray Dataset2D, not a DataFrame, so any pandas-only call + in the join helpers (``reset_index``, ``groupby``) breaks these paths. That is not + covered by simply reading a table lazily, which is why it is asserted here. + """ + from spatialdata import bounding_box_query, get_values, join_spatialelement_table + + with tempfile.TemporaryDirectory() as tmpdir: + path = os.path.join(tmpdir, "data.zarr") + sdata_with_shapes_and_table.write(path) + + try: + lazy = SpatialData.read(path, lazy=True) + except ImportError: + pytest.skip("anndata.experimental.read_lazy not available") + eager = SpatialData.read(path, lazy=False) + + # The table is genuinely lazy, otherwise the rest of this test proves nothing. + assert isinstance(lazy.tables["table"].X, da.Array) + assert not isinstance(lazy.tables["table"].obs, pd.DataFrame) + + for how in ["left", "inner", "right", "left_exclusive", "right_exclusive"]: + element_dict, joined = join_spatialelement_table( + spatial_element_names=["circles"], + spatial_elements=[lazy.shapes["circles"]], + table=lazy.tables["table"], + how=how, + ) + assert element_dict is not None + + kwargs = { + "min_coordinate": [-1, -1], + "max_coordinate": [9.5, 1], + "axes": ("x", "y"), + "target_coordinate_system": "global", + } + assert len(bounding_box_query(lazy, **kwargs).shapes["circles"]) == len( + bounding_box_query(eager, **kwargs).shapes["circles"] + ) + + lazy_values = get_values("g3", sdata=lazy, element_name="circles", table_name="table") + eager_values = get_values("g3", sdata=eager, element_name="circles", table_name="table") + np.testing.assert_allclose( + np.asarray(lazy_values).ravel().astype(float), + np.asarray(eager_values).ravel().astype(float), + ) + + # X is still lazy after all of the above, and slicing it agrees with the eager read. + assert isinstance(lazy.tables["table"].X, da.Array) + lazy_block = lazy.tables["table"].X[:4, :3].compute() + eager_block = eager.tables["table"].X[:4, :3] + np.testing.assert_allclose(lazy_block.toarray(), eager_block.toarray()) From 4b1da506a8ca086ca4e0b77cca47e9092f4f5835 Mon Sep 17 00:00:00 2001 From: Tomatokeftes <129113023+Tomatokeftes@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:54:02 +0200 Subject: [PATCH 5/5] docs: note that dask reductions over a sparse X are unsupported A lazily-read sparse table backs X with a dask array whose blocks are scipy.sparse matrices. Dask builds a reduction's result metadata by calling the matching NumPy reduction on one block, and scipy.sparse rejects the keepdims and ndmin arguments NumPy forwards, so X.sum(), .mean(), .max() and .std() all raise while the graph is built, before compute() is reached. This is a dask/scipy.sparse interoperability limitation rather than something the lazy reader introduces, and it is not fixable here. Document it on read_zarr() and SpatialData.read() along with the workarounds that do work, slicing before reducing and map_blocks, so it does not arrive as a bug report. --- src/spatialdata/_core/spatialdata.py | 17 +++++++++++++++++ src/spatialdata/_io/io_zarr.py | 16 ++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/src/spatialdata/_core/spatialdata.py b/src/spatialdata/_core/spatialdata.py index b1e9a485c..4b5633add 100644 --- a/src/spatialdata/_core/spatialdata.py +++ b/src/spatialdata/_core/spatialdata.py @@ -1863,6 +1863,23 @@ def read( Note: Images, labels, and points are always read lazily (using Dask). This parameter only affects tables, which are normally loaded into memory. + When the stored ``X`` is sparse, the lazy table's ``X`` is a Dask array whose + blocks are ``scipy.sparse`` matrices, and Dask's array reductions + (``X.sum()``, ``X.mean()``, ``X.max()``, ``X.std()``, ...) are **not + supported**. They raise a ``TypeError`` or ``IndexError`` while the graph is + being built -- before ``.compute()`` is ever reached -- because Dask derives + the result metadata by calling the corresponding NumPy reduction on a + ``scipy.sparse`` block, and ``scipy.sparse`` does not accept the + ``keepdims``/``ndmin`` arguments NumPy passes down. This is a + Dask/``scipy.sparse`` interoperability limitation, not something this reader + introduces. Slicing and ``.compute()`` work normally, so reduce a materialized + block instead:: + + table.X[:1000].compute().sum() # works + table.X.sum() # raises + + or use ``dask.array.map_blocks`` with a function that handles sparse blocks. + Returns ------- The SpatialData object. diff --git a/src/spatialdata/_io/io_zarr.py b/src/spatialdata/_io/io_zarr.py index 0e36b2a04..d8a2a3aea 100644 --- a/src/spatialdata/_io/io_zarr.py +++ b/src/spatialdata/_io/io_zarr.py @@ -150,6 +150,22 @@ def read_zarr( Note: Images, labels, and points are always read lazily (using Dask). This parameter only affects tables, which are normally loaded into memory. + When the stored ``X`` is sparse, the lazy table's ``X`` is a Dask array whose + blocks are ``scipy.sparse`` matrices, and Dask's array reductions + (``X.sum()``, ``X.mean()``, ``X.max()``, ``X.std()``, ...) are **not supported**. + They raise a ``TypeError`` or ``IndexError`` while the graph is being built -- + before ``.compute()`` is ever reached -- because Dask derives the result metadata + by calling the corresponding NumPy reduction on a ``scipy.sparse`` block, and + ``scipy.sparse`` does not accept the ``keepdims``/``ndmin`` arguments NumPy passes + down. This is a Dask/``scipy.sparse`` interoperability limitation, not something + this reader introduces. Slicing and ``.compute()`` work normally, so reduce a + materialized block instead:: + + table.X[:1000].compute().sum() # works + table.X.sum() # raises + + or use ``dask.array.map_blocks`` with a function that handles sparse blocks. + Returns ------- A SpatialData object.