feat: Add lazy table loading via anndata.experimental.read_lazy

[Tomatokeftes]

Summary

This PR adds support for lazy loading of tables in SpatialData using anndata's experimental read_lazy() function.

Motivation

Currently, all elements in SpatialData (images, labels, points) are loaded lazily using Dask, except for tables which are always loaded into memory. For large datasets, particularly Mass Spectrometry Imaging (MSI) data where tables can contain millions of pixels with hundreds of thousands of m/z bins, this creates memory bottlenecks.

Changes

Core lazy loading support:

• Add lazy: bool = False parameter to SpatialData.read() and read_zarr()
• Add lazy: bool = False parameter to _read_table() in io_table.py
• Use anndata.experimental.read_lazy() when lazy=True
• Add _is_lazy_anndata() helper function to detect lazy AnnData objects
• Modify validation to skip eager checks for lazy tables (prevents defeating lazy loading)
• Add fallback with warning if anndata version doesn't support read_lazy

Query API compatibility fixes:

• Fix _filter_table_by_element_names() to handle lazy Dataset2D obs
• Fix _filter_table_by_elements() to handle lazy Dataset2D obs
• Fix get_values() to handle lazy Dataset2D obs
• Fix _inplace_fix_subset_categorical_obs() to handle lazy tables

The query fixes ensure that bounding_box_query, aggregate, and other APIs work correctly with lazy-loaded tables. The issue was that pd.DataFrame(table.obs) doesn't correctly convert lazy Dataset2D objects - it produces a malformed DataFrame. The fix uses table.obs.to_memory() for lazy tables instead.

Usage

from spatialdata import SpatialData

# Load tables lazily (keeps large tables out of memory)
sdata = SpatialData.read("large_dataset.zarr", lazy=True)

# Access table - data is loaded on-demand
table = sdata.tables["my_table"]
# table.X is now backed by Dask/Zarr, not loaded into memory

# Query APIs work with lazy tables
from spatialdata import bounding_box_query
result = bounding_box_query(sdata, min_coordinate=[0, 0], max_coordinate=[100, 100], 
                            axes=("x", "y"), target_coordinate_system="global")

Benchmark Results

Test configuration: 100,000 pixels x 100,000 m/z bins, 3,000 peaks/pixel (~296M non-zeros)

Metric	Lazy Loading	Eager Loading	Improvement
Memory	15.4 MB	2,270.7 MB	99% savings
Time	0.13s	1.57s	12x faster

Reproducible Example

import numpy as np
from scipy import sparse
import anndata as ad
import psutil
import tempfile
from pathlib import Path

# Create synthetic sparse data (100k pixels x 100k m/z bins, 3000 peaks/pixel)
rng = np.random.default_rng(42)
n_pixels, n_mz, peaks_per_pixel = 100000, 100000, 3000
nnz = int(n_pixels * peaks_per_pixel)

X = sparse.csc_matrix(
    (rng.lognormal(7, 1.5, nnz).astype(np.float32),
     (rng.integers(0, n_pixels, nnz), rng.integers(0, n_mz, nnz))),
    shape=(n_pixels, n_mz)
)

# Create and write AnnData
adata = ad.AnnData(X=X)
adata.obs_names = [f"pixel_{i}" for i in range(n_pixels)]
adata.var_names = [f"mz_{i}" for i in range(n_mz)]

zarr_path = Path(tempfile.mkdtemp()) / "test.zarr"
adata.write_zarr(str(zarr_path))

# Compare lazy vs eager loading
from anndata.experimental import read_lazy
from anndata import read_zarr

def get_mem():
    return psutil.Process().memory_info().rss / 1e6

mem_before = get_mem()
adata_lazy = read_lazy(str(zarr_path))
print(f"Lazy:  +{get_mem() - mem_before:.1f} MB")

mem_before = get_mem()
adata_eager = read_zarr(str(zarr_path))
print(f"Eager: +{get_mem() - mem_before:.1f} MB")

Requirements

• Requires anndata >= 0.12 for lazy loading support
• Falls back to eager loading with a warning if anndata version is older

Real-world use case

This feature was developed for Thyra, a Mass Spectrometry Imaging converter. MSI datasets can have:

• Millions of pixels (observations)
• Hundreds of thousands of m/z bins (variables)
• Resulting in tables that exceed available RAM

With lazy loading, users can work with these datasets without loading the full table into memory.

Test plan

• test_lazy_read_basic - Verify lazy=True creates a SpatialData object without errors
• test_lazy_false_loads_normally - Verify lazy=False maintains current behavior
• test_read_zarr_lazy_parameter - Verify lazy parameter is passed through correctly
• All 29 relational query tests pass (verifies query API fixes)
• All 277 spatial query tests pass (verifies bounding_box_query works)
• Manual testing: lazy_loading, table_ops, query, aggregate, compute all work

