feat: migrate to zarr v3 defaults (#2439)

Co-authored-by: Philipp A. <flying-sheep@web.de>

Author: ilan-gold

docs/release-notes/2439.feat.md

@@ -0,0 +1,8 @@
+{attr}`anndata.settings.zarr_write_format` is now 3 and {func}`anndata.settings.auto_shard_zarr_v3` is now `True` (with `None` removed as an option).
+This means that sharded zarr v3 stores will be written with a target uncompressed shard size of 1GB unless the user either
+
+- changes one of these flags, or
+- explicitly creates a {class}`zarr.Group` with `zarr_format=2`, or
+- overrides the `shards` in `dataset_kwargs` passed in to {func}`~anndata.io.write_elem`
+
+{user}`ilan-gold`

src/anndata/_io/specs/methods.py

@@ -125,12 +125,7 @@ def zarr_v3_sharding(dataset_kwargs: dict, format: Literal[2, 3]) -> Generator[d
         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:
+    if auto_sharding:
         dataset_kwargs = {**dataset_kwargs, "shards": "auto"}
     # Auto shard sizes are a relatively recent feature
     supports_auto_shard_size = Version(version("zarr")) >= Version("3.1.4")

src/anndata/_io/specs/registry.py

@@ -483,8 +483,11 @@ def read_elem_lazy(
 
     Reading a dense matrix from a zarr store lazily:
 
+    ..
+        TODO: remove “SKIP” once https://github.com/zarr-developers/zarr-python/issues/3602 becomes minimum zarr (3.1.6)
+
     >>> adata.layers["dense"] = ad.experimental.read_elem_lazy(g["layers/dense"])
-    >>> adata.layers["dense"]
+    >>> adata.layers["dense"]  # doctest: +SKIP
     dask.array<from-zarr, shape=(2700, 32738), dtype=float32, chunksize=(169, 2047), chunktype=numpy.ndarray>
 
     Making a new anndata object from on-disk, with custom chunks:

src/anndata/_settings.py

@@ -460,7 +460,7 @@ def validate_zarr_write_format(format: int, settings: SettingsManager):
 
 settings.register(
     "zarr_write_format",
-    default_value=2,
+    default_value=3,
     description="Which version of zarr to write to when anndata must internally open a write-able zarr group.",
     validate=validate_zarr_write_format,
     get_from_env=lambda name, default: check_and_get_environ_var(
@@ -507,11 +507,11 @@ def validate_sparse_settings(val: Any, settings: SettingsManager) -> None:
 
 settings.register(
     "auto_shard_zarr_v3",
-    default_value=None,
+    default_value=True,
     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=gen_validator((bool, NoneType)),
+    validate=validate_bool,
     option_type=bool | None,
-    get_from_env=check_and_get_bool_or_none,
+    get_from_env=check_and_get_bool,
 )
 
 

src/anndata/_settings.pyi

@@ -41,11 +41,11 @@ class _AnnDataSettingsManager(SettingsManager):
     check_uniqueness: bool = True
     copy_on_write_X: bool = False
     allow_write_nullable_strings: bool | None = None
-    zarr_write_format: Literal[2, 3] = 2
+    zarr_write_format: Literal[2, 3] = 3
     use_sparse_array_on_read: bool = False
     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 | None = None
+    auto_shard_zarr_v3: bool = True
 
 settings: _AnnDataSettingsManager

tests/conftest.py

@@ -33,12 +33,6 @@ 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"],

tests/lazy/conftest.py

@@ -141,7 +141,8 @@ def adata_remote_with_store_tall_skinny_path(
         g,
         "obs",
         obs,
-        dataset_kwargs=dict(chunks=(250,)),
+        # No shards so we can track chunking exactly.
+        dataset_kwargs=dict(chunks=(250,), shards=None),
     )
     zarr.consolidate_metadata(g.store)
     return orig_path

tests/lazy/test_read.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import json
 from importlib.util import find_spec
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -23,7 +25,6 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable
-    from pathlib import Path
 
     from anndata._types import AnnDataElem
 
@@ -188,14 +189,19 @@ def test_unconsolidated(tmp_path: Path, mtx_format):
     adata = gen_adata((10, 10), mtx_format, **GEN_ADATA_NO_XARRAY_ARGS)
     orig_pth = tmp_path / "orig.zarr"
     adata.write_zarr(orig_pth)
-    (orig_pth / ".zmetadata").unlink()
+    with Path.open(orig_pth / "zarr.json", "r+") as f:
+        data = json.load(f)
+        del data["consolidated_metadata"]
+        f.seek(0)
+        json.dump(data, f)
+        f.truncate()
     store = AccessTrackingStore(orig_pth, read_only=True)
-    store.initialize_key_trackers(["obs/.zgroup", ".zgroup"])
+    store.initialize_key_trackers(["obs/zarr.json", "zarr.json"])
     with pytest.warns(UserWarning, match=r"Did not read zarr as consolidated"):
         remote = read_lazy(store)
     remote_to_memory = remote.to_memory()
     assert_equal(remote_to_memory, adata)
-    store.assert_access_count("obs/.zgroup", 1)
+    store.assert_access_count("obs/zarr.json", 1)
 
 
 @pytest.mark.zarr_io

tests/test_backed_sparse.py

@@ -397,10 +397,7 @@ def test_lazy_array_cache(
     a_disk[3:5]
     a_disk[6:7]
     a_disk[8:9]
-    # Three hits for metadata in zarr v3:
-    # see https://github.com/zarr-developers/zarr-python/discussions/2760 for more info on the difference.
-    # Then there is actual data access, 1 more when cached, 4 more otherwise.
-    c_expected = 4 if should_cache_indptr else 7
+    c_expected = 2 if should_cache_indptr else 5
     assert store.get_access_count("X/indptr") == c_expected
     for elem_not_indptr in elems - {"indptr"}:
         assert (
@@ -540,9 +537,9 @@ def test_data_access(
         subset.X.compute(scheduler="single-threaded")
     # zarr v2 fetches all and not just metadata for that node in 3.X.X python package
     # TODO: https://github.com/zarr-developers/zarr-python/discussions/2760
-    if (
-        ad.settings.zarr_write_format == 2 and read_data
-    ) or open_func is not sparse_dataset:
+    if ad.settings.zarr_write_format == 2 and (
+        read_data or open_func is not sparse_dataset
+    ):
         exp = [*exp, "X/data/.zgroup", "X/data/.zattrs"]
 
     assert store.get_access_count("X/data") == len(exp), store.get_accessed_keys(

tests/test_concatenate_disk.py

@@ -11,7 +11,7 @@
 from scipy import sparse
 
 import anndata as ad
-from anndata import AnnData, concat, settings
+from anndata import AnnData, concat
 from anndata._core import merge
 from anndata._core.merge import _resolve_axis
 from anndata.experimental.merge import as_group, concat_on_disk
@@ -261,11 +261,10 @@ 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)
 
 
-def test_concatenate_zarr_v3_shard(xxxm_adatas, tmp_path):
+def test_concatenate_zarr_stays_sharded_v3(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")
+    assert_eq_concat_on_disk(xxxm_adatas, tmp_path, file_format="zarr")
     g = zarr.open(tmp_path)
     assert g.metadata.zarr_format == 3