Support for consolidated remote zarr (#278)

Co-authored-by: Luca Marconato <[email protected]>

src/spatialdata/__main__.py

@@ -3,39 +3,48 @@
 
 This module provides command line interface (CLI) interactions for the SpatialData library, allowing users to perform
 various operations through a terminal. Currently, it implements the "peek" function, which allows users to inspect
-the contents of a SpatialData .zarr file. Additional CLI functionalities will be implemented in the future.
+the contents of a SpatialData .zarr dataset. Additional CLI functionalities will be implemented in the future.
 """
-import os
+from typing import Literal
 
 import click
 
 
-@click.command(help="Peek inside the SpatialData .zarr file")
+@click.command(help="Peek inside the SpatialData .zarr dataset")
 @click.argument("path", default=False, type=str)
-def peek(path: str) -> None:
+@click.argument("selection", type=click.Choice(["images", "labels", "points", "shapes", "table"]), nargs=-1)
+def peek(path: str, selection: tuple[Literal["images", "labels", "points", "shapes", "table"]]) -> None:
     """
-    Peek inside the SpatialData .zarr file.
+    Peek inside the SpatialData .zarr dataset.
 
-    This function takes a path to a .zarr file, checks if it is a valid directory, and then reads and prints
-    its contents using the SpatialData library.
+    This function takes a path to a local or remote .zarr dataset, reads and prints
+    its contents using the SpatialData library. If any ValueError is raised, it is caught and printed to the
+    terminal along with a help message.
 
     Parameters
     ----------
     path
-        The path to the .zarr file to be inspected.
+        The path to the .zarr dataset to be inspected.
+    selection
+        Optional, a list of keys (among images, labels, points, shapes, table) to load only a subset of the dataset.
+        Example: `python -m spatialdata peek data.zarr images labels`
     """
-    if not os.path.isdir(path):
+    import spatialdata as sd
+
+    try:
+        sdata = sd.SpatialData.read(path, selection=selection)
+        print(sdata)  # noqa: T201
+    except ValueError as e:
+        # checking if a valid path was provided is difficult given the various ways in which
+        # a possibly remote path and storage access options can be specified
+        # so we just catch the ValueError and print a help message
+        print(e)  # noqa: T201
         print(  # noqa: T201
             f"Error: .zarr storage not found at {path}. Please specify a valid OME-NGFF spatial data (.zarr) file. "
-            "Example "
-            '"python -m '
-            'spatialdata peek data.zarr"'
+            "Examples "
+            '"python -m spatialdata peek data.zarr"'
+            '"python -m spatialdata peek https://remote/.../data.zarr labels table"'
         )
-    else:
-        import spatialdata as sd
-
-        sdata = sd.SpatialData.read(path)
-        print(sdata)  # noqa: T201
 
 
 @click.group()

src/spatialdata/_core/spatialdata.py

@@ -9,13 +9,13 @@
 
 import zarr
 from anndata import AnnData
+from dask.dataframe import read_parquet
 from dask.dataframe.core import DataFrame as DaskDataFrame
 from dask.delayed import Delayed
 from geopandas import GeoDataFrame
 from multiscale_spatial_image.multiscale_spatial_image import MultiscaleSpatialImage
 from ome_zarr.io import parse_url
 from ome_zarr.types import JSONDict
-from pyarrow.parquet import read_table
 from spatial_image import SpatialImage
 
 from spatialdata._io import (
@@ -917,6 +917,7 @@ def write(
         file_path: str | Path,
         storage_options: JSONDict | list[JSONDict] | None = None,
         overwrite: bool = False,
+        consolidate_metadata: bool = True,
     ) -> None:
         """Write the SpatialData object to Zarr."""
         if isinstance(file_path, str):
@@ -1052,6 +1053,12 @@ def write(
             self.path = None
             raise e
 
+        if consolidate_metadata:
+            # consolidate metadata to more easily support remote reading
+            # bug in zarr, 'zmetadata' is written instead of '.zmetadata'
+            # see discussion https://github.com/zarr-developers/zarr-python/issues/1121
+            zarr.consolidate_metadata(store, metadata_key=".zmetadata")
+
         # old code to support overwriting the backing file
         # if target_path is not None:
         #     if os.path.isdir(file_path):
@@ -1133,10 +1140,24 @@ def table(self) -> None:
             del root["table/table"]
 
     @staticmethod
-    def read(file_path: str) -> SpatialData:
+    def read(file_path: str, selection: tuple[str] | None = None) -> SpatialData:
+        """
+        Read a SpatialData object from a Zarr storage (on-disk or remote).
+
+        Parameters
+        ----------
+        file_path
+            The path or URL to the Zarr storage.
+        selection
+            The elements to read (images, labels, points, shapes, table). If None, all elements are read.
+
+        Returns
+        -------
+        The SpatialData object.
+        """
         from spatialdata import read_zarr
 
-        return read_zarr(file_path)
+        return read_zarr(file_path, selection=selection)
 
     @property
     def images(self) -> dict[str, SpatialImage | MultiscaleSpatialImage]:
@@ -1238,7 +1259,7 @@ def h(s: str) -> str:
                                 assert isinstance(t, tuple)
                                 assert len(t) == 1
                                 parquet_file = t[0]
-                                table = read_table(parquet_file)
+                                table = read_parquet(parquet_file)
                                 length = len(table)
                             else:
                                 # length = len(v)

src/spatialdata/_io/io_points.py

@@ -1,3 +1,4 @@
+import os
 from collections.abc import MutableMapping
 from pathlib import Path
 from typing import Union
@@ -29,8 +30,10 @@ def _read_points(
     assert isinstance(store, (str, Path))
     f = zarr.open(store, mode="r")
 
-    path = Path(f._store.path) / f.path / "points.parquet"
-    table = read_parquet(path)
+    path = os.path.join(f._store.path, 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
+    table = read_parquet("simplecache::" + path if "http" in path else path)
     assert isinstance(table, DaskDataFrame)
 
     transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])

src/spatialdata/_io/io_zarr.py

@@ -1,74 +1,145 @@
 import logging
+import os
 from pathlib import Path
 from typing import Optional, Union
 
 import numpy as np
 import zarr
 from anndata import AnnData
 from anndata import read_zarr as read_anndata_zarr
+from anndata.experimental import read_elem
 
 from spatialdata import SpatialData
 from spatialdata._io._utils import ome_zarr_logger
 from spatialdata._io.io_points import _read_points
 from spatialdata._io.io_raster import _read_multiscale
 from spatialdata._io.io_shapes import _read_shapes
+from spatialdata._logging import logger
 from spatialdata.models import TableModel
 
 
-def read_zarr(store: Union[str, Path, zarr.Group]) -> SpatialData:
-    if isinstance(store, str):
-        store = Path(store)
+def _open_zarr_store(store: Union[str, Path, zarr.Group]) -> tuple[zarr.Group, str]:
+    """
+    Open a zarr store (on-disk or remote) and return the zarr.Group object and the path to the store.
+
+    Parameters
+    ----------
+    store
+        Path to the zarr store (on-disk or remote) or a zarr.Group object.
+
+    Returns
+    -------
+    A tuple of the zarr.Group object and the path to the store.
+    """
+    f = store if isinstance(store, zarr.Group) else zarr.open(store, mode="r")
+    # workaround: .zmetadata is being written as zmetadata (https://github.com/zarr-developers/zarr-python/issues/1121)
+    if isinstance(store, (str, Path)) and str(store).startswith("http") and len(f) == 0:
+        f = zarr.open_consolidated(store, mode="r", metadata_key="zmetadata")
+    f_store_path = f.store.store.path if isinstance(f.store, zarr.storage.ConsolidatedMetadataStore) else f.store.path
+    return f, f_store_path
+
+
+def read_zarr(store: Union[str, Path, zarr.Group], selection: Optional[tuple[str]] = None) -> SpatialData:
+    """
+    Read a SpatialData dataset from a zarr store (on-disk or remote).
+
+    Parameters
+    ----------
+    store
+        Path to the zarr store (on-disk or remote) or a zarr.Group object.
+
+    selection
+        List of elements to read from the zarr store (images, labels, points, shapes, table). If None, all elements are
+        read.
+
+    Returns
+    -------
+    A SpatialData object.
+    """
+    f, f_store_path = _open_zarr_store(store)
 
-    f = zarr.open(store, mode="r")
     images = {}
     labels = {}
     points = {}
     table: Optional[AnnData] = None
     shapes = {}
 
+    selector = {"images", "labels", "points", "shapes", "table"} if not selection else set(selection or [])
+    logger.debug(f"Reading selection {selector}")
+
     # read multiscale images
-    images_store = store / "images"
-    if images_store.exists():
-        f = zarr.open(images_store, mode="r")
-        for k in f:
-            f_elem = f[k].name
-            f_elem_store = f"{images_store}{f_elem}"
-            images[k] = _read_multiscale(f_elem_store, raster_type="image")
+    if "images" in selector and "images" in f:
+        group = f["images"]
+        count = 0
+        for subgroup_name in group:
+            if Path(subgroup_name).name.startswith("."):
+                # skip hidden files like .zgroup or .zmetadata
+                continue
+            f_elem = group[subgroup_name]
+            f_elem_store = os.path.join(f_store_path, f_elem.path)
+            element = _read_multiscale(f_elem_store, raster_type="image")
+            images[subgroup_name] = element
+            count += 1
+        logger.debug(f"Found {count} elements in {group}")
 
     # read multiscale labels
     with ome_zarr_logger(logging.ERROR):
-        labels_store = store / "labels"
-        if labels_store.exists():
-            f = zarr.open(labels_store, mode="r")
-            for k in f:
-                f_elem = f[k].name
-                f_elem_store = f"{labels_store}{f_elem}"
-                labels[k] = _read_multiscale(f_elem_store, raster_type="labels")
+        if "labels" in selector and "labels" in f:
+            group = f["labels"]
+            count = 0
+            for subgroup_name in group:
+                if Path(subgroup_name).name.startswith("."):
+                    # skip hidden files like .zgroup or .zmetadata
+                    continue
+                f_elem = group[subgroup_name]
+                f_elem_store = os.path.join(f_store_path, f_elem.path)
+                labels[subgroup_name] = _read_multiscale(f_elem_store, raster_type="labels")
+                count += 1
+            logger.debug(f"Found {count} elements in {group}")
 
     # now read rest of the data
-    points_store = store / "points"
-    if points_store.exists():
-        f = zarr.open(points_store, mode="r")
-        for k in f:
-            f_elem = f[k].name
-            f_elem_store = f"{points_store}{f_elem}"
-            points[k] = _read_points(f_elem_store)
-
-    shapes_store = store / "shapes"
-    if shapes_store.exists():
-        f = zarr.open(shapes_store, mode="r")
-        for k in f:
-            f_elem = f[k].name
-            f_elem_store = f"{shapes_store}{f_elem}"
-            shapes[k] = _read_shapes(f_elem_store)
-
-    table_store = store / "table"
-    if table_store.exists():
-        f = zarr.open(table_store, mode="r")
-        for k in f:
-            f_elem = f[k].name
-            f_elem_store = f"{table_store}{f_elem}"
-            table = read_anndata_zarr(f_elem_store)
+    if "points" in selector and "points" in f:
+        group = f["points"]
+        count = 0
+        for subgroup_name in group:
+            f_elem = group[subgroup_name]
+            if Path(subgroup_name).name.startswith("."):
+                # skip hidden files like .zgroup or .zmetadata
+                continue
+            f_elem_store = os.path.join(f_store_path, f_elem.path)
+            points[subgroup_name] = _read_points(f_elem_store)
+            count += 1
+        logger.debug(f"Found {count} elements in {group}")
+
+    if "shapes" in selector and "shapes" in f:
+        group = f["shapes"]
+        count = 0
+        for subgroup_name in group:
+            if Path(subgroup_name).name.startswith("."):
+                # skip hidden files like .zgroup or .zmetadata
+                continue
+            f_elem = group[subgroup_name]
+            f_elem_store = os.path.join(f_store_path, f_elem.path)
+            shapes[subgroup_name] = _read_shapes(f_elem_store)
+            count += 1
+        logger.debug(f"Found {count} elements in {group}")
+
+    if "table" in selector and "table" in f:
+        group = f["table"]
+        count = 0
+        for subgroup_name in group:
+            if Path(subgroup_name).name.startswith("."):
+                # skip hidden files like .zgroup or .zmetadata
+                continue
+            f_elem = group[subgroup_name]
+            f_elem_store = os.path.join(f_store_path, f_elem.path)
+            if isinstance(f.store, zarr.storage.ConsolidatedMetadataStore):
+                table = read_elem(f_elem)
+                # we can replace read_elem with read_anndata_zarr after this PR gets into a release (>= 0.6.5)
+                # https://github.com/scverse/anndata/pull/1057#pullrequestreview-1530623183
+                # table = read_anndata_zarr(f_elem)
+            else:
+                table = read_anndata_zarr(f_elem_store)
             if TableModel.ATTRS_KEY in table.uns:
                 # fill out eventual missing attributes that has been omitted because their value was None
                 attrs = table.uns[TableModel.ATTRS_KEY]
@@ -81,6 +152,8 @@ def read_zarr(store: Union[str, Path, zarr.Group]) -> SpatialData:
                 # fix type for region
                 if "region" in attrs and isinstance(attrs["region"], np.ndarray):
                     attrs["region"] = attrs["region"].tolist()
+            count += 1
+        logger.debug(f"Found {count} elements in {group}")
 
     sdata = SpatialData(
         images=images,

src/spatialdata/_logging.py

@@ -1,16 +1,18 @@
 import logging
+import os
 
 
 def _setup_logger() -> "logging.Logger":
     from rich.console import Console
     from rich.logging import RichHandler
 
     logger = logging.getLogger(__name__)
-    logger.setLevel(logging.INFO)
+    level = os.environ.get("LOGLEVEL", logging.INFO)
+    logger.setLevel(level=level)
     console = Console(force_terminal=True)
     if console.is_jupyter is True:
         console.is_jupyter = False
-    ch = RichHandler(show_path=False, console=console, show_time=False)
+    ch = RichHandler(show_path=False, console=console, show_time=logger.level == logging.DEBUG)
     logger.addHandler(ch)
 
     # this prevents double outputs