fix: warn for upcoming zarr autosharding + fix settings dependency (#2427)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Philipp A. <flying-sheep@web.de>

Author: 3 people

docs/release-notes/2427.fix.md

@@ -0,0 +1 @@
+{attr}`anndata.settings.auto_shard_zarr_v3` and {attr}`anndata.settings.zarr_write_format` are no longer dependent on each other, as stated in the docs {user}`ilan-gold`

docs/release-notes/2427.perf.md

@@ -0,0 +1 @@
+{attr}`anndata.settings.auto_shard_zarr_v3` now utilizes zarr's support for a target shard via `zarr.config.set({"array.target_shard_size_bytes" ...})` (only `zarr` version `>=3.1.4`) to make the target shard size 1GB (uncompressed) if not otherwise set. {user}`ilan-gold`

pyproject.toml

@@ -176,7 +176,6 @@ filterwarnings_when_strict = [
     "default:Consolidated metadata is:UserWarning",
     "default:.*Struct:zarr.core.dtype.common.UnstableSpecificationWarning",
     "default:.*FixedLengthUTF32:zarr.core.dtype.common.UnstableSpecificationWarning",
-    "default:Automatic shard shape inference is experimental",
     "default:Writing zarr v2:UserWarning",
     # TODO: Remove in conjunction with or before https://github.com/scverse/anndata/pull/1707
     "default:.*will obey copy-on-write semantics:FutureWarning",

src/anndata/_io/specs/methods.py

@@ -1,16 +1,21 @@
 from __future__ import annotations
 
+import warnings
 from collections.abc import Mapping, MutableMapping
+from contextlib import contextmanager, nullcontext
 from copy import copy
-from functools import partial
+from functools import partial, wraps
+from importlib.metadata import version
 from itertools import product
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Protocol
 
 import h5py
 import numpy as np
 import pandas as pd
+import zarr
 from numcodecs import VLenUTF8
+from packaging.version import Version
 from scipy import sparse
 from zarr.core.dtype import VariableLengthUTF8
 
@@ -21,6 +26,7 @@
 from anndata._core.merge import intersect_keys
 from anndata._core.sparse_dataset import _CSCDataset, _CSRDataset, sparse_dataset
 from anndata._io.utils import check_key, zero_dim_array_as_scalar
+from anndata._types import StorageType
 from anndata._warnings import OldFormatWarning
 from anndata.compat import (
     AwkArray,
@@ -45,16 +51,16 @@
 from .registry import _REGISTRY, IOSpec, read_elem, read_elem_partial
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
+    from collections.abc import Generator, Iterator
     from os import PathLike
     from typing import Any, Literal
 
     from numpy import typing as npt
     from numpy.typing import NDArray
 
-    from anndata._types import _ArrayStorageType, _GroupStorageType
+    from anndata._types import _ArrayStorageType, _GroupStorageType, _WriteInternal
     from anndata.compat import CSArray, CSMatrix, CupyCSMatrix
-    from anndata.typing import AxisStorable, _InMemoryArrayOrScalarType
+    from anndata.typing import AxisStorable, RWAble, _InMemoryArrayOrScalarType
 
     from .registry import Reader, Writer
 
@@ -112,10 +118,35 @@ def zarr_v3_compressor_compat(dataset_kwargs: dict) -> dict:
     return dataset_kwargs
 
 
-def zarr_v3_sharding(dataset_kwargs) -> dict:
-    if "shards" not in dataset_kwargs and ad.settings.auto_shard_zarr_v3:
+@contextmanager
+def zarr_v3_sharding(dataset_kwargs: dict, format: Literal[2, 3]) -> Generator[dict]:
+    auto_sharding = (
+        "shards" not in dataset_kwargs
+        and ad.settings.auto_shard_zarr_v3
+        and format == 3
+    )
+    if ad.settings.auto_shard_zarr_v3 is None and format == 3:
+        warn(
+            "zarr v3 autosharding will be the default in the next minor release.",
+            UserWarning,
+        )
+    elif auto_sharding:
         dataset_kwargs = {**dataset_kwargs, "shards": "auto"}
-    return dataset_kwargs
+    # Auto shard sizes are a relatively recent feature
+    supports_auto_shard_size = Version(version("zarr")) >= Version("3.1.4")
+    has_auto_shard_size = supports_auto_shard_size and isinstance(
+        zarr.config.get("array.target_shard_size_bytes"), int
+    )
+    # 1GB uncompressed shard size seems reasonable.
+    # Shards need to generally held completely in memory before writing.
+    # Even at a compression ration of 6x, that's still a ~20x improvement on number of files.
+    # Users can ovetrride this nonetheless, hence the above checks.
+    with (
+        zarr.config.set({"array.target_shard_size_bytes": 1_000_000_000})
+        if supports_auto_shard_size and not has_auto_shard_size and auto_sharding
+        else nullcontext()
+    ):
+        yield dataset_kwargs
 
 
 def _to_cpu_mem_wrapper(write_func):
@@ -140,6 +171,32 @@ def wrapper(
     return wrapper
 
 
+def suppress_autoshard_warning[S: StorageType, T: RWAble](
+    func: _WriteInternal[S, T],
+) -> _WriteInternal[S, T]:
+    @wraps(func)
+    def wrapper(
+        f: S,
+        k: str,
+        val: T,
+        *,
+        _writer: Writer,
+        dataset_kwargs: Mapping[str, Any] = MappingProxyType({}),
+    ):
+        with warnings.catch_warnings():
+            # Suppress warnings only if the user has opted into autosharding at the top level.
+            # If someone provides `shards` explicitly, then they should get the warning.
+            if ad.settings.auto_shard_zarr_v3 and "shards" not in dataset_kwargs:
+                warnings.filterwarnings(
+                    "ignore",
+                    r"Automatic shard shape inference is experimental",
+                    UserWarning,
+                )
+            return func(f, k, val, _writer=_writer, dataset_kwargs=dataset_kwargs)
+
+    return wrapper
+
+
 ################################
 # Fallbacks / backwards compat #
 ################################
@@ -277,6 +334,7 @@ def _read_partial(group, *, items=None, indices=(slice(None), slice(None))):
 
 @_REGISTRY.register_write(ZarrGroup, AnnData, IOSpec("anndata", "0.1.0"))
 @_REGISTRY.register_write(H5Group, AnnData, IOSpec("anndata", "0.1.0"))
+@suppress_autoshard_warning
 def write_anndata(
     f: _GroupStorageType,
     k: str,
@@ -329,6 +387,7 @@ def read_anndata(elem: _GroupStorageType | H5File, *, _reader: Reader) -> AnnDat
 
 @_REGISTRY.register_write(H5Group, Raw, IOSpec("raw", "0.1.0"))
 @_REGISTRY.register_write(ZarrGroup, Raw, IOSpec("raw", "0.1.0"))
+@suppress_autoshard_warning
 def write_raw(
     f: _GroupStorageType,
     k: str,
@@ -361,6 +420,7 @@ def write_null_h5py(f, k, _v, _writer, dataset_kwargs=MappingProxyType({})):
 
 
 @_REGISTRY.register_write(ZarrGroup, type(None), IOSpec("null", "0.1.0"))
+@suppress_autoshard_warning
 def write_null_zarr(f, k, _v, _writer, dataset_kwargs=MappingProxyType({})):
     dataset_kwargs = _remove_scalar_compression_args(dataset_kwargs)
     # zarr has no first-class null dataset
@@ -384,6 +444,7 @@ def read_mapping(
 
 @_REGISTRY.register_write(H5Group, dict, IOSpec("dict", "0.1.0"))
 @_REGISTRY.register_write(ZarrGroup, dict, IOSpec("dict", "0.1.0"))
+@suppress_autoshard_warning
 def write_mapping(
     f: _GroupStorageType,
     k: str,
@@ -404,6 +465,7 @@ def write_mapping(
 
 @_REGISTRY.register_write(H5Group, list, IOSpec("array", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, list, IOSpec("array", "0.2.0"))
+@suppress_autoshard_warning
 def write_list(
     f: _GroupStorageType,
     k: str,
@@ -425,6 +487,7 @@ def write_list(
 @_REGISTRY.register_write(ZarrGroup, np.ma.MaskedArray, IOSpec("array", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, ZarrArray, IOSpec("array", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, H5Array, IOSpec("array", "0.2.0"))
+@suppress_autoshard_warning
 @zero_dim_array_as_scalar
 def write_basic(
     f: _GroupStorageType,
@@ -441,8 +504,10 @@ def write_basic(
         f.create_dataset(k, data=elem, shape=elem.shape, dtype=dtype, **dataset_kwargs)
     else:
         dataset_kwargs = zarr_v3_compressor_compat(dataset_kwargs)
-        dataset_kwargs = zarr_v3_sharding(dataset_kwargs)
-        f.create_array(k, shape=elem.shape, dtype=dtype, **dataset_kwargs)
+        with zarr_v3_sharding(
+            dataset_kwargs, format=f.metadata.zarr_format
+        ) as dataset_kwargs:
+            f.create_array(k, shape=elem.shape, dtype=dtype, **dataset_kwargs)
         # see https://github.com/zarr-developers/zarr-python/discussions/2712
         if isinstance(elem, ZarrArray | H5Array):
             f[k][...] = elem[...]
@@ -499,14 +564,15 @@ def write_chunked_dense_array_to_group(
     _to_cpu_mem_wrapper(write_basic)
 )
 _REGISTRY.register_write(ZarrGroup, CupyArray, IOSpec("array", "0.2.0"))(
-    _to_cpu_mem_wrapper(write_basic)
+    suppress_autoshard_warning(_to_cpu_mem_wrapper(write_basic))
 )
 
 
 @_REGISTRY.register_write(ZarrGroup, views.DaskArrayView, IOSpec("array", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, DaskArray, IOSpec("array", "0.2.0"))
 @_REGISTRY.register_write(H5Group, views.DaskArrayView, IOSpec("array", "0.2.0"))
 @_REGISTRY.register_write(H5Group, DaskArray, IOSpec("array", "0.2.0"))
+@suppress_autoshard_warning
 def write_basic_dask_dask_dense(
     f: ZarrGroup | H5Group,
     k: str,
@@ -521,11 +587,13 @@ def write_basic_dask_dask_dense(
     is_h5 = isinstance(f, H5Group)
     if not is_h5:
         dataset_kwargs = zarr_v3_compressor_compat(dataset_kwargs)
-        dataset_kwargs = zarr_v3_sharding(dataset_kwargs)
     if is_h5:
         g = f.require_dataset(k, shape=elem.shape, dtype=elem.dtype, **dataset_kwargs)
     else:
-        g = f.require_array(k, shape=elem.shape, dtype=elem.dtype, **dataset_kwargs)
+        with zarr_v3_sharding(
+            dataset_kwargs, format=f.metadata.zarr_format
+        ) as dataset_kwargs:
+            g = f.require_array(k, shape=elem.shape, dtype=elem.dtype, **dataset_kwargs)
     da.store(elem, g, scheduler="threads")
 
 
@@ -590,6 +658,7 @@ def write_vlen_string_array(
 @_REGISTRY.register_write(ZarrGroup, (np.ndarray, "U"), IOSpec("string-array", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, (np.ndarray, "O"), IOSpec("string-array", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, (np.ndarray, "T"), IOSpec("string-array", "0.2.0"))
+@suppress_autoshard_warning
 @zero_dim_array_as_scalar
 def write_vlen_string_array_zarr(
     f: ZarrGroup,
@@ -605,15 +674,17 @@ def write_vlen_string_array_zarr(
     filters, fill_value = None, None
     if f.metadata.zarr_format == 2:
         filters, fill_value = [VLenUTF8()], ""
-    dataset_kwargs = zarr_v3_sharding(dataset_kwargs)
-    f.create_array(
-        k,
-        shape=elem.shape,
-        dtype=dtype,
-        filters=filters,
-        fill_value=fill_value,
-        **dataset_kwargs,
-    )
+    with zarr_v3_sharding(
+        dataset_kwargs, format=f.metadata.zarr_format
+    ) as dataset_kwargs:
+        f.create_array(
+            k,
+            shape=elem.shape,
+            dtype=dtype,
+            filters=filters,
+            fill_value=fill_value,
+            **dataset_kwargs,
+        )
     f[k][:] = elem
 
 
@@ -674,8 +745,9 @@ def write_recarray_zarr(
     dataset_kwargs = dict(dataset_kwargs)
     dataset_kwargs = zarr_v3_compressor_compat(dataset_kwargs)
     # https://github.com/zarr-developers/zarr-python/issues/3546
-    # if "shards" not in dataset_kwargs and ad.settings.auto_shard_zarr_v3:
-    #     dataset_kwargs = {**dataset_kwargs, "shards": "auto"}
+    # with zarr_v3_sharding(
+    #     dataset_kwargs, format=f.metadata.zarr_format
+    # ) as dataset_kwargs:
     f.create_array(k, shape=elem.shape, dtype=elem.dtype, **dataset_kwargs)
     f[k][...] = elem
 
@@ -730,16 +802,20 @@ def write_sparse_compressed(
                 attr_name, data=attr, shape=attr.shape, dtype=dtype, **dataset_kwargs
             )
         else:
-            dataset_kwargs = zarr_v3_sharding(dataset_kwargs)
-            arr = g.create_array(
-                attr_name, shape=attr.shape, dtype=dtype, **dataset_kwargs
-            )
+            with zarr_v3_sharding(
+                dataset_kwargs, format=f.metadata.zarr_format
+            ) as dataset_kwargs:
+                arr = g.create_array(
+                    attr_name, shape=attr.shape, dtype=dtype, **dataset_kwargs
+                )
             # see https://github.com/zarr-developers/zarr-python/discussions/2712
             arr[...] = attr[...]
 
 
-write_csr = partial(write_sparse_compressed, fmt="csr")
-write_csc = partial(write_sparse_compressed, fmt="csc")
+write_csr, write_csc = (
+    suppress_autoshard_warning(partial(write_sparse_compressed, fmt=fmt))
+    for fmt in ["csr", "csc"]
+)
 
 for store_type, (cls, spec, func) in product(
     (H5Group, ZarrGroup),
@@ -776,6 +852,7 @@ def write_sparse_compressed(
 @_REGISTRY.register_write(H5Group, _CSCDataset, IOSpec("csc_matrix", "0.1.0"))
 @_REGISTRY.register_write(ZarrGroup, _CSRDataset, IOSpec("csr_matrix", "0.1.0"))
 @_REGISTRY.register_write(ZarrGroup, _CSCDataset, IOSpec("csc_matrix", "0.1.0"))
+@suppress_autoshard_warning
 def write_sparse_dataset(
     f: _GroupStorageType,
     k: str,
@@ -902,6 +979,7 @@ def read_sparse_partial(elem, *, items=None, indices=(slice(None), slice(None)))
 @_REGISTRY.register_write(
     ZarrGroup, views.AwkwardArrayView, IOSpec("awkward-array", "0.1.0")
 )
+@suppress_autoshard_warning
 def write_awkward(
     f: _GroupStorageType,
     k: str,
@@ -945,6 +1023,7 @@ def read_awkward(elem: _GroupStorageType, *, _reader: Reader) -> AwkArray:
 @_REGISTRY.register_write(H5Group, pd.DataFrame, IOSpec("dataframe", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, views.DataFrameView, IOSpec("dataframe", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, pd.DataFrame, IOSpec("dataframe", "0.2.0"))
+@suppress_autoshard_warning
 def write_dataframe(
     f: _GroupStorageType,
     key: str,
@@ -1086,6 +1165,7 @@ def read_partial_dataframe_0_1_0(
 
 @_REGISTRY.register_write(H5Group, pd.Categorical, IOSpec("categorical", "0.2.0"))
 @_REGISTRY.register_write(ZarrGroup, pd.Categorical, IOSpec("categorical", "0.2.0"))
+@suppress_autoshard_warning
 def write_categorical(
     f: _GroupStorageType,
     k: str,

src/anndata/_io/specs/registry.py

@@ -114,7 +114,7 @@ def __init__(self) -> None:
 
     def register_write[S: StorageType, T: RWAble](
         self,
-        dest_type: type,
+        dest_type: type[S],
         src_type: type | tuple[type, str],
         spec: IOSpec | Mapping[str, str],
         modifiers: Iterable[str] = frozenset(),

src/anndata/_settings.py

@@ -456,16 +456,6 @@ def validate_zarr_write_format(format: int, settings: SettingsManager):
     if format not in {2, 3}:
         msg = "non-v2 zarr on-disk format not supported"
         raise ValueError(msg)
-    if format == 2 and getattr(settings, "auto_shard_zarr_v3", False):
-        msg = "Cannot set `zarr_write_format` to 2 with autosharding on.  Please set to `False` `anndata.settings.auto_shard_zarr_v3`"
-        raise ValueError(msg)
-
-
-def validate_zarr_sharding(auto_shard: bool, settings: SettingsManager):  # noqa: FBT001
-    validate_bool(auto_shard, settings)
-    if auto_shard and settings.zarr_write_format == 2:
-        msg = "Cannot shard v2 format data. Please set `anndata.settings.zarr_write_format` to 3."
-        raise ValueError(msg)
 
 
 settings.register(
@@ -517,10 +507,11 @@ def validate_sparse_settings(val: Any, settings: SettingsManager) -> None:
 
 settings.register(
     "auto_shard_zarr_v3",
-    default_value=False,
+    default_value=None,
     description="Whether or not to use zarr's auto computation of sharding for v3.  For v2 this setting will be ignored. The setting will apply to all calls to anndata's writing mechanism (write_zarr / write_elem) and will **not** override any user-defined kwargs for shards.",
-    validate=validate_zarr_sharding,
-    get_from_env=check_and_get_bool,
+    validate=gen_validator((bool, NoneType)),
+    option_type=bool | None,
+    get_from_env=check_and_get_bool_or_none,
 )
 
 

src/anndata/_settings.pyi

@@ -46,6 +46,6 @@ class _AnnDataSettingsManager(SettingsManager):
     min_rows_for_chunked_h5_copy: int = 1000
     disallow_forward_slash_in_h5ad: bool = False
     write_csr_csc_indices_with_min_possible_dtype: bool = False
-    auto_shard_zarr_v3: bool = False
+    auto_shard_zarr_v3: bool | None = None
 
 settings: _AnnDataSettingsManager

tests/conftest.py

@@ -33,6 +33,12 @@ def backing_h5ad(tmp_path: Path) -> Path:
     return tmp_path / "test.h5ad"
 
 
+@pytest.fixture(autouse=True)
+def zarr_shard(request: pytest.FixtureRequest):
+    with ad.settings.override(auto_shard_zarr_v3=True):
+        yield
+
+
 @pytest.fixture(
     params=[("h5ad", None), ("zarr", 2), ("zarr", 3)],
     ids=["h5ad", "zarr2", "zarr3"],