Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Tomatokeftes wants to merge 4 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
dbece29
feat: expose a table shard size via `table_shard_size_bytes`
Tomatokeftes Sep 1, 2026
9a2e82f
test: skip the shard budget tests when the backend cannot honour one
Tomatokeftes Sep 1, 2026
c74ccee
refactor: drop the runtime type check and shorten the shard budget do…
Tomatokeftes Sep 2, 2026
bbffcac
fix: normalize the shard budget to an int before setting it
Tomatokeftes Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions src/spatialdata/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -119,6 +121,8 @@
"SpatialData",
# _io._utils
"get_dask_backing_files",
# _io.exceptions
"TableWriteOptionsError",
# _io.format
"SpatialDataFormatType",
# _io.io_zarr
Expand Down Expand Up @@ -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

Expand Down
49 changes: 47 additions & 2 deletions src/spatialdata/_core/spatialdata.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand All @@ -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:
Expand All @@ -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

Expand Down Expand Up @@ -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}")
Expand All @@ -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.
Expand All @@ -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:
Expand All @@ -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

Expand Down Expand Up @@ -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():
Expand Down
92 changes: 92 additions & 0 deletions src/spatialdata/_io/_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 "

@Tomaz-Vieira Tomaz-Vieira Sep 1, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@Tomatokeftes Tomatokeftes Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see the problem, and I think there might be a third option to bump the floors. This should remove rather than rename around it.

The two version checks only because pyproject.toml still says anndata>=0.9.1 and zarr.=3.0.0. The pieces this argument needs are older than they look:

  • anndata.settings.auto_shard_zarr_v3 exists since anndata 0.12.5 (anndata#2167)
  • the byte budget via zarr.config exists since anndata 0.12.14 (anndata#2427)
  • zarr 3.1.6 is the first release that sizes the inner chunk correctly (zarr-python#3603)

The test matrix here already only runs anndata 0.12 and 0.13, and dask is pinned at >=2026.3.0, so anndata>=0.12.14 and zarr>=3.1.6 would not be unusual. With those floors both version checks are deleted and the signature is honest on every supported install.

What remains is for zarr v2 tabe format, and I suggest we keep raising there, since it is not a dependency proble, but a user asking for a format without sharding, and it already carries a deprecation warning. It also raises at the top of write() before a single element is written, so nobodoy ends up with a half-written store.

On the two alternatives: I prefer not to do _hint, as a silently ignored budgest is exactly the failure that motivate the issue, a table landing as a few hundred thousand files with nobody told. If you prefer not to touch the pins, I can do the TableShardBudget object instead. I will then validate in __init__ and raise rather than a try_create that returns None, but that is a detail. Tell me which and I will push it.

@Tomatokeftes Tomatokeftes Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One caveat against my own suggestion, since it changes how much the floor bump actually buys.

I checked the versions against the tags rather than trusting the PR numbers: auto_shard_zarr_v3 is absent in anndata 0.12.4 and present in 0.12.5; target_shard_size_bytes is absent in 0.12.13 and present in 0.12.14; and the has_auto_shard_size yield-to-caller branch is intact at 0.12.14. So the floors themselves are right.

But the matrix pins anndata>=0.12,<0.13, which resolves to the newest 0.12.x, 0.12.19 today. An anndata>=0.12.14 floor would therefore never actually be exercised by CI. It would be a supported-version claim we do not test. That is true of whatever floor we pick rather than an argument against this one, and it does not change my preference, but you should weigh it now rather than hear it from me later.

Still happy to do TableShardBudget instead if you would rather not move the pins. Either way I would keep the zarr v2 case raising, since that is a user asking for a format without sharding rather than a dependency problem.

@Tomatokeftes Tomatokeftes Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small correction to my own caveat: the anndata>=0.12,<0.13 pin is in the hatch test-anndata-pandas matrix in pyproject.toml, and no workflow runs that matrix. The GitHub legs do uv sync --group=test, so every one of them resolves the newest anndata, 0.13.3.post0 today, and the prerelease leg installs git main. So anndata 0.12 is not exercised by CI at any version, and a floor bump neither adds nor removes coverage. Same conclusion, stated more precisely.

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
4 changes: 4 additions & 0 deletions src/spatialdata/_io/exceptions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
48 changes: 31 additions & 17 deletions src/spatialdata/_io/io_table.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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):

@Tomaz-Vieira Tomaz-Vieira Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@Tomatokeftes Tomatokeftes Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I checked this before, and here is the short version: for a byte budget it is the only way I could find, and it is the way anndata itself already uses.

  1. zarr has no per-call budget. create_array takes no such parameter. The only place the number is read is zarr.config.get("array.target_shard_size_bytes") in zarr/core/chunk_grids.py (line 916 on 3.3.0). This is global by zarr's design as far as I know.

  2. anndata already sets it globally on every default write. zarr_v3_sharding in anndata/_io/specs/methods.py wraps every array creation in zarr.config.set({"array.target_shard_size_bytes": 1_000_000_000}) whenever nothing is set (lines 139-143 on 0.13.3). So the context manager here does not add a new mechanism. It replaces anndata's hard-coded 1 GB with the user's number for the duration of one table and then hands back whatever was there before.

  3. Nesting is safe. donfig's ConfigSet records the previous values and restores them on exit, and anndata.settings.override does the same. If a caller is already inside a zarr.config.set, the innermost wins for one table and the outer value comes back afterwards. I think an explicit argument beats ambient config, so that is the behaviour I would want anyway.

  4. Threads might be a problem. zarr.config is one process-wide dict with a lock around mutation, not thread-local storage, so a table written on another thread during this scope would see this budget. That problem also exists today on anndata's default path in the same shape. This PR does not widen it.

The per-call alternative is an explicit shards tuple in dataset_kwargs, and that does not work for a table. anndata broadcasts one dataset_kwargs to every array of the AnnData, the ranks differ, and a shards entry reaching the 0-d scalars in uns makes zarr hang (zarr-python#4304, fix in zarr-python#4305).

The real fix for your concern is an argument on anndata's side, something like a shard budget on write_zarr / write_elem. It does not exist yet, but I can open that issue on anndata. If it lands, the signature here stays as it is and only the body of _table_shard_budget would change.

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
Expand Down
Loading