From dbece29ac6c0cbb3521d075c7a8993485d4d7855 Mon Sep 17 00:00:00 2001
From: Tomatokeftes <129113023+Tomatokeftes@users.noreply.github.com>
Date: Tue, 1 Sep 2026 14:58:43 +0200
Subject: [PATCH 1/4] feat: expose a table shard size via
 `table_shard_size_bytes`

Add a keyword-only `table_shard_size_bytes` to `SpatialData.write`,
`SpatialData.write_element` and `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.

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 a flat
chunks/shards tuple cannot be honoured: 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, and `shards` without `chunks` raises on
divisibility. A scalar byte budget is the one shape that executes.

Nothing is passed into `dataset_kwargs`. Two process globals are scoped
around the existing anndata call for the duration of one table write:
zarr's `array.target_shard_size_bytes`, and anndata's `zarr_write_format`
and `auto_shard_zarr_v3` settings. anndata then injects `shards="auto"`
itself, only at the writers where that is safe, and yields to the caller
set budget instead of installing its own 1 GB default. zarr derives the
shard shape from the chunk shape, so `shard % chunk == 0` and
`shard <= array` hold by construction at every rank.

`shards` deliberately never reaches `dataset_kwargs`: zarr's
`_guess_num_chunks_per_axis_shard` does not terminate on a rank-0 array
while a shard budget is set (zarr-developers/zarr-python#4304), and every
SpatialData table carries rank-0 string scalars in `uns/spatialdata_attrs`.

`zarr_write_format` is overridden alongside the sharding setting because
`AnnData.write_zarr` reopens the group with `zarr_format` taken from that
setting; leaving it at 2 silently produces a zarr v2 table group and no
sharding at all.

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

The argument is validated up front in `write` and `write_element`, before
any element reaches disk, and raises `TableWriteOptionsError` (a
`ValueError` subclass) when it is not a positive int, when zarr is older
than 3.1.6, when anndata does not support zarr v3 auto-sharding, or when
the table format is zarr v2. The zarr and anndata gates are runtime checks,
so no dependency pins change.

Closes #1178
---
 src/spatialdata/__init__.py          |   7 +
 src/spatialdata/_core/spatialdata.py |  49 +++++-
 src/spatialdata/_io/_utils.py        | 106 +++++++++++
 src/spatialdata/_io/exceptions.py    |   4 +
 src/spatialdata/_io/io_table.py      |  48 +++--
 tests/io/test_readwrite.py           | 251 ++++++++++++++++++++++++++-
 6 files changed, 441 insertions(+), 24 deletions(-)

diff --git a/src/spatialdata/__init__.py b/src/spatialdata/__init__.py
index 7ba66e710..87ba6eb11 100644
--- a/src/spatialdata/__init__.py
+++ b/src/spatialdata/__init__.py
@@ -61,6 +61,8 @@
     "SpatialData": "spatialdata._core.spatialdata",
     # _io._utils
     "get_dask_backing_files": "spatialdata._io._utils",
+    # _io.exceptions
+    "TableWriteOptionsError": "spatialdata._io.exceptions",
     # _io.format
     "SpatialDataFormatType": "spatialdata._io.format",
     # _io.io_zarr
@@ -119,6 +121,8 @@
     "SpatialData",
     # _io._utils
     "get_dask_backing_files",
+    # _io.exceptions
+    "TableWriteOptionsError",
     # _io.format
     "SpatialDataFormatType",
     # _io.io_zarr
@@ -211,6 +215,9 @@ def __dir__() -> list[str]:
     # _io._utils
     from spatialdata._io._utils import get_dask_backing_files
 
+    # _io.exceptions
+    from spatialdata._io.exceptions import TableWriteOptionsError
+
     # _io.format
     from spatialdata._io.format import SpatialDataFormatType
 
diff --git a/src/spatialdata/_core/spatialdata.py b/src/spatialdata/_core/spatialdata.py
index 6ee2296c8..d7eb50f45 100644
--- a/src/spatialdata/_core/spatialdata.py
+++ b/src/spatialdata/_core/spatialdata.py
@@ -1115,6 +1115,8 @@ def write(
         shapes_geometry_encoding: Literal["WKB", "geoarrow"] | None = None,
         raster_compressor: dict[Literal["lz4", "zstd"], int] | None = None,
         convert_table_strings_to_categoricals: bool = False,
+        *,
+        table_shard_size_bytes: int | None = None,
     ) -> None:
         """
         Write the `SpatialData` object to a Zarr store.
@@ -1170,12 +1172,33 @@ def write(
         convert_table_strings_to_categoricals
             If True, convert string columns of all tables to categoricals before writing.
             Note that this will have a side effect of modifying string columns into categoricals in place.
-        """
-        from spatialdata._io._utils import _resolve_zarr_store, _validate_compressor_args
+        table_shard_size_bytes
+            The target size in bytes of uncompressed data for a single zarr shard of every array inside every
+            table group. If `None` (default), no shard budget is requested and the write is left entirely to the
+            backend defaults. Requires a zarr v3 table format, zarr >= 3.1.6 and an anndata that supports zarr v3
+            auto-sharding; a `TableWriteOptionsError` is raised otherwise, before anything is written to disk.
+            Setting it forces anndata's zarr v3 auto-sharding on for the duration of each table write, overriding
+            an explicit :attr:`anndata.settings.auto_shard_zarr_v3` of `False`; there is no value that turns
+            sharding off. The budget is a target, not a bound: if it is below the automatically chosen inner
+            chunk (about 1 MiB), it degenerates to one chunk per shard, which is larger than the budget.
+            Non-table elements ignore this parameter.
+
+        Raises
+        ------
+        TableWriteOptionsError
+            If `table_shard_size_bytes` is not a positive `int`, or if the active backend cannot honour a shard
+            budget.
+        """
+        from spatialdata._io._utils import (
+            _resolve_zarr_store,
+            _validate_compressor_args,
+            _validate_table_shard_size_bytes,
+        )
         from spatialdata._io.format import _parse_formats
 
         parsed = _parse_formats(sdata_formats)
         _validate_compressor_args(raster_compressor)
+        _validate_table_shard_size_bytes(table_shard_size_bytes, tables_zarr_format=parsed["tables"].zarr_format)
 
         if isinstance(file_path, str):
             file_path = Path(file_path)
@@ -1199,6 +1222,7 @@ def write(
                 shapes_geometry_encoding=shapes_geometry_encoding,
                 raster_compressor=raster_compressor,
                 convert_table_strings_to_categoricals=convert_table_strings_to_categoricals,
+                table_shard_size_bytes=table_shard_size_bytes,
             )
 
         if self.path != file_path and update_sdata_path:
@@ -1218,6 +1242,8 @@ def _write_element(
         shapes_geometry_encoding: Literal["WKB", "geoarrow"] | None = None,
         raster_compressor: dict[Literal["lz4", "zstd"], int] | None = None,
         convert_table_strings_to_categoricals: bool = False,
+        *,
+        table_shard_size_bytes: int | None = None,
     ) -> None:
         from spatialdata._io.io_zarr import _get_groups_for_element
 
@@ -1286,6 +1312,7 @@ def _write_element(
                 name=element_name,
                 element_format=parsed_formats["tables"],
                 convert_strings_to_categoricals=convert_table_strings_to_categoricals,
+                shard_size_bytes=table_shard_size_bytes,
             )
         else:
             raise ValueError(f"Unknown element type: {element_type}")
@@ -1298,6 +1325,8 @@ def write_element(
         shapes_geometry_encoding: Literal["WKB", "geoarrow"] | None = None,
         raster_compressor: dict[Literal["lz4", "zstd"], int] | None = None,
         convert_table_strings_to_categoricals: bool = False,
+        *,
+        table_shard_size_bytes: int | None = None,
     ) -> None:
         """
         Write a single element, or a list of elements, to the Zarr store used for backing.
@@ -1324,15 +1353,29 @@ def write_element(
         convert_table_strings_to_categoricals
             If True, and if element to be written is a table, convert string columns to categoricals before writing.
             Note that this will have a side effect of modifying string columns into categoricals in place.
+        table_shard_size_bytes
+            The target size in bytes of uncompressed data for a single zarr shard of every array inside the
+            table group. See :meth:`spatialdata.SpatialData.write` for the full semantics. Ignored if the
+            element being written is not a table.
+
+        Raises
+        ------
+        TableWriteOptionsError
+            If `table_shard_size_bytes` is not a positive `int`, or if the active backend cannot honour a shard
+            budget.
 
         Notes
         -----
         If you pass a list of names, the elements will be written one by one. If an error occurs during the writing of
         an element, the writing of the remaining elements will not be attempted.
         """
+        from spatialdata._io._utils import _validate_table_shard_size_bytes
         from spatialdata._io.format import _parse_formats
 
         parsed_formats = _parse_formats(formats=sdata_formats)
+        _validate_table_shard_size_bytes(
+            table_shard_size_bytes, tables_zarr_format=parsed_formats["tables"].zarr_format
+        )
 
         if isinstance(element_name, list):
             for name in element_name:
@@ -1344,6 +1387,7 @@ def write_element(
                     shapes_geometry_encoding=shapes_geometry_encoding,
                     raster_compressor=raster_compressor,
                     convert_table_strings_to_categoricals=convert_table_strings_to_categoricals,
+                    table_shard_size_bytes=table_shard_size_bytes,
                 )
             return
 
@@ -1381,6 +1425,7 @@ def write_element(
             shapes_geometry_encoding=shapes_geometry_encoding,
             raster_compressor=raster_compressor,
             convert_table_strings_to_categoricals=convert_table_strings_to_categoricals,
+            table_shard_size_bytes=table_shard_size_bytes,
         )
         # After every write, metadata should be consolidated, otherwise this can lead to IO problems like when deleting.
         if self.has_consolidated_metadata():
diff --git a/src/spatialdata/_io/_utils.py b/src/spatialdata/_io/_utils.py
index fa5af1dd7..d69e1d216 100644
--- a/src/spatialdata/_io/_utils.py
+++ b/src/spatialdata/_io/_utils.py
@@ -11,15 +11,18 @@
 from contextlib import contextmanager
 from enum import Enum
 from functools import singledispatch
+from importlib.metadata import version
 from pathlib import Path
 from typing import Any, Literal
 
+import anndata as ad
 import zarr
 from anndata import AnnData
 from dask._task_spec import Task
 from dask.array import Array as DaskArray
 from dask.dataframe import DataFrame as DaskDataFrame
 from geopandas import GeoDataFrame
+from packaging.version import Version
 from upath import UPath
 from upath.implementations.local import PosixUPath, WindowsUPath
 from xarray import DataArray, DataTree
@@ -597,3 +600,106 @@ def _validate_compressor_args(compressor_dict: dict[Literal["lz4", "zstd"], int]
             )
         if not isinstance(value := list(compressor_dict.values())[0], int) or not (0 <= value <= 9):
             raise ValueError(f"The compression level must be an integer inclusive between 0 and 9. Got: {value}")
+
+
+_MIN_ZARR_FOR_SHARD_BUDGET = Version("3.1.6")
+
+
+def _validate_table_shard_size_bytes(table_shard_size_bytes: int | None, tables_zarr_format: int | None = None) -> None:
+    """Validate `table_shard_size_bytes` against the argument itself and against the active backend.
+
+    Parameters
+    ----------
+    table_shard_size_bytes
+        The requested target size in bytes of uncompressed data for a single zarr shard. `None` disables the
+        validation entirely, since nothing will be requested from the backend.
+    tables_zarr_format
+        The zarr format of the table element format that will be used for the write, if already known. Sharding does
+        not exist in zarr format 2, so a budget cannot be honoured there.
+
+    Raises
+    ------
+    TableWriteOptionsError
+        If the value is not a positive `int`, or if zarr, anndata or the table format cannot honour a shard budget.
+    """
+    if table_shard_size_bytes is None:
+        return
+
+    from spatialdata._io.exceptions import TableWriteOptionsError
+
+    if isinstance(table_shard_size_bytes, bool) or not isinstance(table_shard_size_bytes, int):
+        raise TableWriteOptionsError(
+            f"`table_shard_size_bytes` must be a positive int, got {table_shard_size_bytes!r}."
+        )
+    if table_shard_size_bytes <= 0:
+        raise TableWriteOptionsError(
+            f"`table_shard_size_bytes` must be a positive int, got {table_shard_size_bytes!r}."
+        )
+
+    zarr_version = Version(version("zarr"))
+    if zarr_version < _MIN_ZARR_FOR_SHARD_BUDGET:
+        raise TableWriteOptionsError(
+            f"`table_shard_size_bytes` requires zarr >= {_MIN_ZARR_FOR_SHARD_BUDGET}, got {zarr_version}. zarr 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 zarr-python#3603), which would put ~130k inner chunks "
+            "in a 128 MiB shard."
+        )
+
+    # `anndata.settings` did not exist before anndata 0.10, and `pyproject.toml` pins anndata>=0.9.1, so the attribute
+    # has to be looked up defensively rather than assumed to be there.
+    settings_obj = getattr(ad, "settings", None)
+    if settings_obj is None or not hasattr(settings_obj, "auto_shard_zarr_v3"):
+        raise TableWriteOptionsError(
+            "`table_shard_size_bytes` requires an anndata that supports zarr v3 auto-sharding, got "
+            f"{version('anndata')}."
+        )
+
+    if tables_zarr_format == 2:
+        raise TableWriteOptionsError(
+            "`table_shard_size_bytes` requires a zarr v3 table format, but the table format in use has "
+            "zarr_format=2. Sharding does not exist in zarr format 2."
+        )
+
+
+@contextmanager
+def _table_shard_budget(shard_size_bytes: int | None) -> Generator[None, None, None]:
+    """Scope a zarr shard budget and anndata's zarr v3 auto-sharding around a single table write.
+
+    Nothing is passed into anndata's `dataset_kwargs`. Instead two process globals are set for the duration of one
+    table write: zarr's `array.target_shard_size_bytes`, and anndata's `zarr_write_format` and `auto_shard_zarr_v3`
+    settings. anndata then injects `shards="auto"` itself, at the writers where that is safe, and yields to the budget
+    set here instead of installing its own 1 GB default. zarr derives the shard shape from the chunk shape, so
+    `shard % chunk == 0` and `shard <= array` hold by construction at every rank, length and dtype.
+
+    `shards` deliberately never reaches `dataset_kwargs`. A `shards` entry there would be forwarded to the rank-0
+    string scalars every SpatialData table carries in `uns/spatialdata_attrs`, and zarr's
+    `_guess_num_chunks_per_axis_shard` does not terminate on a rank-0 array while `array.target_shard_size_bytes` is
+    set (zarr-developers/zarr-python#4304). Today the rank-0 arrays are kept away from that code path by two
+    independent anndata mechanisms: `write_scalar_zarr`/`write_null_zarr` never call `zarr_v3_sharding` at all, and
+    `@zero_dim_array_as_scalar` re-dispatches 0-d ndarrays before `write_basic`'s sharding is reached.
+
+    `zarr_write_format` is overridden alongside the sharding setting because `AnnData.write_zarr` opens the group with
+    `mode="w"` and `zarr_format=settings.zarr_write_format`, i.e. it destroys and recreates the group spatialdata just
+    made; leaving that setting at 2 would silently produce a zarr v2 table group and no sharding at all.
+
+    Parameters
+    ----------
+    shard_size_bytes
+        The target size in bytes of uncompressed data for a single zarr shard. If `None`, nothing is set and this
+        context manager is a no-op.
+
+    Yields
+    ------
+    None
+    """
+    if shard_size_bytes is None:
+        yield
+        return
+
+    # `override` is order-preserving in both anndata implementations, so the zarr write format is set before the
+    # sharding setting, which is required because sharding cannot be enabled while the write format is 2.
+    with (
+        zarr.config.set({"array.target_shard_size_bytes": shard_size_bytes}),
+        ad.settings.override(zarr_write_format=3, auto_shard_zarr_v3=True),
+    ):
+        yield
diff --git a/src/spatialdata/_io/exceptions.py b/src/spatialdata/_io/exceptions.py
index 66f5802b7..398000cd8 100644
--- a/src/spatialdata/_io/exceptions.py
+++ b/src/spatialdata/_io/exceptions.py
@@ -24,3 +24,7 @@ class WritingToZarrV2DeprecationWarning(DeprecationWarning):
         "and will be removed in a future version. "
         "Please consider writing to zarr v3."
     )
+
+
+class TableWriteOptionsError(ValueError):
+    """Exception raised when table write options cannot be honoured by the active backend."""
diff --git a/src/spatialdata/_io/io_table.py b/src/spatialdata/_io/io_table.py
index d7795015d..f3628b265 100644
--- a/src/spatialdata/_io/io_table.py
+++ b/src/spatialdata/_io/io_table.py
@@ -12,7 +12,7 @@
 from ome_zarr.format import Format
 from packaging.version import Version
 
-from spatialdata._io._utils import _resolve_zarr_store
+from spatialdata._io._utils import _resolve_zarr_store, _table_shard_budget, _validate_table_shard_size_bytes
 from spatialdata._io.exceptions import FormatVersionUnknownError, WritingToZarrV2DeprecationWarning
 from spatialdata._io.format import (
     CurrentTablesFormat,
@@ -61,6 +61,8 @@ def write_table(
     group_type: str = "ngff:regions_table",
     element_format: Format = CurrentTablesFormat(),
     convert_strings_to_categoricals: bool = False,
+    *,
+    shard_size_bytes: int | None = None,
 ) -> None:
     """
     Write a table to a Zarr store.
@@ -80,7 +82,18 @@ def write_table(
     convert_strings_to_categoricals
         If True, convert string columns to categoricals before writing.
         Note that this will have a side effect of modifying dtypes of the input table in place.
+    shard_size_bytes
+        The target size in bytes of uncompressed data for a single zarr shard of every array of the table. If `None`
+        (default), no shard budget is requested and the write is left entirely to the backend defaults. Requires a
+        zarr v3 table format, zarr >= 3.1.6 and an anndata that supports zarr v3 auto-sharding.
+
+    Raises
+    ------
+    TableWriteOptionsError
+        If `shard_size_bytes` is not a positive `int`, or if the active backend cannot honour a shard budget.
     """
+    _validate_table_shard_size_bytes(shard_size_bytes, tables_zarr_format=element_format.zarr_format)
+
     if element_format.zarr_format == 2:
         warnings.warn(
             message=WritingToZarrV2DeprecationWarning.message, category=WritingToZarrV2DeprecationWarning, stacklevel=2
@@ -98,26 +111,27 @@ def write_table(
     if element_format not in TablesFormats.values():
         raise FormatVersionUnknownError(element_type="table", version_encountered=element_format)
 
-    if element_format.zarr_format == 3 and Version(version("anndata")) >= Version("0.13"):
-        # `write_zarr` in anndata v0.13 and above can only write to zarr v3
-        # solution of passing resolved store directly roughly based on:
-        # https://github.com/scverse/anndata/issues/1548#issuecomment-2199801855
+    with _table_shard_budget(shard_size_bytes):
+        if element_format.zarr_format == 3 and Version(version("anndata")) >= Version("0.13"):
+            # `write_zarr` in anndata v0.13 and above can only write to zarr v3
+            # solution of passing resolved store directly roughly based on:
+            # https://github.com/scverse/anndata/issues/1548#issuecomment-2199801855
 
-        # resolve the store from the group
-        resolved_store = _resolve_zarr_store(table_group)
+            # resolve the store from the group
+            resolved_store = _resolve_zarr_store(table_group)
 
-        # Write the table to the path of the table group
-        table.write_zarr(
-            store=resolved_store,
-            consolidate_metadata=False,
-            convert_strings_to_categoricals=convert_strings_to_categoricals,
-        )
+            # Write the table to the path of the table group
+            table.write_zarr(
+                store=resolved_store,
+                consolidate_metadata=False,
+                convert_strings_to_categoricals=convert_strings_to_categoricals,
+            )
 
-    else:
-        if convert_strings_to_categoricals:
-            table.strings_to_categoricals()
+        else:
+            if convert_strings_to_categoricals:
+                table.strings_to_categoricals()
 
-        write_adata(group, name, table)
+            write_adata(group, name, table)
 
     # Re-fetch the group before setting the attributes below: the handle obtained above was cached while the group
     # was still empty, and zarr writes attributes as a whole document based on the handle's cached view, so writing
diff --git a/tests/io/test_readwrite.py b/tests/io/test_readwrite.py
index c018d887d..59a8b8393 100644
--- a/tests/io/test_readwrite.py
+++ b/tests/io/test_readwrite.py
@@ -8,6 +8,7 @@
 from pathlib import Path
 from typing import Any, Literal
 
+import anndata as ad
 import dask.array as da
 import dask.dataframe as dd
 import numpy as np
@@ -20,13 +21,14 @@
 from numpy.random import default_rng
 from packaging.version import Version
 from pandas.testing import assert_series_equal
+from scipy import sparse
 from shapely import MultiPolygon, Polygon
 from upath import UPath
 from xarray import DataArray
 from zarr.errors import GroupNotFoundError
 
 import spatialdata.config
-from spatialdata import SpatialData, deepcopy, read_zarr
+from spatialdata import SpatialData, TableWriteOptionsError, deepcopy, read_zarr
 from spatialdata._core.validation import ValidationError
 from spatialdata._io._utils import _are_directories_identical, get_dask_backing_files
 from spatialdata._io.format import (
@@ -37,7 +39,7 @@
 )
 from spatialdata._io.io_raster import write_image
 from spatialdata.datasets import blobs
-from spatialdata.models import Image2DModel
+from spatialdata.models import Image2DModel, TableModel
 from spatialdata.models._utils import get_channel_names
 from spatialdata.testing import assert_spatial_data_objects_are_identical
 from spatialdata.transformations.operations import (
@@ -688,13 +690,29 @@ def test_incremental_io_in_memory(
             sdata["poly"] = v
 
 
-def test_table_group_keeps_anndata_encoding_metadata(tmp_path: str, table_single_annotation: SpatialData) -> None:
+@pytest.mark.parametrize(
+    "table_shard_size_bytes",
+    [
+        None,
+        pytest.param(
+            2 * 1024 * 1024,
+            marks=pytest.mark.skipif(
+                Version(version("zarr")) < Version("3.1.6"),
+                reason="`array.target_shard_size_bytes` only sizes the inner chunk correctly from zarr 3.1.6 on",
+            ),
+        ),
+    ],
+)
+def test_table_group_keeps_anndata_encoding_metadata(
+    tmp_path: str, table_single_annotation: SpatialData, table_shard_size_bytes: int | None
+) -> None:
     # https://github.com/scverse/spatialdata/issues/1183
     # Writing the spatialdata attributes on the table group must not erase the
     # `encoding-type`/`encoding-version` metadata that anndata writes on the same
-    # group; anndata-level readers (read_elem, read_lazy) dispatch on it.
+    # group; anndata-level readers (read_elem, read_lazy) dispatch on it. A shard budget must not change that: it is
+    # scoped around the array writes only, and released before the attributes are written.
     tmpdir = Path(tmp_path) / "tmp.zarr"
-    table_single_annotation.write(tmpdir)
+    table_single_annotation.write(tmpdir, table_shard_size_bytes=table_shard_size_bytes)
 
     on_disk = json.loads((tmpdir / "tables" / "table" / "zarr.json").read_text())["attributes"]
     assert on_disk["encoding-type"] == "anndata"
@@ -1368,3 +1386,226 @@ def test_sdata_with_nan_in_obs(tmp_path: Path, convert_strings_to_categoricals:
             assert pd.isna(r1.iloc[1])
         else:
             assert r1.iloc[1] == "nan"
+
+
+SHARD_BUDGET_SMALL = 512 * 1024
+SHARD_BUDGET_LARGE = 2 * 1024 * 1024
+SDATA_FORMATS_ZARR_V3 = [f for f in SDATA_FORMATS if f.zarr_format == 3]
+SDATA_FORMATS_ZARR_V2 = [f for f in SDATA_FORMATS if f.zarr_format == 2]
+requires_shard_budget_support = pytest.mark.skipif(
+    Version(version("zarr")) < Version("3.1.6"),
+    reason="`array.target_shard_size_bytes` only sizes the inner chunk correctly from zarr 3.1.6 on",
+)
+
+
+def _shard_table(region: str | list[str] = "labels2d") -> AnnData:
+    """Build a table large enough that the shard budget measurably changes the on-disk geometry."""
+    n_obs, n_var = 4000, 2000
+    x = sparse.random(n_obs, n_var, density=0.05, format="csr", random_state=0, dtype=np.float64)
+    obs = pd.DataFrame(index=[f"cell_{i}" for i in range(n_obs)])
+    obs["instance_id"] = np.arange(n_obs)
+    obs["region"] = pd.Categorical(
+        [region] * n_obs if isinstance(region, str) else RNG.choice(region, size=n_obs).tolist()
+    )
+    adata = AnnData(X=x, obs=obs, var=pd.DataFrame(index=[f"gene_{i}" for i in range(n_var)]))
+    adata.obsm["spatial"] = RNG.normal(size=(n_obs, 2))
+    return TableModel.parse(adata, region=region, region_key="region", instance_key="instance_id")
+
+
+@pytest.fixture(scope="module")
+def shard_table() -> AnnData:
+    return _shard_table()
+
+
+def _write_shard_table(
+    path: Path,
+    table: AnnData,
+    sdata_container_format: SpatialDataContainerFormatType,
+    **kwargs: Any,
+) -> None:
+    SpatialData(tables={"table": table}).write(path, sdata_formats=sdata_container_format, **kwargs)
+
+
+def _table_only_sdata() -> SpatialData:
+    """A small table-only object: writing a labels-backed one twice trips a Windows file lock unrelated to this."""
+    return SpatialData(tables={"table": _get_table(region="labels2d")})
+
+
+def _x_data_array(path: Path) -> zarr.Array:
+    array = zarr.open_group(path, mode="r")["tables"]["table"]["X"]["data"]
+    assert isinstance(array, zarr.Array)
+    return array
+
+
+def _global_shard_state() -> tuple[Any, Any, Any]:
+    return (
+        zarr.config.get("array.target_shard_size_bytes", None),
+        ad.settings.auto_shard_zarr_v3,
+        ad.settings.zarr_write_format,
+    )
+
+
+@requires_shard_budget_support
+@pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
+@pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS_ZARR_V3)
+def test_table_shard_size_bytes_bounds_shard_size(
+    tmp_path: Path,
+    shard_table: AnnData,
+    sdata_container_format: SpatialDataContainerFormatType,
+) -> None:
+    arrays = {}
+    for label, budget in (("small", SHARD_BUDGET_SMALL), ("large", SHARD_BUDGET_LARGE)):
+        path = tmp_path / f"{label}.zarr"
+        _write_shard_table(path, shard_table, sdata_container_format, table_shard_size_bytes=budget)
+
+        array = _x_data_array(path)
+        assert array.shards is not None
+        assert array.shards[0] % array.chunks[0] == 0
+        assert array.shards[0] <= array.shape[0]
+        chunk_bytes = array.chunks[0] * array.dtype.itemsize
+        shard_bytes = array.shards[0] * array.dtype.itemsize
+        # the budget is a target, not a bound: when it is below the automatically chosen inner chunk, the shard
+        # degenerates to a single chunk, which is allowed to exceed the budget
+        assert shard_bytes <= max(budget, chunk_bytes)
+        arrays[label] = array
+
+    # the actual lever: a smaller budget must produce a smaller shard
+    assert arrays["small"].shards[0] < arrays["large"].shards[0]
+
+
+@requires_shard_budget_support
+@pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
+@pytest.mark.parametrize("region", ["labels2d", ["labels2d", "labels3d"]])
+@pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS_ZARR_V3)
+def test_table_scalars_are_never_sharded(
+    tmp_path: Path,
+    region: str | list[str],
+    sdata_container_format: SpatialDataContainerFormatType,
+) -> None:
+    # zarr cannot derive a shard shape for a rank-0 array while a shard budget is set: it loops forever
+    # (https://github.com/zarr-developers/zarr-python/issues/4304). Every SpatialData table carries rank-0 string
+    # scalars in `uns/spatialdata_attrs`, so this asserts none of them ever reaches that code path.
+    path = tmp_path / "scalars.zarr"
+    _write_shard_table(path, _shard_table(region), sdata_container_format, table_shard_size_bytes=SHARD_BUDGET_SMALL)
+
+    scalars = []
+
+    def collect(group: zarr.Group, prefix: str) -> None:
+        for key, member in group.members():
+            member_path = f"{prefix}/{key}"
+            if isinstance(member, zarr.Array):
+                if member.shape == ():
+                    scalars.append((member_path, member.shards))
+            else:
+                collect(member, member_path)
+
+    collect(zarr.open_group(path, mode="r")["tables"]["table"], "tables/table")
+
+    # `region_key` and `instance_key` are always rank-0; `region` only is when a single region is annotated
+    assert any(name.endswith("/region_key") for name, _ in scalars)
+    assert [(name, shards) for name, shards in scalars if shards is not None] == []
+
+
+@pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
+@pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS_ZARR_V3)
+def test_no_shards_key_reaches_anndata(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    sdata_container_format: SpatialDataContainerFormatType,
+) -> None:
+    # a fast guard for the same non-termination: passing `shards` down to anndata would hang the write rather than
+    # fail it, and a hang under `pytest -n auto` stalls a worker with no diagnostic
+    captured: list[dict[str, Any]] = []
+
+    def record_write_adata(group: zarr.Group, name: str, table: AnnData, **kwargs: Any) -> None:
+        captured.append(kwargs)
+
+    def record_write_zarr(self: AnnData, *args: Any, **kwargs: Any) -> None:
+        captured.append(kwargs)
+
+    monkeypatch.setattr("spatialdata._io.io_table.write_adata", record_write_adata)
+    monkeypatch.setattr(AnnData, "write_zarr", record_write_zarr)
+
+    _table_only_sdata().write(
+        tmp_path / "data.zarr",
+        sdata_formats=sdata_container_format,
+        table_shard_size_bytes=SHARD_BUDGET_SMALL,
+    )
+
+    assert len(captured) == 1
+    for kwargs in captured:
+        assert "shards" not in kwargs
+        assert "shards" not in kwargs.get("dataset_kwargs", {})
+
+
+@pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
+@pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS)
+def test_table_shard_size_bytes_none_is_unchanged(
+    tmp_path: Path,
+    sdata_container_format: SpatialDataContainerFormatType,
+) -> None:
+    sdata = _table_only_sdata()
+    absent = tmp_path / "absent.zarr"
+    explicit_none = tmp_path / "explicit_none.zarr"
+    sdata.write(absent, sdata_formats=sdata_container_format, update_sdata_path=False)
+    sdata.write(
+        explicit_none, sdata_formats=sdata_container_format, update_sdata_path=False, table_shard_size_bytes=None
+    )
+
+    assert _are_directories_identical(absent, explicit_none)
+
+
+@requires_shard_budget_support
+@pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
+def test_table_shard_budget_restores_global_state(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    before = _global_shard_state()
+
+    _table_only_sdata().write(tmp_path / "ok.zarr", table_shard_size_bytes=SHARD_BUDGET_SMALL)
+    assert _global_shard_state() == before
+
+    def boom(*args: Any, **kwargs: Any) -> None:
+        raise RuntimeError("write failed")
+
+    monkeypatch.setattr("spatialdata._io.io_table.write_adata", boom)
+    monkeypatch.setattr(AnnData, "write_zarr", boom)
+
+    with pytest.raises(RuntimeError, match="write failed"):
+        _table_only_sdata().write(tmp_path / "boom.zarr", table_shard_size_bytes=SHARD_BUDGET_SMALL)
+    assert _global_shard_state() == before
+
+
+@pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
+@pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS_ZARR_V2)
+def test_table_shard_size_bytes_rejected_on_zarr_v2(
+    tmp_path: Path,
+    sdata_container_format: SpatialDataContainerFormatType,
+) -> None:
+    path = tmp_path / "v2.zarr"
+    with pytest.raises(TableWriteOptionsError, match="requires a zarr v3 table format"):
+        _table_only_sdata().write(
+            path,
+            sdata_formats=sdata_container_format,
+            table_shard_size_bytes=SHARD_BUDGET_SMALL,
+        )
+    # the argument is validated before any element reaches disk
+    assert not path.exists()
+
+
+@pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
+@pytest.mark.parametrize("invalid", [0, -1, 1.5, True])
+def test_table_shard_size_bytes_validation(tmp_path: Path, invalid: Any) -> None:
+    sdata = _table_only_sdata()
+    path = tmp_path / "invalid.zarr"
+    with pytest.raises(TableWriteOptionsError, match="must be a positive int"):
+        sdata.write(path, table_shard_size_bytes=invalid)
+    assert not path.exists()
+
+    backing = tmp_path / "backed.zarr"
+    sdata.write(backing)
+    sdata["table2"] = deepcopy(sdata["table"])
+    with pytest.raises(TableWriteOptionsError, match="must be a positive int"):
+        sdata.write_element("table2", table_shard_size_bytes=invalid)
+    assert not (backing / "tables" / "table2").exists()

From 9a2e82f3cdc60ae834fb268bc05c98d5d1f46e2c Mon Sep 17 00:00:00 2001
From: Tomatokeftes <129113023+Tomatokeftes@users.noreply.github.com>
Date: Tue, 1 Sep 2026 15:08:43 +0200
Subject: [PATCH 2/4] test: skip the shard budget tests when the backend cannot
 honour one

The skip condition now also covers anndata, so a leg without zarr v3
auto-sharding support skips instead of failing on an AttributeError or on
the wrong validation message. Applied to the two tests that push a budget
through validation without sharding anything, and reused by the issue #1183
guard, so all the shard tests share one predicate.
---
 tests/io/test_readwrite.py | 31 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/tests/io/test_readwrite.py b/tests/io/test_readwrite.py
index 59a8b8393..736ee9784 100644
--- a/tests/io/test_readwrite.py
+++ b/tests/io/test_readwrite.py
@@ -58,6 +58,17 @@
 
 RNG = default_rng(0)
 SDATA_FORMATS = list(SpatialDataContainerFormats.values())
+SHARD_BUDGET_SMALL = 512 * 1024
+SHARD_BUDGET_LARGE = 2 * 1024 * 1024
+SDATA_FORMATS_ZARR_V3 = [f for f in SDATA_FORMATS if f.zarr_format == 3]
+SDATA_FORMATS_ZARR_V2 = [f for f in SDATA_FORMATS if f.zarr_format == 2]
+requires_shard_budget_support = pytest.mark.skipif(
+    Version(version("zarr")) < Version("3.1.6") or not hasattr(getattr(ad, "settings", None), "auto_shard_zarr_v3"),
+    reason=(
+        "a shard budget needs zarr >= 3.1.6, which is the first release to size the inner chunk correctly, and an "
+        "anndata that supports zarr v3 auto-sharding"
+    ),
+)
 
 
 @pytest.mark.filterwarnings("ignore:SpatialData is not stored in the most current format:UserWarning")
@@ -694,13 +705,7 @@ def test_incremental_io_in_memory(
     "table_shard_size_bytes",
     [
         None,
-        pytest.param(
-            2 * 1024 * 1024,
-            marks=pytest.mark.skipif(
-                Version(version("zarr")) < Version("3.1.6"),
-                reason="`array.target_shard_size_bytes` only sizes the inner chunk correctly from zarr 3.1.6 on",
-            ),
-        ),
+        pytest.param(SHARD_BUDGET_LARGE, marks=requires_shard_budget_support),
     ],
 )
 def test_table_group_keeps_anndata_encoding_metadata(
@@ -1388,16 +1393,6 @@ def test_sdata_with_nan_in_obs(tmp_path: Path, convert_strings_to_categoricals:
             assert r1.iloc[1] == "nan"
 
 
-SHARD_BUDGET_SMALL = 512 * 1024
-SHARD_BUDGET_LARGE = 2 * 1024 * 1024
-SDATA_FORMATS_ZARR_V3 = [f for f in SDATA_FORMATS if f.zarr_format == 3]
-SDATA_FORMATS_ZARR_V2 = [f for f in SDATA_FORMATS if f.zarr_format == 2]
-requires_shard_budget_support = pytest.mark.skipif(
-    Version(version("zarr")) < Version("3.1.6"),
-    reason="`array.target_shard_size_bytes` only sizes the inner chunk correctly from zarr 3.1.6 on",
-)
-
-
 def _shard_table(region: str | list[str] = "labels2d") -> AnnData:
     """Build a table large enough that the shard budget measurably changes the on-disk geometry."""
     n_obs, n_var = 4000, 2000
@@ -1506,6 +1501,7 @@ def collect(group: zarr.Group, prefix: str) -> None:
     assert [(name, shards) for name, shards in scalars if shards is not None] == []
 
 
+@requires_shard_budget_support
 @pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
 @pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS_ZARR_V3)
 def test_no_shards_key_reaches_anndata(
@@ -1577,6 +1573,7 @@ def boom(*args: Any, **kwargs: Any) -> None:
     assert _global_shard_state() == before
 
 
+@requires_shard_budget_support
 @pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
 @pytest.mark.parametrize("sdata_container_format", SDATA_FORMATS_ZARR_V2)
 def test_table_shard_size_bytes_rejected_on_zarr_v2(

From c74ccee3f0a9f630fbac13f224eb1a0544f7b579 Mon Sep 17 00:00:00 2001
From: Tomatokeftes <129113023+Tomatokeftes@users.noreply.github.com>
Date: Wed, 2 Sep 2026 11:35:57 +0200
Subject: [PATCH 3/4] refactor: drop the runtime type check and shorten the
 shard budget docstring

Review follow-up on #1199. `_validate_table_shard_size_bytes` no longer checks
the type of `table_shard_size_bytes`, only that it is positive, so the `1.5`
and `True` cases leave the validation test. The `_table_shard_budget` docstring
is cut to what is needed at the call site; the rank-0 rationale now lives in the
PR description.
---
 src/spatialdata/_io/_utils.py | 32 +++++++-------------------------
 tests/io/test_readwrite.py    |  4 ++--
 2 files changed, 9 insertions(+), 27 deletions(-)

diff --git a/src/spatialdata/_io/_utils.py b/src/spatialdata/_io/_utils.py
index d69e1d216..6933339f6 100644
--- a/src/spatialdata/_io/_utils.py
+++ b/src/spatialdata/_io/_utils.py
@@ -627,10 +627,6 @@ def _validate_table_shard_size_bytes(table_shard_size_bytes: int | None, tables_
 
     from spatialdata._io.exceptions import TableWriteOptionsError
 
-    if isinstance(table_shard_size_bytes, bool) or not isinstance(table_shard_size_bytes, int):
-        raise TableWriteOptionsError(
-            f"`table_shard_size_bytes` must be a positive int, got {table_shard_size_bytes!r}."
-        )
     if table_shard_size_bytes <= 0:
         raise TableWriteOptionsError(
             f"`table_shard_size_bytes` must be a positive int, got {table_shard_size_bytes!r}."
@@ -665,32 +661,18 @@ def _validate_table_shard_size_bytes(table_shard_size_bytes: int | None, tables_
 def _table_shard_budget(shard_size_bytes: int | None) -> Generator[None, None, None]:
     """Scope a zarr shard budget and anndata's zarr v3 auto-sharding around a single table write.
 
-    Nothing is passed into anndata's `dataset_kwargs`. Instead two process globals are set for the duration of one
-    table write: zarr's `array.target_shard_size_bytes`, and anndata's `zarr_write_format` and `auto_shard_zarr_v3`
-    settings. anndata then injects `shards="auto"` itself, at the writers where that is safe, and yields to the budget
-    set here instead of installing its own 1 GB default. zarr derives the shard shape from the chunk shape, so
-    `shard % chunk == 0` and `shard <= array` hold by construction at every rank, length and dtype.
+    Sets zarr's `array.target_shard_size_bytes` and overrides anndata's `zarr_write_format=3` and
+    `auto_shard_zarr_v3=True` for the duration of the block, restoring all three on exit. anndata then injects
+    `shards="auto"` itself where that is safe and uses this budget instead of its own 1 GB default. `zarr_write_format`
+    is overridden too because `AnnData.write_zarr` recreates the group with that setting, and at 2 no sharding happens.
 
-    `shards` deliberately never reaches `dataset_kwargs`. A `shards` entry there would be forwarded to the rank-0
-    string scalars every SpatialData table carries in `uns/spatialdata_attrs`, and zarr's
-    `_guess_num_chunks_per_axis_shard` does not terminate on a rank-0 array while `array.target_shard_size_bytes` is
-    set (zarr-developers/zarr-python#4304). Today the rank-0 arrays are kept away from that code path by two
-    independent anndata mechanisms: `write_scalar_zarr`/`write_null_zarr` never call `zarr_v3_sharding` at all, and
-    `@zero_dim_array_as_scalar` re-dispatches 0-d ndarrays before `write_basic`'s sharding is reached.
-
-    `zarr_write_format` is overridden alongside the sharding setting because `AnnData.write_zarr` opens the group with
-    `mode="w"` and `zarr_format=settings.zarr_write_format`, i.e. it destroys and recreates the group spatialdata just
-    made; leaving that setting at 2 would silently produce a zarr v2 table group and no sharding at all.
+    `shards` is deliberately never put into `dataset_kwargs`: it would reach the rank-0 scalars in `uns` and hang zarr
+    (https://github.com/zarr-developers/zarr-python/issues/4304).
 
     Parameters
     ----------
     shard_size_bytes
-        The target size in bytes of uncompressed data for a single zarr shard. If `None`, nothing is set and this
-        context manager is a no-op.
-
-    Yields
-    ------
-    None
+        The target size in bytes of uncompressed data for a single zarr shard. If `None`, nothing is set.
     """
     if shard_size_bytes is None:
         yield
diff --git a/tests/io/test_readwrite.py b/tests/io/test_readwrite.py
index 736ee9784..20bb108ad 100644
--- a/tests/io/test_readwrite.py
+++ b/tests/io/test_readwrite.py
@@ -1592,8 +1592,8 @@ def test_table_shard_size_bytes_rejected_on_zarr_v2(
 
 
 @pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
-@pytest.mark.parametrize("invalid", [0, -1, 1.5, True])
-def test_table_shard_size_bytes_validation(tmp_path: Path, invalid: Any) -> None:
+@pytest.mark.parametrize("invalid", [0, -1])
+def test_table_shard_size_bytes_validation(tmp_path: Path, invalid: int) -> None:
     sdata = _table_only_sdata()
     path = tmp_path / "invalid.zarr"
     with pytest.raises(TableWriteOptionsError, match="must be a positive int"):

From bbffcac38b9044fdca2f46c7053907afe0897b5e Mon Sep 17 00:00:00 2001
From: Tomatokeftes <129113023+Tomatokeftes@users.noreply.github.com>
Date: Wed, 2 Sep 2026 11:41:13 +0200
Subject: [PATCH 4/4] fix: normalize the shard budget to an int before setting
 it

anndata only honours `array.target_shard_size_bytes` when it reads back as
an `int`, because of the `isinstance` check in `zarr_v3_sharding`. A float
failed that check and anndata quietly installed its own 1 GB default
instead, so the same numeric value behaved differently depending on its
type, with no error and no warning.

Measured on a 4000 x 2000 CSR table before this change: an int budget of
2 MiB gave `shards=(250000,)`, while the identical `2.0 * 1024**2` and a
plain `1e8` both gave `shards=(400000,)`, the whole array as one shard.
`1e8` is a natural way to write a budget, so this was easy to hit.

Dropping the runtime type check was right, but it left this case silent.
Normalizing is not a type check: the signature stays honest and nothing is
raised. The new test asserts a float budget produces the same on-disk
geometry as the equivalent int.
---
 src/spatialdata/_io/_utils.py |  6 +++++-
 tests/io/test_readwrite.py    | 15 +++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/spatialdata/_io/_utils.py b/src/spatialdata/_io/_utils.py
index 6933339f6..55d7839bd 100644
--- a/src/spatialdata/_io/_utils.py
+++ b/src/spatialdata/_io/_utils.py
@@ -678,10 +678,14 @@ def _table_shard_budget(shard_size_bytes: int | None) -> Generator[None, None, N
         yield
         return
 
+    # anndata only honours the budget when it reads back as an `int` (its `isinstance` check); a float would be
+    # dropped in favour of anndata's own 1 GB default, silently, so normalize instead of leaving that to chance.
+    budget = int(shard_size_bytes)
+
     # `override` is order-preserving in both anndata implementations, so the zarr write format is set before the
     # sharding setting, which is required because sharding cannot be enabled while the write format is 2.
     with (
-        zarr.config.set({"array.target_shard_size_bytes": shard_size_bytes}),
+        zarr.config.set({"array.target_shard_size_bytes": budget}),
         ad.settings.override(zarr_write_format=3, auto_shard_zarr_v3=True),
     ):
         yield
diff --git a/tests/io/test_readwrite.py b/tests/io/test_readwrite.py
index 20bb108ad..636830b9a 100644
--- a/tests/io/test_readwrite.py
+++ b/tests/io/test_readwrite.py
@@ -1468,6 +1468,21 @@ def test_table_shard_size_bytes_bounds_shard_size(
     assert arrays["small"].shards[0] < arrays["large"].shards[0]
 
 
+@requires_shard_budget_support
+@pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
+def test_table_shard_size_bytes_accepts_an_equivalent_float(tmp_path: Path, shard_table: AnnData) -> None:
+    # anndata only honours the budget when it reads back as an `int`, so an un-normalized float would be dropped in
+    # favour of anndata's own 1 GB default with no error and no warning; `1e8` is a natural way to write a budget
+    geometries = []
+    for label, budget in (("int", SHARD_BUDGET_LARGE), ("float", float(SHARD_BUDGET_LARGE))):
+        path = tmp_path / f"{label}.zarr"
+        _write_shard_table(path, shard_table, CurrentSpatialDataContainerFormat(), table_shard_size_bytes=budget)
+        array = _x_data_array(path)
+        geometries.append((array.chunks, array.shards))
+
+    assert geometries[0] == geometries[1]
+
+
 @requires_shard_budget_support
 @pytest.mark.filterwarnings("ignore:The table is annotating:UserWarning")
 @pytest.mark.parametrize("region", ["labels2d", ["labels2d", "labels3d"]])