feat: add table_shard_size_bytes to control the zarr shard size of tables

[Tomatokeftes]

Closes #1178.

Adds a keyword-only table_shard_size_bytes: int | None to SpatialData.write and
SpatialData.write_element, forwarded to write_table as shard_size_bytes. It is a target size in
bytes of uncompressed data for a single zarr shard of every array inside a table group.

sdata.write("data.zarr", table_shard_size_bytes=128 * 1024**2)
sdata.write_element("table", table_shard_size_bytes=8 * 1024**2)  # per table

Why a byte budget and not a chunks/shards tuple

This is deliberately not symmetric with #1106. A table is a heterogeneous tree of zarr arrays of
mixed rank, length and dtype that all receive one shared dataset_kwargs from anndata, so no single
tuple can be honoured by all of them. Measured on anndata 0.12.16 and zarr 3.2.1 with an ordinary
table:

• a 2-D chunks raises on obs/_index
• a 1-D chunks raises on 2-D obsm
• any shards tuple raises on the uns scalars
• shards without chunks raises on divisibility
• chunks=<int> broadcasts, but shards=<int> raises TypeError

A table_write_kwargs mirroring raster_write_kwargs would therefore ship an API whose documented
happy path cannot execute. A scalar budget avoids that: zarr derives shard = chunk * n per array, so
shard % chunk == 0 and shard <= array hold by construction at every rank, length and dtype.

I offered raster_shard_size_bytes as a symmetric form on the issue and have since withdrawn it. The
mechanism does not carry over: array.target_shard_size_bytes is read only when zarr is asked for an
automatic shard shape, and nothing injects shards="auto" on the raster side (shards does not appear
in _io/io_raster.py), so the same construction there would be a silently inert argument. A byte budget
is also the wrong shape for raster, where rank is uniform and storage_options already carries an
explicit per-level chunks. #1106's raster_write_kwargs looks like the right form for that side, so
this PR stays table only.

How it is delivered

Nothing is passed into dataset_kwargs. Two process globals are scoped around the existing anndata
call, for one table's write:

• zarr.config["array.target_shard_size_bytes"]
• anndata.settings.override(zarr_write_format=3, auto_shard_zarr_v3=True)

anndata then injects shards="auto" itself, only at the four writers where that is safe, and yields
to the caller-set budget instead of installing its own 1 GB default. The existing table.write_zarr(...)
and write_adata(group, name, table) calls are unchanged, and the #1183 re-fetch is outside the
scoped block, so the encoding attributes are untouched.

To be clear about what this does and does not do: it narrows the global, it does not remove it. Today
a downstream writer has to hold zarr.config open across a whole sdata.write; after this it is
scoped to one element and restored on exit, including on exceptions. It is still a process global
underneath.

Both write branches are wrapped, so the semantics are uniform across the supported anndata range with
no version-conditional code.

Two things worth flagging

shards must never reach dataset_kwargs. zarr's _guess_num_chunks_per_axis_shard does not
terminate on a rank-0 array while array.target_shard_size_bytes is set, and every SpatialData table
carries rank-0 string scalars in uns/spatialdata_attrs. It is an unbounded pure-Python loop, not an
error, so it would hang the write rather than fail it. Filed upstream as
zarr-developers/zarr-python#4304. Nothing here can reach it, and there is a fast test asserting that.

zarr_write_format has to be overridden alongside the sharding setting. AnnData.write_zarr reopens
the group with mode="w" and zarr_format=settings.zarr_write_format, destroying and recreating the
group spatialdata just made; with that setting left at 2, the argument would be silently inert
(measured: table group format 3 before, 2 after, X/data shards None, no error and no warning).

Validation

All errors are TableWriteOptionsError, a new ValueError subclass re-exported from the top level.
All four are raised up front in write and write_element, before anything reaches disk, because
_write_element creates the element group and write writes every preceding element before the
table is reached.

