feat: generalize Zarr I/O in spatialdata to local and add remote (fsspec / UPath) backends

src/spatialdata/_io/io_points.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from pathlib import Path
-
 import zarr
 from dask.dataframe import DataFrame as DaskDataFrame
 from dask.dataframe import read_parquet
@@ -13,34 +11,43 @@
     overwrite_coordinate_transformations_non_raster,
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
+from spatialdata._store import arrow_fs_and_path
 from spatialdata.models import get_axes_names
 from spatialdata.transformations._utils import (
     _get_transformations,
     _set_transformations,
 )
 
 
-def _read_points(
-    store: str | Path,
-) -> DaskDataFrame:
-    """Read points from a zarr store."""
-    f = zarr.open(Path(store), mode="r")  # Path avoids zarr v3 URL-parsing special chars (e.g. #) in names
-
-    version = _parse_version(f, expect_attrs_key=True)
+def _read_points(group: zarr.Group) -> DaskDataFrame:
+    """Read a points element from an open zarr group."""
+    version = _parse_version(group, expect_attrs_key=True)
     assert version is not None
     points_format = PointsFormats[version]
 
-    store_root = f.store_path.store.root
-    path = 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::" + str(path) if str(path).startswith("http") else path)
+    fs, parquet_path = arrow_fs_and_path(group, "points.parquet")
+    # Passing filesystem= to read_parquet makes pyarrow convert dictionary columns into pandas
+    # categoricals eagerly per partition and marks them known=True with an empty category list.
+    # This happens for ANY pyarrow filesystem (both LocalFileSystem and PyFileSystem(FSSpecHandler(.))
+    # return the same broken categorical), so it is a property of the filesystem= handoff itself,
+    # not of local-vs-remote. Left as is, it would make write_points' cat.as_known() a no-op and
+    # the next to_parquet(filesystem=.) would fail with a per-partition schema mismatch
+    # (dictionary<values=null> vs dictionary<values=string>). We demote the categoricals back to
+    # "unknown" right here so that write_points recomputes categories consistently across partitions.
+    # TODO: allow reading in the metadata without materializing the data.
+    points = read_parquet(parquet_path, filesystem=fs)
     assert isinstance(points, DaskDataFrame)
+    for column_name in points.columns:
+        c = points[column_name]
+        if c.dtype == "category" and c.cat.known:
+            points[column_name] = c.cat.as_unknown()
+    if points.index.name == "__null_dask_index__":
+        points = points.rename_axis(None)
 
-    transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
+    transformations = _get_transformations_from_ngff_dict(group.attrs.asdict()["coordinateTransformations"])
     _set_transformations(points, transformations)
 
-    attrs = points_format.attrs_from_dict(f.attrs.asdict())
+    attrs = points_format.attrs_from_dict(group.attrs.asdict())
     if len(attrs):
         points.attrs["spatialdata_attrs"] = attrs
     return points
