Zarr v3 support

.github/workflows/test.yaml

@@ -18,7 +18,7 @@ jobs:
         strategy:
             fail-fast: false
             matrix:
-                python: ["3.10", "3.12"]
+                python: ["3.11", "3.13"]
                 os: [ubuntu-latest]
                 include:
                     - os: macos-latest

pyproject.toml

@@ -33,7 +33,8 @@ dependencies = [
     "networkx",
     "numba>=0.55.0",
     "numpy",
-    "ome_zarr>=0.8.4",
+    "ome_zarr>=0.12rc1",
+#    "ome_zarr>=0.8.4",
     "pandas",
     "pooch",
     "pyarrow",
@@ -47,7 +48,7 @@ dependencies = [
     "xarray>=2024.10.0",
     "xarray-schema",
     "xarray-spatial>=0.3.5",
-    "zarr<3",
+    "zarr>=3.0.0",
 ]
 
 [project.optional-dependencies]

src/spatialdata/_core/spatialdata.py

@@ -621,10 +621,13 @@ def _get_groups_for_element(
         -------
         either the existing Zarr subgroup or a new one.
         """
+        from spatialdata._io.format import SpatialDataFormat
+
         if not isinstance(zarr_path, Path):
             raise ValueError("zarr_path should be a Path object")
-        store = parse_url(zarr_path, mode="r+").store
-        root = zarr.group(store=store)
+        store = SpatialDataFormat().init_store(str(zarr_path), mode="r+")
+        # store = parse_url(zarr_path, mode="r+", fmt=SpatialDataFormat()).store
+        root = zarr.open_group(store=store, mode="r+")
         if element_type not in ["images", "labels", "points", "polygons", "shapes", "tables"]:
             raise ValueError(f"Unknown element type {element_type}")
         element_type_group = root.require_group(element_type)
@@ -646,8 +649,10 @@ def _group_for_element_exists(self, zarr_path: Path, element_type: str, element_
         -------
         True if the group exists, False otherwise.
         """
-        store = parse_url(zarr_path, mode="r").store
-        root = zarr.group(store=store)
+        from spatialdata._io.format import SpatialDataFormat
+
+        store = parse_url(zarr_path, mode="r", fmt=SpatialDataFormat()).store
+        root = zarr.open_group(store=store, mode="r")
         assert element_type in ["images", "labels", "points", "polygons", "shapes", "tables"]
         exists = element_type in root and element_name in root[element_type]
         store.close()
@@ -1068,19 +1073,27 @@ def elements_paths_on_disk(self) -> list[str]:
         -------
         A list of paths of the elements saved in the Zarr store.
         """
+        from spatialdata._io.format import SpatialDataFormat
+
         if self.path is None:
             raise ValueError("The SpatialData object is not backed by a Zarr store.")
-        store = parse_url(self.path, mode="r").store
-        root = zarr.group(store=store)
+
+        store = parse_url(self.path, mode="r", fmt=SpatialDataFormat()).store
+        root = zarr.open_group(store=store, mode="r")
         elements_in_zarr = []
 
         def find_groups(obj: zarr.Group, path: str) -> None:
-            # with the current implementation, a path of a zarr group if the path for an element if and only if its
+            # with the current implementation, a path of a zarr group is the path for an element if and only if its
             # string representation contains exactly one "/"
             if isinstance(obj, zarr.Group) and path.count("/") == 1:
                 elements_in_zarr.append(path)
 
-        root.visit(lambda path: find_groups(root[path], path))
+        for element_type in root:
+            if element_type in ["images", "labels", "points", "shapes", "tables"]:
+                for element_name in root[element_type]:
+                    path = f"{element_type}/{element_name}"
+                    elements_in_zarr.append(path)
+        # root.visit(lambda path: find_groups(root[path], path))
         store.close()
         return elements_in_zarr
 
@@ -1115,6 +1128,7 @@ def _validate_can_safely_write_to_path(
         saving_an_element: bool = False,
     ) -> None:
         from spatialdata._io._utils import _backed_elements_contained_in_path, _is_subfolder
+        from spatialdata._io.format import SpatialDataFormat
 
         if isinstance(file_path, str):
             file_path = Path(file_path)
@@ -1123,7 +1137,7 @@ def _validate_can_safely_write_to_path(
             raise ValueError(f"file_path must be a string or a Path object, type(file_path) = {type(file_path)}.")
 
         if os.path.exists(file_path):
-            if parse_url(file_path, mode="r") is None:
+            if parse_url(file_path, mode="r", fmt=SpatialDataFormat()) is None:
                 raise ValueError(
                     "The target file path specified already exists, and it has been detected to not be a Zarr store. "
                     "Overwriting non-Zarr stores is not supported to prevent accidental data loss."
@@ -1205,13 +1219,15 @@ def write(
             :class:`~spatialdata._io.format.CurrentRasterFormat`, :class:`~spatialdata._io.format.CurrentShapesFormat`,
             :class:`~spatialdata._io.format.CurrentPointsFormat`, :class:`~spatialdata._io.format.CurrentTablesFormat`.
         """
+        from spatialdata._io.format import SpatialDataFormat
+
         if isinstance(file_path, str):
             file_path = Path(file_path)
         self._validate_can_safely_write_to_path(file_path, overwrite=overwrite)
         self._validate_all_elements()
 
-        store = parse_url(file_path, mode="w").store
-        zarr_group = zarr.group(store=store, overwrite=overwrite)
+        store = parse_url(file_path, mode="w", fmt=SpatialDataFormat()).store
+        zarr_group = zarr.open_group(store=store, mode="w" if overwrite else "a")
         self.write_attrs(zarr_group=zarr_group)
         store.close()
 
@@ -1370,14 +1386,15 @@ def delete_element_from_disk(self, element_name: str | list[str]) -> None:
         environment (e.g. operating system, local vs network storage, file permissions, ...) and call this function
         appropriately (or implement a tailored solution), to prevent data loss.
         """
+        from spatialdata._io._utils import _backed_elements_contained_in_path
+        from spatialdata._io.format import SpatialDataFormat
+
         if isinstance(element_name, list):
             for name in element_name:
                 assert isinstance(name, str)
                 self.delete_element_from_disk(name)
             return
 
-        from spatialdata._io._utils import _backed_elements_contained_in_path
-
         if self.path is None:
             raise ValueError("The SpatialData object is not backed by a Zarr store.")
 
@@ -1417,8 +1434,8 @@ def delete_element_from_disk(self, element_name: str | list[str]) -> None:
             )
 
         # delete the element
-        store = parse_url(self.path, mode="r+").store
-        root = zarr.group(store=store)
+        store = parse_url(self.path, mode="r+", fmt=SpatialDataFormat()).store
+        root = zarr.open_group(store=store, mode="r+")
         root[element_type].pop(element_name)
         store.close()
 
@@ -1438,15 +1455,19 @@ def _check_element_not_on_disk_with_different_type(self, element_type: str, elem
                 )
 
     def write_consolidated_metadata(self) -> None:
-        store = parse_url(self.path, mode="r+").store
+        from spatialdata._io.format import SpatialDataFormat
+
+        store = parse_url(self.path, mode="r+", fmt=SpatialDataFormat()).store
         # consolidate metadata to more easily support remote reading bug in zarr. In reality, 'zmetadata' is written
         # instead of '.zmetadata' see discussion https://github.com/zarr-developers/zarr-python/issues/1121
-        zarr.consolidate_metadata(store, metadata_key=".zmetadata")
+        zarr.consolidate_metadata(store)
         store.close()
 
     def has_consolidated_metadata(self) -> bool:
+        from spatialdata._io.format import SpatialDataFormat
+
         return_value = False
-        store = parse_url(self.path, mode="r").store
+        store = parse_url(self.path, mode="r", fmt=SpatialDataFormat()).store
         if "zmetadata" in store:
             return_value = True
         store.close()
@@ -1614,16 +1635,16 @@ def _element_type_and_name_from_element_path(self, element_path: str) -> tuple[s
         return element_type, element_name
 
     def write_attrs(self, format: SpatialDataFormat | None = None, zarr_group: zarr.Group | None = None) -> None:
-        from spatialdata._io.format import _parse_formats
+        from spatialdata._io.format import SpatialDataFormat, _parse_formats
 
         parsed = _parse_formats(formats=format)
 
         store = None
 
         if zarr_group is None:
             assert self.is_backed(), "The SpatialData object must be backed by a Zarr store to write attrs."
-            store = parse_url(self.path, mode="r+").store
-            zarr_group = zarr.group(store=store, overwrite=False)
+            store = parse_url(self.path, mode="r+", fmt=SpatialDataFormat()).store
+            zarr_group = zarr.open_group(store=store, overwrite=False, mode="r+")
 
         version = parsed["SpatialData"].spatialdata_format_version
         version_specific_attrs = parsed["SpatialData"].attrs_to_dict()

src/spatialdata/_io/format.py

@@ -6,7 +6,7 @@
 import ome_zarr.format
 import zarr
 from anndata import AnnData
-from ome_zarr.format import CurrentFormat, Format, FormatV01, FormatV02, FormatV03, FormatV04
+from ome_zarr.format import Format, FormatV01, FormatV02, FormatV03, FormatV04
 from pandas.api.types import CategoricalDtype
 from shapely import GeometryType
 
@@ -46,7 +46,7 @@ def _parse_version(group: zarr.Group, expect_attrs_key: bool) -> str | None:
     return version
 
 
-class SpatialDataFormat(CurrentFormat):
+class SpatialDataFormat(FormatV04):
     pass
 
 

src/spatialdata/_io/io_points.py

@@ -31,7 +31,8 @@ def _read_points(
     assert version is not None
     format = PointsFormats[version]
 
-    path = os.path.join(f._store.path, f.path, "points.parquet")
+    store_root = f.store_path.store.root
+    path = os.path.join(store_root, f.path, "points.parquet")
     # cache on remote file needed for parquet reader to work
     # TODO: allow reading in the metadata without caching all the data
     points = read_parquet("simplecache::" + path if path.startswith("http") else path)
@@ -57,7 +58,9 @@ def write_points(
     t = _get_transformations(points)
 
     points_groups = group.require_group(name)
-    path = Path(points_groups._store.path) / points_groups.path / "points.parquet"
+    store_root = points_groups.store_path.store.root
+    group_path = points_groups.path
+    path = Path(store_root) / group_path / "points.parquet"
 
     # The following code iterates through all columns in the 'points' DataFrame. If the column's datatype is
     # 'category', it checks whether the categories of this column are known. If not, it explicitly converts the

src/spatialdata/_io/io_shapes.py

@@ -50,7 +50,8 @@ def _read_shapes(
             geometry = from_ragged_array(typ, coords, offsets)
             geo_df = GeoDataFrame({"geometry": geometry}, index=index)
     elif isinstance(format, ShapesFormatV02):
-        path = Path(f._store.path) / f.path / "shapes.parquet"
+        store_root = f.store_path.store.root
+        path = Path(store_root) / f.path / "shapes.parquet"
         geo_df = read_parquet(path)
     else:
         raise ValueError(
@@ -93,7 +94,8 @@ def write_shapes(
         attrs = format.attrs_to_dict(geometry)
         attrs["version"] = format.spatialdata_format_version
     elif isinstance(format, ShapesFormatV02):
-        path = Path(shapes_group._store.path) / shapes_group.path / "shapes.parquet"
+        store_root = shapes_group.store_path.store.root
+        path = Path(store_root) / shapes_group.path / "shapes.parquet"
         shapes.to_parquet(path)
 
         attrs = format.attrs_to_dict(shapes.attrs)

src/spatialdata/_io/io_table.py

@@ -10,8 +10,8 @@
 from anndata import read_zarr as read_anndata_zarr
 from anndata._io.specs import write_elem as write_adata
 from ome_zarr.format import Format
-from zarr.errors import ArrayNotFoundError
 
+# from zarr.errors import ArrayNotFoundError  # removed in zarr 3.0
 from spatialdata._io._utils import BadFileHandleMethod, handle_read_errors
 from spatialdata._io.format import CurrentTablesFormat, TablesFormats, _parse_version
 from spatialdata._logging import logger
@@ -53,7 +53,12 @@ def _read_table(
         with handle_read_errors(
             on_bad_files=on_bad_files,
             location=f"{subgroup.path}/{table_name}",
-            exc_types=(JSONDecodeError, KeyError, ValueError, ArrayNotFoundError),
+            exc_types=(
+                JSONDecodeError,
+                KeyError,
+                ValueError,
+                # ArrayNotFoundError,  # removed in zarr 3.0
+            ),
         ):
             tables[table_name] = read_anndata_zarr(f_elem_store)