Update Zarr to v3

.github/workflows/test.yaml

@@ -73,6 +73,10 @@ jobs:
           cache-dependency-glob: pyproject.toml
       - name: create hatch environment
         run: uvx hatch env create ${{ matrix.env.name }}
+      - name: Show installed packages
+        if: matrix.env.python == '3.13'
+        run: |
+          uvx hatch run hatch-test.py3.13:uv pip freeze
       - name: run tests using hatch
         env:
           MPLBACKEND: agg

.gitignore

@@ -21,12 +21,12 @@ __pycache__/
 /docs/generated/
 /docs/_build/
 /docs/generated/
-docs/api/ehrdata
-docs/api/data
-docs/api/tools
-docs/api/io
-docs/api/plot
-!docs/api/api.md
+/docs/api/ehrdata.*.rst
+/docs/api/data
+/docs/api/tools
+/docs/api/io
+/docs/api/plot
+!/docs/api/api.md
 
 ehrapy_data
 test.ipynb

.readthedocs.yaml

@@ -10,5 +10,5 @@ build:
       - asdf global uv latest
     build:
       html:
-        - uvx hatch run docs:build
+        - uvx hatch run docs:build -W # do not remove -W, fix warnings.
         - mv docs/_build $READTHEDOCS_OUTPUT

CHANGELOG.md

