Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 8 additions & 0 deletions docs/release-notes/2439.feat.md
Original file line number Diff line number Diff line change
@@ -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`
7 changes: 1 addition & 6 deletions src/anndata/_io/specs/methods.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
5 changes: 4 additions & 1 deletion src/anndata/_io/specs/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
8 changes: 4 additions & 4 deletions src/anndata/_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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,
)


Expand Down
4 changes: 2 additions & 2 deletions src/anndata/_settings.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
6 changes: 0 additions & 6 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
Expand Down
3 changes: 2 additions & 1 deletion tests/lazy/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 10 additions & 4 deletions tests/lazy/test_read.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -23,7 +25,6 @@

if TYPE_CHECKING:
from collections.abc import Callable
from pathlib import Path

from anndata._types import AnnDataElem

Expand Down Expand Up @@ -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
Expand Down
11 changes: 4 additions & 7 deletions tests/test_backed_sparse.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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(
Expand Down
7 changes: 3 additions & 4 deletions tests/test_concatenate_disk.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
36 changes: 16 additions & 20 deletions tests/test_io_elementwise.py
Original file line number Diff line number Diff line change
Expand Up @@ -937,15 +937,12 @@ def test_write_auto_sharded_size(tmp_path: Path):


@pytest.mark.zarr_io
def test_write_auto_sharded_default_warns(tmp_path: Path):
def test_write_shards_by_default(tmp_path: Path):
path = tmp_path / "check.zarr"
adata = gen_adata((100, 10), **GEN_ADATA_NO_XARRAY_ARGS)
ad.settings.reset("auto_shard_zarr_v3")
with (
ad.settings.override(zarr_write_format=3),
pytest.warns(UserWarning, match=r"zarr v3 autosharding will be the default"),
):
adata.write_zarr(path)
adata.write_zarr(path)
check_all_sharded(zarr.open(path))


@pytest.mark.zarr_io
Expand All @@ -954,20 +951,19 @@ def test_write_auto_sharded_does_not_override(tmp_path: Path):
X = sparse.random(
100, 100, density=0.1, format="csr", random_state=np.random.default_rng(42)
)
with ad.settings.override(auto_shard_zarr_v3=True, zarr_write_format=3):
ad.io.write_elem(z, "X_default", X)
shards_default = z["X_default"]["indices"].shards
new_shards = shards_default[0] // 2
new_shards = int(new_shards - new_shards % 2)
ad.io.write_elem(
z,
"X_manually_set",
X,
dataset_kwargs={
"shards": (new_shards,),
"chunks": (int(new_shards / 2),),
},
)
ad.io.write_elem(z, "X_default", X)
shards_default = z["X_default"]["indices"].shards
new_shards = shards_default[0] // 2
new_shards = int(new_shards - new_shards % 2)
ad.io.write_elem(
z,
"X_manually_set",
X,
dataset_kwargs={
"shards": (new_shards,),
"chunks": (int(new_shards / 2),),
},
)

def visitor(key: str, array: zarr.Array):
assert array.shards == (new_shards,)
Expand Down
Loading