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

src/spatialdata/_io/io_points.py

@@ -6,13 +6,16 @@
 from dask.dataframe import DataFrame as DaskDataFrame
 from dask.dataframe import read_parquet
 from ome_zarr.format import Format
+from upath import UPath
 
 from spatialdata._io._utils import (
     _get_transformations_from_ngff_dict,
+    _resolve_zarr_store,
     _write_metadata,
     overwrite_coordinate_transformations_non_raster,
 )
 from spatialdata._io.format import CurrentPointsFormat, PointsFormats, _parse_version
+from spatialdata._store import ZarrStore, make_zarr_store, make_zarr_store_from_group, open_zarr_for_read
 from spatialdata.models import get_axes_names
 from spatialdata.transformations._utils import (
     _get_transformations,
@@ -21,21 +24,38 @@
 
 
 def _read_points(
-    store: str | Path,
+    store: str | Path | UPath | ZarrStore,
 ) -> DaskDataFrame:
-    """Read points from a zarr store."""
-    f = zarr.open(store, mode="r")
+    """Read points from a zarr store (path, hierarchical URI string, or remote ``UPath``)."""
+    zarr_store = store if isinstance(store, ZarrStore) else make_zarr_store(store)
+    resolved_store = _resolve_zarr_store(zarr_store.path)
+    f = open_zarr_for_read(resolved_store, as_group=False)
 
     version = _parse_version(f, 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)
+    parquet_store = zarr_store.child("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_store.arrow_path(),
+        filesystem=parquet_store.arrow_filesystem(),
+    )
     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"])
     _set_transformations(points, transformations)
@@ -68,8 +88,7 @@ def write_points(
     axes = get_axes_names(points)
     transformations = _get_transformations(points)
 
-    store_root = group.store_path.store.root
-    path = store_root / group.path / "points.parquet"
+    parquet_store = make_zarr_store_from_group(group).child("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
@@ -84,7 +103,10 @@ 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_store.arrow_path(),
+        filesystem=parquet_store.arrow_filesystem(),
+    )
 
     attrs = element_format.attrs_to_dict(points.attrs)
     attrs["version"] = element_format.spatialdata_format_version

src/spatialdata/_io/io_raster.py

@@ -16,17 +16,20 @@
 from ome_zarr.writer import write_labels as write_labels_ngff
 from ome_zarr.writer import write_multiscale as write_multiscale_ngff
 from ome_zarr.writer import write_multiscale_labels as write_multiscale_labels_ngff
+from upath import UPath
 from xarray import DataArray, DataTree
 
 from spatialdata._io._utils import (
     _get_transformations_from_ngff_dict,
+    _resolve_zarr_store,
     overwrite_coordinate_transformations_raster,
 )
 from spatialdata._io.format import (
     CurrentRasterFormat,
     RasterFormatType,
     get_ome_zarr_format,
 )
+from spatialdata._store import ZarrStore, make_zarr_store
 from spatialdata._utils import get_pyramid_levels
 from spatialdata.models._utils import get_channel_names
 from spatialdata.models.models import ATTRS_KEY
@@ -160,13 +163,14 @@ def _prepare_storage_options(
 
 
 def _read_multiscale(
-    store: str | Path, raster_type: Literal["image", "labels"], reader_format: Format
+    store: str | Path | UPath | ZarrStore, raster_type: Literal["image", "labels"], reader_format: Format
 ) -> DataArray | DataTree:
-    assert isinstance(store, str | Path)
     assert raster_type in ["image", "labels"]
+    zarr_store = store if isinstance(store, ZarrStore) else make_zarr_store(store)
+    resolved_store = _resolve_zarr_store(zarr_store.path)
 
     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

@@ -9,9 +9,11 @@
 from natsort import natsorted
 from ome_zarr.format import Format
 from shapely import from_ragged_array, to_ragged_array
+from upath import UPath
 
 from spatialdata._io._utils import (
     _get_transformations_from_ngff_dict,
+    _resolve_zarr_store,
     _write_metadata,
     overwrite_coordinate_transformations_non_raster,
 )
@@ -23,6 +25,7 @@
     ShapesFormatV03,
     _parse_version,
 )
+from spatialdata._store import ZarrStore, make_zarr_store, make_zarr_store_from_group, open_zarr_for_read
 from spatialdata.models import ShapesModel, get_axes_names
 from spatialdata.transformations._utils import (
     _get_transformations,
@@ -31,10 +34,12 @@
 
 
 def _read_shapes(
-    store: str | Path,
+    store: str | Path | UPath | ZarrStore,
 ) -> GeoDataFrame:
-    """Read shapes from a zarr store."""
-    f = zarr.open(store, mode="r")
+    """Read shapes from a zarr store (path, hierarchical URI string, or remote ``UPath``)."""
+    zarr_store = store if isinstance(store, ZarrStore) else make_zarr_store(store)
+    resolved_store = _resolve_zarr_store(zarr_store.path)
+    f = open_zarr_for_read(resolved_store, as_group=False)
     version = _parse_version(f, expect_attrs_key=True)
     assert version is not None
     shape_format = ShapesFormats[version]
@@ -54,9 +59,9 @@ def _read_shapes(
             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)
+        parquet_store = zarr_store.child("shapes.parquet")
+        with parquet_store.arrow_filesystem().open_input_file(parquet_store.arrow_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."
@@ -169,13 +174,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"
+    parquet_store = make_zarr_store_from_group(group).child("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 parquet_store.arrow_filesystem().open_output_stream(parquet_store.arrow_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

@@ -8,21 +8,26 @@
 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 upath import UPath
 
+from spatialdata._io._utils import _resolve_zarr_store
 from spatialdata._io.format import (
     CurrentTablesFormat,
     TablesFormats,
     TablesFormatV01,
     TablesFormatV02,
     _parse_version,
 )
+from spatialdata._store import ZarrStore, make_zarr_store, open_zarr_for_read
 from spatialdata.models import TableModel, get_table_keys
 
 
-def _read_table(store: str | Path) -> AnnData:
-    table = read_anndata_zarr(str(store))
+def _read_table(store: str | Path | UPath | ZarrStore) -> AnnData:
+    zarr_store = store if isinstance(store, ZarrStore) else make_zarr_store(store)
+    resolved_store = _resolve_zarr_store(zarr_store.path)
+    table = read_anndata_zarr(resolved_store)
 
-    f = zarr.open(store, mode="r")
+    f = open_zarr_for_read(resolved_store, as_group=False)
     version = _parse_version(f, expect_attrs_key=False)
     assert version is not None
     table_format = TablesFormats[version]