@@ -14,14 +14,16 @@ and this project adheres to [Semantic Versioning][].
 
 ### Maintenance
 - Enhanced {doc}`tutorials/getting_started` ([#184](https://github.com/theislab/ehrdata/pull/184)) @eroell
+- Move from zarr<3 to zarr>=3 ([#185](https://github.com/theislab/ehrdata/pull/185)) @eroell
 
 ### Fixed
 
 ### Modified
 - `EHRData` drops the `.R` field in favor of using `.layers` for any 3D data arrays ([#184](https://github.com/theislab/ehrdata/pull/184)) @eroell
 - `EHRData`'s shape property will always return a 3 dimensional shape. If an `EHRData` object has flat arrays only, the third dimension will be 1. ([#184](https://github.com/theislab/ehrdata/pull/184)) @eroell
 - The following functions now take a `layer` argument: {func}`~ehrdata.io.read_csv`, {func}`~ehrdata.io.from_pandas`, {func}`~ehrdata.io.to_pandas`, {func}`~ehrdata.io.omop.setup_variables`, {func}`~ehrdata.io.omop.setup_interval_variables`, {func}`~ehrdata.dt.ehrdata_blobs`, {func}`~ehrdata.dt.physionet2012`. If it is let to its default, `None`, the `.X` field of `EHRData` is used. Since `.X` is 2D in this release, in cases with 3D data, the `layer` argument needs to be used. ([#184](https://github.com/theislab/ehrdata/pull/184)) @eroell
-
+- {func}`~ehrdata.io.write_zarr` now writes an `EHRData` specific store encoding, with `anndata` as a substore. This change allows to use `AnnData` with its change to consolidated Zarr metadata, and better isolates `AnnData`'s io. ([#185](https://github.com/theislab/ehrdata/pull/185)) @eroell
+- {func}`~ehrdata.io.read_zarr` is adapted to read the new store encoding, and can also deal with `AnnData` stores. ([#185](https://github.com/theislab/ehrdata/pull/185)) @eroell
 
 
 ## [0.0.9]

docs/conf.py

@@ -81,16 +81,18 @@
 }
 
 intersphinx_mapping = {
-    "python": ("https://docs.python.org/3", None),
     "anndata": ("https://anndata.readthedocs.io/en/stable", None),
+    "fsspec": ("https://filesystem-spec.readthedocs.io/en/stable", None),
     "h5py": ("https://docs.h5py.org/en/latest", None),
+    "lamin": ("https://docs.lamin.ai", None),
     "numpy": ("https://numpy.org/doc/stable", None),
-    "torch": ("https://pytorch.org/docs/main/", None),
-    "scipy": ("https://docs.scipy.org/doc/scipy", None),
     "pandas": ("https://pandas.pydata.org/docs", None),
-    "zarr": ("https://zarr.readthedocs.io/en/stable", None),
+    "python": ("https://docs.python.org/3", None),
+    "scanpy": ("https://scanpy.readthedocs.io/en/stable", None),
+    "scipy": ("https://docs.scipy.org/doc/scipy", None),
+    "torch": ("https://docs.pytorch.org/docs/main", None),
     "vitessce": ("https://python-docs.vitessce.io", None),
-    "lamin": ("https://docs.lamin.ai", None),
+    "zarr": ("https://zarr.readthedocs.io/en/stable", None),
 }
 
 # List of patterns, relative to source directory, that match files and
@@ -105,7 +107,7 @@
 html_static_path = ["_static"]
 html_css_files = ["css/custom.css"]
 html_title = project_name
-html_logo = "_static/tutorial_images/logo.png"
+html_logo = "_static/tutorial_images/ehrdata_logo.png"
 html_theme_options = {
     "repository_url": repository_url,
     "use_repository_button": True,
@@ -127,6 +129,10 @@
 nitpick_ignore = [
     ("py:class", "pathlib._local.Path"),
     ("py:class", "types.EllipsisType"),
+    # TODO: remove once https://github.com/sphinx-doc/sphinx/pull/13508 is released
+    ("py:class", "ehrdata._types.TypeAliasType"),
+    # typing.Union fails in tutorials/tutorial_time_series_with_pypots
+    ("py:data", "typing.Union"),
     # https://github.com/duckdb/duckdb-web/issues/3806
     ("py:class", "duckdb.duckdb.DuckDBPyConnection"),
     ("py:class", "_duckdb.DuckDBPyConnection"),
@@ -136,18 +142,18 @@
     ("py:class", "awkward.highlevel.Array"),
     ("py:class", "h5py._hl.dataset.Dataset"),
     ("py:class", "zarr.core.Array"),
+    ("py:class", "zarr.core.buffer.core.Buffer"),
     ("py:class", "ehrdata._compat.ZappyArray"),
     ("py:class", "dask.array.core.Array"),
     ("py:class", "anndata.compat.CupyArray"),
     ("py:class", "anndata.compat.CupySparseMatrix"),
     ("py:class", "sparse.numba_backend._coo.core.COO"),
     ("py:class", "sparse._coo.core.COO"),
-    ("py:data", "typing.Union"),  # typing.Union fails in tutorials/tutorial_time_series_with_pypots
 ]
 
 # Redirect broken parameter annotation classes
 qualname_overrides = {
-    "zarr._storage.store.Store": "zarr.storage.MemoryStore",
+    "zarr.storage._common.StorePath": "zarr.storage.StorePath",
     "zarr.core.group.Group": "zarr.group.Group",
     "lnschema_core.models.Artifact": "lamindb.Artifact",
 }

pyproject.toml

@@ -30,15 +30,18 @@ dependencies = [
   "requests",
   "rich",
   "xarray",
-  "zarr<3",
+  "zarr>=3",
+]
+optional-dependencies.dask = [
+  "dask",
 ]
 optional-dependencies.dev = [
   "pre-commit",
   "twine",
 ]
 optional-dependencies.doc = [
   "docutils>=0.8,!=0.18.*,!=0.19.*",
-  "ehrdata[lamin,torch,vitessce]",
+  "ehrdata[lamin,torch]",            # do not include vitessce[all] in doc: https://github.com/keller-mark/esbuild-py/issues/19
   "ipykernel",
   "ipython",
   "myst-nb>=1.1",
@@ -51,17 +54,18 @@ optional-dependencies.doc = [
   "sphinx-tabs",
   "sphinxcontrib-bibtex>=1",
   "sphinxext-opengraph",
+  "vitessce",
 ]
 optional-dependencies.ehrapy = [
   "ehrapy",
 ]
 optional-dependencies.lamin = [
-  "lamindb",
+  "lamindb>=1.1",
   #"omop", # Not needed and causing an issue in the readthedocs build https://github.com/laminlabs/lamindb/issues/3081
 ]
 optional-dependencies.test = [
   "coverage",
-  "ehrdata[lamin,torch,vitessce]",
+  "ehrdata[dask,lamin,torch,vitessce]",
   "pytest",
   "pytest-cov",
 ]
@@ -71,7 +75,6 @@ optional-dependencies.torch = [
 
 optional-dependencies.vitessce = [
   "vitessce[all]>=3.4", # the actual dependency
-  "zarr>=2.18.7,<3",    # vitessce does not support zarr>=3
 ]
 # https://docs.pypi.org/project_metadata/#project-urls
 urls.Documentation = "https://ehrdata.readthedocs.io/"

src/ehrdata/core/constants.py

@@ -11,6 +11,8 @@
 
 DEFAULT_TEM_LAYER_NAME = "tem_data"
 
+EHRDATA_ZARR_ENCODING_VERSION = "0.0.1"
+
 # Missing values
 # --------------
 # These values if encountered as strings are considered to represent missing values in the data

src/ehrdata/core/ehrdata.py

@@ -171,7 +171,7 @@ def _get_layers_3d_dim(layers: Mapping[str, Any] | None) -> int:
 class EHRData(AnnData):
     """Model two and three dimensional electronic health record data.
 
-    .. figure:: ../../_static/tutorial_images/logo.png
+    .. figure:: /_static/tutorial_images/ehrdata_logo.png
        :width: 260px
        :align: right
        :class: dark-light

src/ehrdata/integrations/vitessce/_config.py

@@ -11,13 +11,13 @@
 
     from lamindb import Artifact
     from vitessce import VitessceConfig
-    from zarr.storage import Store
+    from zarr.storage import StoreLike
 
 
 def gen_config(
     path: Path | None = None,
     *,
-    store: Path | Store | None = None,
+    store: Path | StoreLike | None = None,
     url: str | None = None,
     artifact: Artifact | None = None,
     # arguments not about how the store goes in:

src/ehrdata/io/zarr.py

@@ -1,15 +1,20 @@
 from __future__ import annotations
 
+import warnings
+from functools import wraps
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Literal
 
 import anndata as ad
 import zarr
 
+import ehrdata as ed
 from ehrdata._logger import logger
+from ehrdata.core.constants import EHRDATA_ZARR_ENCODING_VERSION
 from ehrdata.io._array_casting import _cast_arrays_dtype_to_float_or_str_if_nonnumeric_object, _cast_variables_to_float
 
 if TYPE_CHECKING:
+    from collections.abc import Callable, Mapping
     from os import PathLike
 
     from ehrdata import EHRData
@@ -23,6 +28,8 @@ def read_zarr(
 ) -> EHRData:
     """Read a zarr store into an :class:`~ehrdata.EHRData` object.
 
+    Can also read :class:`~anndata.AnnData` Zarr stores. In this case, a default `.tem` field is created in the `ehrdata object`.
+
     Args:
         filename: The filename, or a Zarr storage class.
         harmonize_missing_values: Whether to call `ehrdata.harmonize_missing_values` on all detected layers.
@@ -36,15 +43,36 @@ def read_zarr(
         >>> ed.io.write_zarr("mimic_2.zarr", edata)
         >>> edata_from_zarr = ed.io.read_zarr("mimic_2.zarr")
     """
-    import ehrdata as ed
     from ehrdata import EHRData
 
     if isinstance(filename, Path):
         filename = str(filename)
 
     f = filename if isinstance(filename, zarr.Group) else zarr.open(filename, mode="r")
 
-    dictionary_for_init = {k: ad.io.read_elem(f[k]) for k, v in dict(f).items() if not k.startswith("raw.")}
+    if "encoding-type" not in f.attrs:
+        err = "The zarr store does not contain an encoding-type attribute."
+        raise ValueError(err)
+
+    if f.attrs["encoding-type"] == "ehrdata":
+        if "anndata" in f:
+            dictionary_for_init = {
+                k: ad.io.read_elem(f["anndata"][k]) for k, v in dict(f["anndata"]).items() if not k.startswith("raw.")
+            }
+        else:
+            err = "The zarr store does not contain the 'anndata' group."
+            raise ValueError(err)
+        if "tem" in f:
+            dictionary_for_init["tem"] = ad.io.read_elem(f["tem"])
+        else:
+            warnings.warn("The zarr store does not contain the 'tem' group.", stacklevel=2)
+
+    elif f.attrs["encoding-type"] == "anndata":
+        dictionary_for_init = {k: ad.io.read_elem(f[k]) for k, v in dict(f).items() if not k.startswith("raw.")}
+
+    else:
+        err = f"Unkown encoding-type '{f.attrs['encoding-type']}'."
+        raise ValueError(err)
 
     edata = EHRData(**dictionary_for_init)
 
@@ -62,37 +90,86 @@ def read_zarr(
     return edata
 
 
+def _allow_write_nullable_strings[T, **P](f: Callable[P, T]) -> Callable[P, T]:
+    @wraps(f)
+    def wrapped(*args: P.args, **kwargs: P.kwargs):
+        with ad.settings.override(allow_write_nullable_strings=True):
+            return f(*args, **kwargs)
+
+    return wrapped
+
+
+@_allow_write_nullable_strings
 def write_zarr(
     edata: EHRData,
     filename: str | Path,
     *,
-    chunks: bool | int | tuple[int, ...] | None = None,
+    chunks: Literal["auto" | "ehrdata_auto"] = "auto",
     convert_strings_to_categoricals: bool = True,
 ) -> None:
     """Write :class:`~ehrdata.EHRData` objects to disk.
 
-    To write to a `.zarr` file, `X` and `layers` cannot be written as  `object` dtype.
+    To write to a `.zarr` file, `X`, and `layers` cannot be written as `object` dtype.
     If any of these fields is of `object` dtype, it this function will attempt to cast it to a numeric dtype; if this fails, the field will be casted to a `str` dtype.
 
     Args:
         edata: Central data object.
         filename: Name of the output file, can also be prefixed with relative or absolute path to save the file to.
-        chunks: Chunk shape, passed to :meth:`zarr.Group.create_dataset` for `Zarr` version 2, or to :meth:`zarr.Group.create_array` for `Zarr` version 3.
-        convert_strings_to_categoricals: Convert columns of `str` dtype in `.obs` and `.var` to `categorical` dtype.
+        chunks: Specify strategy of how data should be chunked. For simplicity, currently only 2 options are available: `"auto"` will write the data with :func:`~anndata.io.write_elem`'s default settings. `"ehrdata_auto"` will write the data chunked (and sharded) based on a heuristic that loosely speaking writes slightly smaller chunks.
+        convert_strings_to_categoricals: Convert columns of `str` dtype in `.obs` and `.var` and `.tem` to `categorical` dtype.
 
     Examples:
         >>> import ehrdata as ed
         >>> edata = ed.dt.mimic_2()
         >>> ed.io.write_zarr("mimic_2.zarr", edata)
     """
     filename = Path(filename)
-
     edata = _cast_arrays_dtype_to_float_or_str_if_nonnumeric_object(edata)
 
-    ad.AnnData(edata).write_zarr(
-        filename,
-        chunks=chunks,
-        convert_strings_to_categoricals=convert_strings_to_categoricals,
-    )
-    f = zarr.open(filename, mode="a")
-    ad.io.write_elem(f, "tem", edata.tem)
+    store = zarr.open_group(filename, mode="a", use_consolidated=False, zarr_format=3)
+
+    adata = ad.AnnData(edata)
+
+    if convert_strings_to_categoricals:
+        adata.strings_to_categoricals(adata.obs)
+        adata.strings_to_categoricals(adata.var)
+        adata.strings_to_categoricals(edata.tem)
+
+    # write_sharded this is a slightly modified version from https://anndata.readthedocs.io/en/stable/tutorials/zarr-v3.html
+    # write_sharded is intended as a future blueprint of implementing better chunking defaults for ehrdata based based on real usecases
+    def write_sharded(group: zarr.Group, adata: ad.AnnData):
+        def callback(
+            func: ad.experimental.Write,
+            g: zarr.Group,
+            k: str,
+            elem: ad.typing.RWAble,
+            dataset_kwargs: Mapping[str, Any],
+            iospec: ad.experimental.IOSpec,
+        ):
+            if iospec.encoding_type in {"array"} and not isinstance(elem, list):
+                dataset_kwargs = {
+                    "shards": tuple(int(2 ** (16 / len(elem.shape))) for _ in elem.shape),
+                    **dataset_kwargs,
+                }
+                dataset_kwargs["chunks"] = tuple(i // 2 for i in dataset_kwargs["shards"])
+            elif iospec.encoding_type in {"csr_matrix", "csc_matrix"}:
+                dataset_kwargs = {"shards": (2**16,), "chunks": (2**8,), **dataset_kwargs}
+            func(g, k, elem, dataset_kwargs=dataset_kwargs)
+
+        return ad.experimental.write_dispatched(group, "/", adata, callback=callback)
+
+    if chunks == "auto":
+        ad.io.write_elem(store, "anndata", adata)
+    elif chunks == "ehrdata_auto":
+        anndata_group = store.create_group("anndata")
+        write_sharded(anndata_group, adata)
+    else:
+        err = (
+            f"chunks={chunks} is not implemented. Currently, only chunks='auto' and chunks='ehrdata_auto' is supported."
+        )
+        raise NotImplementedError(err)
+
+    ad.io.write_elem(store, "tem", edata.tem)
+
+    store.attrs["encoding-version"] = EHRDATA_ZARR_ENCODING_VERSION
+    store.attrs["encoding-type"] = "ehrdata"