• not a positive int (bool included)
• zarr < 3.1.6: 3.1.4 added array.target_shard_size_bytes, but 3.1.4 and 3.1.5 still size the inner
  chunk with max_bytes=1024 where 1 MiB was intended (fixed by fix: auto-chunking when auto-sharding 1MiB number zarr-developers/zarr-python#3603),
  which would put roughly 130k inner chunks in a 128 MiB shard
• an anndata without zarr v3 auto-sharding support
• a zarr v2 table format, where sharding does not exist

The zarr and anndata gates are runtime checks, so zarr>=3.0.0 and anndata>=0.9.1 are unchanged and
no CI leg gains a dependency.

Setting the argument forces auto_shard_zarr_v3=True for the duration of each table write, so it
overrides an explicit False; there is no value that turns sharding off. The budget is a target, not
a bound: below the automatically chosen inner chunk it degenerates to one chunk per shard. Both are
documented.

Tests

tests/io/test_readwrite.py, on a purpose-built 4000 x 2000 CSR table (the shipped _get_table is
8 kB and cannot differentiate any budget):

• on-disk geometry for two budgets: shards present, shards % chunks == 0, and the smaller budget
  produces a strictly smaller shard
• every rank-0 array in the written table group is unsharded, for single-string and list region
• shards never reaches anndata, on both write branches (a fast guard, since the failure mode is a
  hang)
• the default write is byte-identical to a write with the argument absent
• both process globals are restored after a normal write and after one that raises
• rejected on a zarr v2 table format, with the store left uncreated
• invalid values rejected on write and write_element, with nothing written
• the Tables written on main lose the anndata encoding-type/encoding-version attributes #1183 encoding-metadata guard re-run with a budget set

Release notes

Added table_shard_size_bytes to SpatialData.write and SpatialData.write_element, to set a target
uncompressed size in bytes for the zarr shards of table arrays.

[codecov]

Codecov Report

❌ Patch coverage is 92.68293% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.89%. Comparing base (ccf1ea0) to head (9a2e82f).

Files with missing lines	Patch %	Lines
src/spatialdata/_io/_utils.py	92.59%	2 Missing ⚠️
src/spatialdata/_io/io_table.py	88.88%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1199   +/-   ##
=======================================
  Coverage   91.89%   91.89%           
=======================================
  Files          53       53           
  Lines        7942     7975   +33     
=======================================
+ Hits         7298     7329   +31     
- Misses        644      646    +2     

Files with missing lines	Coverage Δ
src/spatialdata/__init__.py	95.65% <ø> (ø)
src/spatialdata/_core/spatialdata.py	93.88% <100.00%> (+0.01%)	⬆️
src/spatialdata/_io/exceptions.py	63.63% <100.00%> (+3.63%)	⬆️
src/spatialdata/_io/io_table.py	89.83% <88.88%> (+0.35%)	⬆️
src/spatialdata/_io/_utils.py	87.23% <92.59%> (+0.56%)	⬆️

🚀 New features to boost your workflow:

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

[Tomaz-Vieira]

isinstance(x, bool) and not isinstance(x, int) are redundant. Plus, I don't think we should be checking the types of non-union arguments anyway =)

[Tomaz-Vieira]

I know we're trying to be nice here by catering to many versions of anndata (and zarr!) but I really dislike that we essentially lie to our users on the function signature, only to immediately disappoint by throwing an exception if the versions of zarr and/or anndata aren't what we need. We also completely defeat the type checker's ability to tell if the arguments are good or not.

One way around it would be to name those parameters as something like table_shard_size_bytes_hint (note the "hint" at the end); This makes it clear that they may or may not apply and we can just do nothing if that feature isn't supported.

Alternatively, we could create the type TableShardBudget, with a method like TableShardBudget.try_create(...), which is clearly visibly fallible, and would go through the validation logic in this function. This way if a client fails to get a TableShardBudget, then they can react accordingly (and locally to their code!), and all functions that use the budget don't have to re-validate. And you could also make the TableShardBudget be itself the context manager.

Maybe there is a way to have different signatures depending on what dependencies we have, but that would have strange impacts in our versioning scheme, so I'm skeptical that this could work.

Curious to see what other people think

[Tomaz-Vieira]

I think this docstring could be a bit more succint; I find it a bit hard to understand in the context of this PR, and would find it even harder when browsing the code out of context.

[Tomaz-Vieira]

I would just like to be absolutely sure that this is the only way to do this. Temporarily setting a global variable is a very dangerous design, even with the context managers (e.g.: how do we even know we're not already inside a context? what happens on multithreaded applications? etc), so if there is any way we could pass these arguments to the a function call, I'd much much much prefer that.