[codecov]

Codecov Report

❌ Patch coverage is 97.43590% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 92.44%. Comparing base (eb4fb3d) to head (4b1da50).

Files with missing lines	Patch %	Lines
src/spatialdata/_core/query/relational_query.py	92.30%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1055   +/-   ##
=======================================
  Coverage   92.44%   92.44%           
=======================================
  Files          51       51           
  Lines        7820     7837   +17     
=======================================
+ Hits         7229     7245   +16     
- Misses        591      592    +1     

Files with missing lines	Coverage Δ
src/spatialdata/_core/spatialdata.py	93.86% <100.00%> (ø)
src/spatialdata/_io/io_table.py	91.11% <100.00%> (+0.63%)	⬆️
src/spatialdata/_io/io_zarr.py	92.45% <100.00%> (+0.07%)	⬆️
src/spatialdata/_utils.py	85.62% <100.00%> (+0.09%)	⬆️
src/spatialdata/models/models.py	88.15% <100.00%> (+0.10%)	⬆️
src/spatialdata/_core/query/relational_query.py	95.17% <92.30%> (-0.15%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Tomatokeftes]

Rebased onto current main (eb4fb3d). This was 5 commits behind; the rebase was
conflict-free and the history is now 5 commits instead of 13 (the merge commits and
pre-commit noise are gone).

While rebasing I found that this branch had quietly broken, and I think this is worth
flagging regardless of what happens to the PR. #1131 rewrote the join helpers to call
reset_index() and groupby() directly on table.obs. A table read through
anndata.experimental.read_lazy has an xarray Dataset2D there, which implements
neither, so join_spatialelement_table, bounding_box_query and get_values all
raised AttributeError on lazy tables. That was true of the previous head too, not
something the rebase introduced. Nothing caught it because no test exercised a query
against a lazy table.

Fixed by routing those call sites through a small _obs_as_dataframe() helper that
materializes obs only when it is not already a DataFrame (obs is the small axis,
X stays lazy), plus computing X in get_values before it reaches pd.DataFrame().
Added 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. It fails without the fix.

Benchmark

Re-measured on the rebased branch: 100,000 pixels x 100,000 m/z bins, 295,547,702
non-zeros. Each read in a fresh process, median of 3 runs.

	lazy	eager
Memory (RSS delta)	16.0 MB	2,383.4 MB
Wall time	0.071 s	1.62 s

anndata 0.13.2, dask 2026.7.1, scipy 1.18.0, zarr 3.3.0, Python 3.13.

Downstream consumer

This is on the critical path for Thyra,
a Mass Spectrometry Imaging converter that writes MSI acquisitions to SpatialData/Zarr.
MSI tables are pixels x m/z bins and routinely exceed RAM, so the table being the one
element type that is always read eagerly is the thing that blocks working with them at
all. Stores written by all three of Thyra's write paths (in-memory, streaming PCS,
streaming COO) read correctly under lazy=True on this branch, with lazy blocks
identical to the eager read; the encoding-type/encoding-version attrs that Thyra
already writes are sufficient, so nothing had to change downstream to consume this.

Known limitation, now documented

When the stored X is sparse, the dask array's blocks are scipy.sparse matrices and
dask's array reductions do not work:

table.X.sum()      # TypeError: _cs_matrix.sum() got an unexpected keyword argument 'keepdims'
table.X.mean()     # IndexError: Index dimension must be 1 or 2
table.X.max()      # TypeError: _cs_matrix.__init__() got an unexpected keyword argument 'ndmin'

Dask derives the result metadata by calling the matching NumPy reduction on one block,
and scipy.sparse rejects the keepdims/ndmin arguments NumPy forwards, so these
raise while the graph is built, before compute(). It is a dask/scipy.sparse
interoperability limitation rather than something this PR causes, and not fixable here.
Slicing and .compute() work normally, which is the realistic access pattern. I have
documented it on read_zarr() and SpatialData.read() with the workarounds, so it does
not arrive later as a "lazy loading is broken" issue.

Status

All 10 checks are green on the rebased head: the full test matrix (ubuntu/macOS/Windows,
Python 3.12/3.13/3.14, plus the min-dask job), pre-commit.ci, Read the Docs and codecov.

Happy to split the query fix into its own PR if you would rather review the two
separately, and equally happy to adjust the approach. This has been open a while without
a maintainer looking at it, so mostly I want to check whether the direction is something
you want at all before putting more into it.