@@ -69,8 +76,7 @@ def write_points(
     transformations = _get_transformations(points)
     assert transformations is not None  # mypy: validate_element() in _write_element guarantees this
 
-    store_root = group.store_path.store.root
-    path = store_root / group.path / "points.parquet"
+    fs, parquet_path = arrow_fs_and_path(group, "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
@@ -85,7 +91,7 @@ def write_points(
 
     points_without_transform = points.copy()
     del points_without_transform.attrs["transform"]
-    points_without_transform.to_parquet(path)
+    points_without_transform.to_parquet(parquet_path, filesystem=fs)
 
     attrs = element_format.attrs_to_dict(points.attrs)
     attrs["version"] = element_format.spatialdata_format_version

src/spatialdata/_io/io_raster.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from pathlib import Path
 from typing import Any, Literal, TypeGuard, cast
 
 import dask.array as da
@@ -27,6 +26,7 @@
     RasterFormatType,
     get_ome_zarr_format,
 )
+from spatialdata._store import store_from_group
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import get_channel_names
 from spatialdata.models.models import ATTRS_KEY
@@ -160,13 +160,15 @@ def _prepare_storage_options(
 
 
 def _read_multiscale(
-    store: str | Path, raster_type: Literal["image", "labels"], reader_format: Format
+    group: zarr.Group, raster_type: Literal["image", "labels"], reader_format: Format
 ) -> DataArray | DataTree:
-    assert isinstance(store, str | Path)
     assert raster_type in ["image", "labels"]
+    # ome_zarr.io.ZarrLocation needs a store rooted at this group's location, not at the
+    # SpatialData container root, so we re-root the parent store at ``group.path``.
+    resolved_store = store_from_group(group, read_only=True)
 
     nodes: list[Node] = []
-    image_loc = ZarrLocation(store, fmt=reader_format)
+    image_loc = ZarrLocation(resolved_store, fmt=reader_format)
     if exists := image_loc.exists():
         image_reader = Reader(image_loc)()
         image_nodes = list(image_reader)

src/spatialdata/_io/io_shapes.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from pathlib import Path
 from typing import Any, Literal
 
 import numpy as np
@@ -23,46 +22,44 @@
     ShapesFormatV03,
     _parse_version,
 )
+from spatialdata._store import arrow_fs_and_path
 from spatialdata.models import ShapesModel, get_axes_names
 from spatialdata.transformations._utils import (
     _get_transformations,
     _set_transformations,
 )
 
 
-def _read_shapes(
-    store: str | Path,
-) -> GeoDataFrame:
-    """Read shapes from a zarr store."""
-    f = zarr.open(Path(store), mode="r")  # Path avoids zarr v3 URL-parsing special chars (e.g. #) in names
-    version = _parse_version(f, expect_attrs_key=True)
+def _read_shapes(group: zarr.Group) -> GeoDataFrame:
+    """Read a shapes element from an open zarr group."""
+    version = _parse_version(group, expect_attrs_key=True)
     assert version is not None
     shape_format = ShapesFormats[version]
 
     if isinstance(shape_format, ShapesFormatV01):
-        coords = np.array(f["coords"])
-        index = np.array(f["Index"])
-        typ = shape_format.attrs_from_dict(f.attrs.asdict())
+        coords = np.array(group["coords"])
+        index = np.array(group["Index"])
+        typ = shape_format.attrs_from_dict(group.attrs.asdict())
         if typ.name == "POINT":
-            radius = np.array(f["radius"])
+            radius = np.array(group["radius"])
             geometry = from_ragged_array(typ, coords)
             geo_df = GeoDataFrame({"geometry": geometry, "radius": radius}, index=index)
         else:
-            offsets_keys = [k for k in f if k.startswith("offset")]
+            offsets_keys = [k for k in group if k.startswith("offset")]
             offsets_keys = natsorted(offsets_keys)
-            offsets = tuple(np.array(f[k]).flatten() for k in offsets_keys)
+            offsets = tuple(np.array(group[k]).flatten() for k in offsets_keys)
             geometry = from_ragged_array(typ, coords, offsets)
             geo_df = GeoDataFrame({"geometry": geometry}, index=index)
     elif isinstance(shape_format, ShapesFormatV02 | ShapesFormatV03):
-        store_root = f.store_path.store.root
-        path = Path(store_root) / f.path / "shapes.parquet"
-        geo_df = read_parquet(path)
+        fs, parquet_path = arrow_fs_and_path(group, "shapes.parquet")
+        with fs.open_input_file(parquet_path) as src:
+            geo_df = read_parquet(src)
     else:
         raise ValueError(
             f"Unsupported shapes format {shape_format} from version {version}. Please update the spatialdata library."
         )
 
-    transformations = _get_transformations_from_ngff_dict(f.attrs.asdict()["coordinateTransformations"])
+    transformations = _get_transformations_from_ngff_dict(group.attrs.asdict()["coordinateTransformations"])
     _set_transformations(geo_df, transformations)
     return geo_df
 
@@ -168,13 +165,13 @@ def _write_shapes_v02_v03(
     """
     from spatialdata.models._utils import TRANSFORM_KEY
 
-    store_root = group.store_path.store.root
-    path = store_root / group.path / "shapes.parquet"
+    fs, parquet_path = arrow_fs_and_path(group, "shapes.parquet")
 
     # Temporarily remove transformations from attrs to avoid serialization issues
     transforms = shapes.attrs[TRANSFORM_KEY]
     del shapes.attrs[TRANSFORM_KEY]
-    shapes.to_parquet(path, geometry_encoding=geometry_encoding)
+    with fs.open_output_stream(parquet_path) as sink:
+        shapes.to_parquet(sink, geometry_encoding=geometry_encoding)
     shapes.attrs[TRANSFORM_KEY] = transforms
 
     attrs = element_format.attrs_to_dict(shapes.attrs)

src/spatialdata/_io/io_table.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from pathlib import Path
-
 import numpy as np
 import zarr
 from anndata import AnnData
@@ -19,16 +17,16 @@
 from spatialdata.models import TableModel, get_table_keys
 
 
-def _read_table(store: str | Path) -> AnnData:
-    table = read_anndata_zarr(str(store))
+def _read_table(group: zarr.Group) -> AnnData:
+    """Read a table element from an open zarr group."""
+    # anndata's read_zarr accepts a StoreLike; pass the group's store sub-rooted at group.path.
+    # The simplest portable way: pass the group itself, which anndata supports.
+    table = read_anndata_zarr(group)
 
-    f = zarr.open(Path(store), mode="r")  # Path avoids zarr v3 URL-parsing special chars (e.g. #) in names
-    version = _parse_version(f, expect_attrs_key=False)
+    version = _parse_version(group, expect_attrs_key=False)
     assert version is not None
     table_format = TablesFormats[version]
 
-    f.store.close()
-
     if isinstance(table_format, TablesFormatV01 | TablesFormatV02):
         if TableModel.ATTRS_KEY in table.uns:
             # fill out eventual missing attributes that has been omitted because their value was None