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..55d7839bd 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,92 @@ 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 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. + + 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` 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. + """ + if shard_size_bytes is None: + 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": budget}), + 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..636830b9a 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 ( @@ -56,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") @@ -688,13 +701,23 @@ 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(SHARD_BUDGET_LARGE, marks=requires_shard_budget_support), + ], +) +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 +1391,233 @@ 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" + + +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") +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"]]) +@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] == [] + + +@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( + 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 + + +@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( + 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]) +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"): + 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()