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,36 @@ 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],
+    element_name: str,
+    element_names: set[str],
+) -> dict[str, Any] | list[dict[str, Any]]:
+    element_raster_write_kwargs = None
+    if isinstance(raster_write_kwargs, dict) and (kwargs := raster_write_kwargs.get(element_name)):
+        element_raster_write_kwargs = kwargs
+
+    if not element_raster_write_kwargs:
+        if isinstance(raster_write_kwargs, dict):
+            for name in element_names:
+                raster_write_kwargs.pop(name, None)
+        if not raster_write_kwargs:
+            element_raster_write_kwargs = {}
+        elif isinstance(raster_write_kwargs, dict) and not all(
+            isinstance(x, (dict, list)) for x in raster_write_kwargs.values()
+        ):
+            element_raster_write_kwargs = raster_write_kwargs
+        elif isinstance(raster_write_kwargs, list):
+            if not all(isinstance(x, dict) for x in raster_write_kwargs):
+                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."
+                )
+            element_raster_write_kwargs = raster_write_kwargs
+        else:
+            raise ValueError(
+                f"Type of raster_write_kwargs should be either dict or list, got {type(raster_write_kwargs)}."
+            )
+    return element_raster_write_kwargs

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
@@ -1113,6 +1114,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,12 +1163,32 @@ 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
+            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
         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
             supported. If not specified, the compression will be `lz4` with compression level 5. Bytes are automatically
             ordered for more efficient compression.
         """
+        from spatialdata._core._utils import create_raster_element_kwargs
         from spatialdata._io._utils import _resolve_zarr_store, _validate_compressor_args
         from spatialdata._io.format import _parse_formats
 
@@ -1185,6 +1207,13 @@ def write(
         store.close()
 
         for element_type, element_name, element in self.gen_elements():
+            element_raster_write_kwargs = None
+            if element_type in ("images", "labels") and raster_write_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
+                )
+
             self._write_element(
                 element=element,
                 zarr_container_path=file_path,
@@ -1193,6 +1222,7 @@ def write(
                 overwrite=False,
                 parsed_formats=parsed,
                 shapes_geometry_encoding=shapes_geometry_encoding,
+                element_raster_write_kwargs=element_raster_write_kwargs,
                 raster_compressor=raster_compressor,
             )
 
@@ -1211,6 +1241,7 @@ def _write_element(
         overwrite: bool,
         parsed_formats: dict[str, SpatialDataFormatType] | None = None,
         shapes_geometry_encoding: Literal["WKB", "geoarrow"] | None = None,
+        element_raster_write_kwargs: JSONDict | list[JSONDict] | None = None,
         raster_compressor: dict[Literal["lz4", "zstd"], int] | None = None,
     ) -> None:
         from spatialdata._io.io_zarr import _get_groups_for_element
@@ -1250,6 +1281,7 @@ def _write_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":
@@ -1258,6 +1290,7 @@ def _write_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":
@@ -1289,6 +1322,7 @@ def write_element(
         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 +1342,25 @@ 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
+            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
          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
@@ -1319,6 +1372,7 @@ def write_element(
         If you pass a list of names, the elements will be written one by one. If an error occurs during the writing of
         an element, the writing of the remaining elements will not be attempted.
         """
+        from spatialdata._core._utils import create_raster_element_kwargs
         from spatialdata._io.format import _parse_formats
 
         parsed_formats = _parse_formats(formats=sdata_formats)
@@ -1331,6 +1385,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
@@ -1359,6 +1414,11 @@ def write_element(
 
         self._check_element_not_on_disk_with_different_type(element_type=element_type, element_name=element_name)
 
+        element_raster_write_kwargs = None
+        if element_type in ("images", "labels") and raster_write_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)
+
         self._write_element(
             element=element,
             zarr_container_path=self.path,
@@ -1367,6 +1427,7 @@ def write_element(
             overwrite=overwrite,
             parsed_formats=parsed_formats,
             shapes_geometry_encoding=shapes_geometry_encoding,
+            element_raster_write_kwargs=element_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,13 @@ def _write_raster(
         for c in channels:
             metadata["metadata"]["omero"]["channels"].append({"label": c})  # type: ignore[union-attr, index, call-overload]
 
+    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,