Backport PR #2167 on branch 0.12.x (perf: auto sharding for zarr v3) (#2188)

Author: ilan-gold

docs/release-notes/2167.perf.md

@@ -0,0 +1 @@
+Enable automatic sharding in zarr v3 via {attr}`anndata.settings.auto_shard_zarr_v3` (via {mod}`zarr`'s own auto sharding mechanism i.e., `shards="auto"`) for all types except {class}`numpy.recarray` {user}`ilan-gold`

docs/tutorials/zarr-v3.md

@@ -38,7 +38,8 @@ There are two ways of opening remote `zarr` stores from the `zarr-python` packag
 Local data generally poses a different set of challenges.
 First, write speeds can be somewhat slow and second, the creation of many small files on a file system can slow down a filesystem.
 For the "many small files" problem, `zarr` has introduced {ref}`sharding <zarr:user-guide-sharding>` in the v3 file format.
-Sharding requires knowledge of the array element you are writing (such as shape or data type), though, and therefore you will need to use {func}`anndata.experimental.write_dispatched` to use sharding.
+We offer {attr}`anndata.settings.auto_shard_zarr_v3` to hook into zarr's ability to automatically compute shards, which is experimental at the moment.
+Manual sharding requires knowledge of the array element you are writing (such as shape or data type), though, and therefore you will need to use {func}`anndata.experimental.write_dispatched` to use custom sharding.
 For example, you cannot shard a 1D array with `shard` sizes `(256, 256)`.
 Here is a short example, although you should tune the sizes to your own use-case and also use the compression that makes the most sense for you:
 

pyproject.toml

@@ -164,6 +164,7 @@ filterwarnings_when_strict = [
     "default:Consolidated metadata is:UserWarning",
     "default:.*Structured:zarr.core.dtype.common.UnstableSpecificationWarning",
     "default:.*FixedLengthUTF32:zarr.core.dtype.common.UnstableSpecificationWarning",
+    "default:Automatic shard shape inference is experimental",
 ]
 python_files = "test_*.py"
 testpaths = [

src/anndata/_io/specs/methods.py

@@ -102,6 +102,12 @@ def zarr_v3_compressor_compat(dataset_kwargs) -> dict:
     return dataset_kwargs
 
 
+def zarr_v3_sharding(dataset_kwargs) -> dict:
+    if "shards" not in dataset_kwargs and ad.settings.auto_shard_zarr_v3:
+        dataset_kwargs = {**dataset_kwargs, "shards": "auto"}
+    return dataset_kwargs
+
+
 def _to_cpu_mem_wrapper(write_func):
     """
     Wrapper to bring cupy types into cpu memory before writing.
@@ -432,6 +438,7 @@ 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)
         # see https://github.com/zarr-developers/zarr-python/discussions/2712
         if isinstance(elem, ZarrArray | H5Array):
@@ -511,6 +518,7 @@ 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_zarr_v2() or is_h5:
         g = f.require_dataset(k, shape=elem.shape, dtype=elem.dtype, **dataset_kwargs)
     else:
@@ -616,6 +624,7 @@ 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,
@@ -684,6 +693,9 @@ def write_recarray_zarr(
     else:
         dataset_kwargs = dataset_kwargs.copy()
         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"}
         f.create_array(k, shape=elem.shape, dtype=elem.dtype, **dataset_kwargs)
         f[k][...] = elem
 
@@ -720,6 +732,7 @@ 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
             )

src/anndata/_settings.py

@@ -17,7 +17,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
-    from typing import Any, TypeGuard
+    from typing import Any, Self, TypeGuard
 
 T = TypeVar("T")
 
@@ -55,7 +55,7 @@ class RegisteredOption(NamedTuple, Generic[T]):
     option: str
     default_value: T
     description: str
-    validate: Callable[[T], None]
+    validate: Callable[[T, SettingsManager], None]
     type: object
 
     describe = describe
@@ -206,7 +206,7 @@ def register(
         *,
         default_value: T,
         description: str,
-        validate: Callable[[T], None],
+        validate: Callable[[T, Self], None],
         option_type: object | None = None,
         get_from_env: Callable[[str, T], T] = lambda x, y: y,
     ) -> None:
@@ -229,7 +229,7 @@ def register(
             Default behavior is to return `default_value` without checking the environment.
         """
         try:
-            validate(default_value)
+            validate(default_value, self)
         except (ValueError, TypeError) as e:
             e.add_note(f"for option {option!r}")
             raise e
@@ -307,7 +307,7 @@ def __setattr__(self, option: str, val: object) -> None:
             )
             raise AttributeError(msg)
         registered_option = self._registered_options[option]
-        registered_option.validate(val)
+        registered_option.validate(val, self)
         self._config[option] = val
 
     def __getattr__(self, option: str) -> object:
@@ -364,10 +364,13 @@ def override(self, **overrides):
         """
         restore = {a: getattr(self, a) for a in overrides}
         try:
-            for attr, value in overrides.items():
-                setattr(self, attr, value)
+            # Preserve order so that settings that depend on each other can be overridden together i.e., always override zarr version before sharding
+            for k in self._config:
+                if k in overrides:
+                    setattr(self, k, overrides.get(k))
             yield None
         finally:
+            # TODO: does the order need to be preserved when restoring?
             for attr, value in restore.items():
                 setattr(self, attr, value)
 
@@ -395,7 +398,7 @@ def __doc__(self):
 
 
 def gen_validator(_type: type[V]) -> Callable[[V], None]:
-    def validate_type(val: V) -> None:
+    def validate_type(val: V, settings: SettingsManager) -> None:
         if not isinstance(val, _type):
             msg = f"{val} not valid {_type}"
             raise TypeError(msg)
@@ -434,14 +437,28 @@ def validate_type(val: V) -> None:
 )
 
 
-def validate_zarr_write_format(format: int):
-    validate_int(format)
+def validate_zarr_write_format(format: int, settings: SettingsManager):
+    validate_int(format, settings)
     if format not in {2, 3}:
         msg = "non-v2 zarr on-disk format not supported"
         raise ValueError(msg)
     if format == 3 and is_zarr_v2():
         msg = "Cannot write v3 format against v2 package"
         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:
+        if is_zarr_v2():
+            msg = "Cannot use sharding with `zarr-python<3`. Please upgrade package and set `anndata.settings.zarr_write_format` to 3."
+            raise ValueError(msg)
+        if 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(
@@ -458,8 +475,8 @@ def validate_zarr_write_format(format: int):
 )
 
 
-def validate_sparse_settings(val: Any) -> None:
-    validate_bool(val)
+def validate_sparse_settings(val: Any, settings: SettingsManager) -> None:
+    validate_bool(val, settings)
 
 
 settings.register(
@@ -486,6 +503,14 @@ def validate_sparse_settings(val: Any) -> None:
     get_from_env=check_and_get_bool,
 )
 
+settings.register(
+    "auto_shard_zarr_v3",
+    default_value=False,
+    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,
+)
+
 
 ##################################################################################
 ##################################################################################

src/anndata/_settings.pyi

@@ -2,7 +2,7 @@ from collections.abc import Callable as Callable
 from collections.abc import Generator, Iterable
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Literal, TypeVar
+from typing import Literal, Self, TypeVar
 
 _T = TypeVar("_T")
 
@@ -25,7 +25,7 @@ class SettingsManager:
         *,
         default_value: _T,
         description: str,
-        validate: Callable[[_T], None],
+        validate: Callable[[_T, Self], None],
         option_type: object | None = None,
         get_from_env: Callable[[str, _T], _T] = ...,
     ) -> None: ...
@@ -46,5 +46,6 @@ class _AnnDataSettingsManager(SettingsManager):
     use_sparse_array_on_read: bool = False
     min_rows_for_chunked_h5_copy: int = 1000
     disallow_forward_slash_in_h5ad: bool = False
+    auto_shard_zarr_v3: bool = False
 
 settings: _AnnDataSettingsManager

src/anndata/tests/helpers.py

@@ -14,6 +14,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import zarr
 from pandas.api.types import is_numeric_dtype
 from scipy import sparse
 
@@ -34,6 +35,7 @@
     XDataArray,
     XDataset,
     ZarrArray,
+    ZarrGroup,
     is_zarr_v2,
 )
 from anndata.utils import asarray
@@ -1187,3 +1189,23 @@ def get_multiindex_columns_df(shape: tuple[int, int]) -> pd.DataFrame:
             + list(itertools.product(["b"], range(shape[1] // 2)))
         ),
     )
+
+
+def visititems_zarr(
+    z: ZarrGroup, visitor: Callable[[str, ZarrGroup | zarr.Array], None]
+) -> None:
+    for key in z:
+        maybe_group = z[key]
+        if isinstance(maybe_group, ZarrGroup):
+            visititems_zarr(maybe_group, visitor)
+        else:
+            visitor(key, maybe_group)
+
+
+def check_all_sharded(g: ZarrGroup):
+    def visit(key: str, arr: zarr.Array | zarr.Group):
+        # Check for recarray via https://numpy.org/doc/stable/user/basics.rec.html#manipulating-and-displaying-structured-datatypes
+        if isinstance(arr, zarr.Array) and arr.shape != () and arr.dtype.names is None:
+            assert arr.shards is not None
+
+    visititems_zarr(g, visitor=visit)

tests/test_concatenate_disk.py

@@ -8,12 +8,13 @@
 import pytest
 from scipy import sparse
 
-from anndata import AnnData, concat
+from anndata import AnnData, concat, settings
 from anndata._core import merge
 from anndata._core.merge import _resolve_axis
+from anndata.compat import is_zarr_v2
 from anndata.experimental.merge import as_group, concat_on_disk
 from anndata.io import read_elem, write_elem
-from anndata.tests.helpers import assert_equal, gen_adata
+from anndata.tests.helpers import assert_equal, check_all_sharded, gen_adata
 from anndata.utils import asarray
 
 if TYPE_CHECKING:
@@ -230,7 +231,7 @@ def gen_index(n):
             X=sparse.csr_matrix((2, 100)),
             obs=pd.DataFrame(index=gen_index(2)),
             obsm={
-                "sparse": np.arange(8).reshape(2, 4),
+                "sparse": sparse.csr_matrix(np.arange(8).reshape(2, 4)),
                 "dense": np.arange(4, 8).reshape(2, 2),
                 "df": pd.DataFrame(
                     {
@@ -253,6 +254,22 @@ def test_concatenate_xxxm(xxxm_adatas, tmp_path, file_format, join_type):
     assert_eq_concat_on_disk(xxxm_adatas, tmp_path, file_format, join=join_type)
 
 
+@pytest.mark.skipif(is_zarr_v2(), reason="auto sharding is allowed only for zarr v3.")
+def test_concatenate_zarr_v3_shard(xxxm_adatas, tmp_path):
+    import zarr
+
+    with settings.override(auto_shard_zarr_v3=True, zarr_write_format=3):
+        assert_eq_concat_on_disk(xxxm_adatas, tmp_path, file_format="zarr")
+    g = zarr.open(tmp_path)
+    assert g.metadata.zarr_format == 3
+
+    def visit(key: str, arr: zarr.Array | zarr.Group):
+        if isinstance(arr, zarr.Array) and arr.shape != ():
+            assert arr.shards is not None
+
+    check_all_sharded(g)
+
+
 def test_output_dir_exists(tmp_path):
     in_pth = tmp_path / "in.h5ad"
     out_pth = tmp_path / "does_not_exist" / "out.h5ad"

tests/test_dask.py

@@ -23,6 +23,7 @@
     as_dense_dask_array,
     as_sparse_dask_array,
     assert_equal,
+    check_all_sharded,
     gen_adata,
 )
 
@@ -109,12 +110,20 @@ def test_dask_write(adata, tmp_path, diskfmt):
 
 @pytest.mark.xdist_group("dask")
 @pytest.mark.dask_distributed
+@pytest.mark.parametrize(
+    "auto_shard_zarr_v3",
+    [pytest.param(True, id="shard"), pytest.param(False, id="no-shard")],
+)
 def test_dask_distributed_write(
     adata: AnnData,
     tmp_path: Path,
     diskfmt: Literal["h5ad", "zarr"],
     local_cluster_addr: str,
+    *,
+    auto_shard_zarr_v3: bool,
 ) -> None:
+    if auto_shard_zarr_v3 and ad.settings.zarr_write_format == 2:
+        pytest.skip(reason="Cannot shard v2 data")
     import dask.array as da
     import dask.distributed as dd
     import numpy as np
@@ -128,9 +137,12 @@ def test_dask_distributed_write(
         adata.obsm["b"] = da.random.random((M, 10))
         adata.varm["a"] = da.random.random((N, 10))
         orig = adata
-        ad.io.write_elem(g, "", orig)
+        with ad.settings.override(auto_shard_zarr_v3=auto_shard_zarr_v3):
+            ad.io.write_elem(g, "", orig)
         # TODO: See https://github.com/zarr-developers/zarr-python/issues/2716
         g = as_group(pth, mode="r")
+        if auto_shard_zarr_v3:
+            check_all_sharded(g)
         curr = ad.io.read_elem(g)
 
     with pytest.raises(AssertionError):

tests/test_io_dispatched.py

@@ -12,10 +12,14 @@
 from anndata._io.zarr import open_write_group
 from anndata.compat import CSArray, CSMatrix, ZarrGroup, is_zarr_v2
 from anndata.experimental import read_dispatched, write_dispatched
-from anndata.tests.helpers import GEN_ADATA_NO_XARRAY_ARGS, assert_equal, gen_adata
+from anndata.tests.helpers import (
+    GEN_ADATA_NO_XARRAY_ARGS,
+    assert_equal,
+    gen_adata,
+    visititems_zarr,
+)
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
     from pathlib import Path
     from typing import Literal
 
@@ -180,18 +184,7 @@ def check_chunking(k: str, v: ZarrGroup | zarr.Array):
     if is_zarr_v2():
         z.visititems(check_chunking)
     else:
-
-        def visititems(
-            z: ZarrGroup, visitor: Callable[[str, ZarrGroup | zarr.Array], None]
-        ) -> None:
-            for key in z:
-                maybe_group = z[key]
-                if isinstance(maybe_group, ZarrGroup):
-                    visititems(maybe_group, visitor)
-                else:
-                    visitor(key, maybe_group)
-
-        visititems(z, check_chunking)
+        visititems_zarr(z, check_chunking)
 
 
 @pytest.mark.zarr_io