Support sharding through config and raster_write_kwargs

pyproject.toml

@@ -27,7 +27,6 @@ dependencies = [
     "click",
     "dask-image",
     "dask>=2026.3.0",
-    "distributed>=2026.3.0",
     "datashader",
     "fsspec[s3,http]",
     "geopandas>=0.14",
@@ -37,6 +36,7 @@ dependencies = [
     "numpy",
     "ome_zarr>=0.16.0",
     "pandas",
+    "platformdirs",
     "pooch",
     "pyarrow",
     "rich",
@@ -60,6 +60,9 @@ extra  = [
     "spatialdata-plot",
     "spatialdata-io",
 ]
+zarrs = [
+    "zarrs"
+]
 
 [dependency-groups]
 dev = [
@@ -71,6 +74,7 @@ test = [
     "pytest-mock",
     "pytest-xdist",
     "torch",
+    "zarrs",
 ]
 docs = [
     "sphinx>=4.5",

src/spatialdata/_core/_utils.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
+from typing import Any
 
 from anndata import AnnData
+from ome_zarr.types import JSONDict
 
 from spatialdata._core.spatialdata import SpatialData
 
@@ -164,3 +166,41 @@ def get_unique_name(name: str, attr: str, is_dataframe_column: bool = False) ->
         setattr(sanitized, attr, new_dict)
 
     return None if inplace else sanitized
+
+
+def create_raster_element_kwargs(
+    raster_write_kwargs: dict[str, JSONDict | list[JSONDict]] | list[JSONDict] | None,
+    element_name: str,
+    element_names: set[str],
+) -> dict[str, Any] | list[dict[str, Any]] | None:
+    """Normalize raster keyword arguments to the kwargs required by `zarr.create_array` for a single raster."""
+    if raster_write_kwargs is None:
+        return {}
+
+    kwargs_copy = raster_write_kwargs.copy()
+    if isinstance(kwargs_copy, dict):
+        element_write_kwargs: JSONDict | list[JSONDict] | None = kwargs_copy.get(element_name)
+        if element_write_kwargs:
+            return element_write_kwargs
+
+        # If we get here it means that we do not have kwargs with the specific element. We need to clear out kwargs
+        # that could be there of other elements.
+        for name in element_names:
+            kwargs_copy.pop(name, None)
+
+    # We return here if there are no kwargs after stripping all kwargs directly corresponding to a given element.
+    if not kwargs_copy:
+        return {}
+
+    if isinstance(kwargs_copy, dict) and not all(isinstance(x, (dict, list)) for x in kwargs_copy.values()):
+        return kwargs_copy
+
+    if isinstance(kwargs_copy, list):
+        if not all(isinstance(x, dict) for x in kwargs_copy):
+            raise ValueError(
+                "If passing raster_write_kwargs as list, it is assumed to be the storage "
+                "options for each scale of a multiscale raster as a dictionary."
+            )
+        return kwargs_copy
+
+    raise ValueError(f"Type of raster_write_kwargs should be either dict or list, got {type(kwargs_copy)}.")

src/spatialdata/_core/spatialdata.py

@@ -16,6 +16,7 @@
 from dask.dataframe import DataFrame as DaskDataFrame
 from dask.dataframe import Scalar
 from geopandas import GeoDataFrame
+from ome_zarr.types import JSONDict
 from shapely import MultiPolygon, Polygon
 from upath import UPath
 from xarray import DataArray, DataTree
@@ -29,9 +30,10 @@
     raise_validation_errors,
     validate_table_attr_keys,
 )
+from spatialdata._docs import docstring_parameter
 from spatialdata._logging import logger
 from spatialdata._types import ArrayLike, Raster_T
-from spatialdata._utils import _deprecation_alias
+from spatialdata._utils import _deprecation_alias, zarrs_context
 from spatialdata.models import (
     Image2DModel,
     Image3DModel,
@@ -57,6 +59,27 @@
         SpatialDataFormatType,
     )
 
+RASTER_WRITE_KWARGS_DOCS = """\
+    Storage options for raster elements. These options are passed to the zarr storage backend for writing and
+    can be provided in several formats:
+
+        1. Single dictionary
+            A dictionary containing all storage options applied globally.
+        2. Dictionary per raster element
+            A dictionary where:
+            - Keys = names of raster elements
+            - Values = storage options for each element
+                - For single-scale data: a dictionary
+                - For multiscale data: a list of dictionaries (one per scale)
+        3. List of dictionaries (multiscale only)
+            A list where each dictionary defines the storage options for one scale of a multiscale raster element.
+
+        Important Notes
+        - The available key–value pairs in these dictionaries depend on the Zarr format used for writing.
+        - For a full list of supported storage options, refer to:
+            https://zarr.readthedocs.io/en/stable/api/zarr/create/#zarr.create_array
+    """
+
 
 class SpatialData:
     """
@@ -1105,6 +1128,7 @@ def _validate_all_elements(self) -> None:
                         validate_table_attr_keys(element, location=element_path)
 
     @_deprecation_alias(format="sdata_formats", version="0.7.0")
+    @docstring_parameter(raster_write_kwargs=RASTER_WRITE_KWARGS_DOCS)
     def write(
         self,
         file_path: str | Path,
@@ -1113,6 +1137,7 @@ def write(
         update_sdata_path: bool = True,
         sdata_formats: SpatialDataFormatType | list[SpatialDataFormatType] | None = None,
         shapes_geometry_encoding: Literal["WKB", "geoarrow"] | None = None,
+        raster_write_kwargs: dict[str, JSONDict | list[JSONDict]] | list[JSONDict] | None = None,
         raster_compressor: dict[Literal["lz4", "zstd"], int] | None = None,
     ) -> None:
         """
@@ -1161,6 +1186,8 @@ def write(
         shapes_geometry_encoding
             Whether to use the WKB or geoarrow encoding for GeoParquet. See :meth:`geopandas.GeoDataFrame.to_parquet`
             for details. If None, uses the value from :attr:`spatialdata.settings.shapes_geometry_encoding`.
+        raster_write_kwargs
+        {RASTER_WRITE_KWARGS_DOCS}
         raster_compressor
             A lenght-1 dictionary with as key the type of compression to use for images and labels and as value the
             compression level which should be inclusive between 0 and 9. For compression, `lz4` and `zstd` are
@@ -1193,6 +1220,7 @@ def write(
                 overwrite=False,
                 parsed_formats=parsed,
                 shapes_geometry_encoding=shapes_geometry_encoding,
+                raster_write_kwargs=raster_write_kwargs,
                 raster_compressor=raster_compressor,
             )
 
@@ -1211,6 +1239,7 @@ def _write_element(
         overwrite: bool,
         parsed_formats: dict[str, SpatialDataFormatType] | None = None,
         shapes_geometry_encoding: Literal["WKB", "geoarrow"] | None = None,
+        raster_write_kwargs: dict[str, JSONDict | list[JSONDict] | Any] | list[JSONDict] | None = None,
         raster_compressor: dict[Literal["lz4", "zstd"], int] | None = None,
     ) -> None:
         from spatialdata._io.io_zarr import _get_groups_for_element
@@ -1244,51 +1273,63 @@ def _write_element(
 
             validate_element(element)
 
-        if element_type == "images":
-            write_image(
-                image=element,
-                group=element_group,
-                name=element_name,
-                element_format=parsed_formats["raster"],
-                raster_compressor=raster_compressor,
-            )
-        elif element_type == "labels":
-            write_labels(
-                labels=element,
-                group=root_group,
-                name=element_name,
-                element_format=parsed_formats["raster"],
-                raster_compressor=raster_compressor,
-            )
-        elif element_type == "points":
-            write_points(
-                points=element,
-                group=element_group,
-                element_format=parsed_formats["points"],
-            )
-        elif element_type == "shapes":
-            write_shapes(
-                shapes=element,
-                group=element_group,
-                element_format=parsed_formats["shapes"],
-                geometry_encoding=shapes_geometry_encoding,
-            )
-        elif element_type == "tables":
-            write_table(
-                table=element,
-                group=element_type_group,
-                name=element_name,
-                element_format=parsed_formats["tables"],
-            )
-        else:
-            raise ValueError(f"Unknown element type: {element_type}")
+        element_raster_write_kwargs = None
+        if element_type in ("images", "labels") and raster_write_kwargs:
+            from spatialdata._core._utils import create_raster_element_kwargs
+
+            element_names = set(self.images.keys()).union(self.labels.keys())
+            element_raster_write_kwargs = create_raster_element_kwargs(raster_write_kwargs, element_name, element_names)
+
+        with zarrs_context():
+            if element_type == "images":
+                write_image(
+                    image=element,
+                    group=element_group,
+                    name=element_name,
+                    element_format=parsed_formats["raster"],
+                    storage_options=element_raster_write_kwargs,
+                    raster_compressor=raster_compressor,
+                )
+            elif element_type == "labels":
+                write_labels(
+                    labels=element,
+                    group=root_group,
+                    name=element_name,
+                    element_format=parsed_formats["raster"],
+                    storage_options=element_raster_write_kwargs,
+                    raster_compressor=raster_compressor,
+                )
+            elif element_type == "points":
+                write_points(
+                    points=element,
+                    group=element_group,
+                    element_format=parsed_formats["points"],
+                )
+            elif element_type == "shapes":
+                write_shapes(
+                    shapes=element,
+                    group=element_group,
+                    element_format=parsed_formats["shapes"],
+                    geometry_encoding=shapes_geometry_encoding,
+                )
+            elif element_type == "tables":
+                write_table(
+                    table=element,
+                    group=element_type_group,
+                    name=element_name,
+                    element_format=parsed_formats["tables"],
+                )
+            else:
+                raise ValueError(f"Unknown element type: {element_type}")
 
+    @docstring_parameter(raster_write_kwargs=RASTER_WRITE_KWARGS_DOCS)
     def write_element(
         self,
         element_name: str | list[str],
         overwrite: bool = False,
         sdata_formats: SpatialDataFormatType | list[SpatialDataFormatType] | None = None,
         shapes_geometry_encoding: Literal["WKB", "geoarrow"] | None = None,
+        raster_write_kwargs: dict[str, JSONDict | list[JSONDict] | Any] | list[JSONDict] | None = None,
         raster_compressor: dict[Literal["lz4", "zstd"], int] | None = None,
     ) -> None:
         """
@@ -1308,6 +1349,8 @@ def write_element(
         shapes_geometry_encoding
             Whether to use the WKB or geoarrow encoding for GeoParquet. See :meth:`geopandas.GeoDataFrame.to_parquet`
             for details. If None, uses the value from :attr:`spatialdata.settings.shapes_geometry_encoding`.
+        raster_write_kwargs
+        {RASTER_WRITE_KWARGS_DOCS}
          raster_compressor
             A lenght-1 dictionary with as key the type of compression to use for images and labels and as value the
             compression level which should be inclusive between 0 and 9. For compression, `lz4` and `zstd` are
@@ -1331,6 +1374,7 @@ def write_element(
                     overwrite=overwrite,
                     sdata_formats=sdata_formats,
                     shapes_geometry_encoding=shapes_geometry_encoding,
+                    raster_write_kwargs=raster_write_kwargs,
                     raster_compressor=raster_compressor,
                 )
             return
@@ -1367,6 +1411,7 @@ def write_element(
             overwrite=overwrite,
             parsed_formats=parsed_formats,
             shapes_geometry_encoding=shapes_geometry_encoding,
+            raster_write_kwargs=raster_write_kwargs,
             raster_compressor=raster_compressor,
         )
         # After every write, metadata should be consolidated, otherwise this can lead to IO problems like when deleting.

src/spatialdata/_io/io_raster.py

@@ -148,13 +148,13 @@ def _prepare_storage_options(
         return None
     if isinstance(storage_options, dict):
         prepared = dict(storage_options)
-        if "chunks" in prepared:
+        if "chunks" in prepared and prepared["chunks"] is not None:
             prepared["chunks"] = _normalize_explicit_chunks(prepared["chunks"])
         return prepared
 
     prepared_options = [dict(options) for options in storage_options]
     for options in prepared_options:
-        if "chunks" in options:
+        if "chunks" in options and options["chunks"] is not None:
             options["chunks"] = _normalize_explicit_chunks(options["chunks"])
     return prepared_options
 
@@ -284,6 +284,19 @@ def _write_raster(
     raster_format
         The format used to write the raster data.
     storage_options
+        Storage options for raster elements, which have been extracted from potentially mixed kwargs dict by
+        `create_raster_element_kwargs`. These options are passed to the zarr storage backend for writing and can be
+        provided in several formats:
+
+            1. Single dictionary
+                A dictionary containing all storage options applied to the raster, either single or multiscale.
+            2. List of dictionaries (multiscale only)
+                A list where each dictionary defines the storage options for one scale of the multiscale raster element.
+
+            Important Notes
+            - The available key–value pairs in these dictionaries depend on the Zarr format used for writing.
+            - For a full list of supported storage options, refer to:
+                https://zarr.readthedocs.io/en/stable/api/zarr/create/#zarr.create_array
         Additional options for writing the raster data, like chunks and compression.
     raster_compressor
         Compression settings as a len-1 dictionary with a single key-value {compression: compression level} pair
@@ -292,6 +305,10 @@ def _write_raster(
     metadata
         Additional metadata for the raster element
     """
+    from dataclasses import asdict
+
+    from spatialdata import settings
+
     if raster_type not in ["image", "labels"]:
         raise ValueError(f"{raster_type} is not a valid raster type. Must be 'image' or 'labels'.")
     # "name" and "label_metadata" are only used for labels. "name" is written in write_multiscale_ngff() but ignored in
@@ -308,6 +325,14 @@ def _write_raster(
         for c in channels:
             metadata["metadata"]["omero"]["channels"].append({"label": c})  # type: ignore[union-attr, index, call-overload]
 
+    # Prefixing with raster_ to account for anndata chunks / shards that will be supported in the future.
+    base_options = {k.split("_")[1]: v for k, v in asdict(settings).items() if k in ("raster_chunks", "raster_shards")}
+
+    if isinstance(storage_options, list):
+        storage_options = [{**base_options, **x} for x in storage_options]
+    else:
+        storage_options = {**base_options, **(storage_options or {})}
+
     if isinstance(raster_data, DataArray):
         _write_raster_dataarray(
             raster_type,