Support sharding through config and raster_write_kwargs

.github/workflows/test.yaml

@@ -19,7 +19,7 @@ jobs:
             fail-fast: false
             matrix:
                 include:
-                    - {os: windows-latest, python: "3.11", dask-version: "2025.12.0", name: "Dask 2025.12.0"}
+                    - {os: windows-latest, python: "3.11", dask-version: "2026.3.0", name: "Dask 2026.3.0"}
                     - {os: windows-latest, python: "3.13", dask-version: "latest", name: "Dask latest"}
                     - {os: ubuntu-latest, python: "3.11", dask-version: "latest", name: "Dask latest"}
                     - {os: ubuntu-latest, python: "3.13", dask-version: "latest", name: "Dask latest"}

pyproject.toml

@@ -26,8 +26,8 @@ dependencies = [
     "annsel>=0.1.2",
     "click",
     "dask-image",
-    "dask>=2025.12.0,<2026.1.2",
-    "distributed<2026.1.2",
+    "dask>=2026.3.0",
+    "distributed>=2026.3.0",
     "datashader",
     "fsspec[s3,http]",
     "geopandas>=0.14",
@@ -50,6 +50,7 @@ dependencies = [
     "xarray>=2024.10.0",
     "xarray-spatial>=0.3.5",
     "zarr>=3.0.0",
+    "zarrs",
 ]
 [project.optional-dependencies]
 torch = [
@@ -62,6 +63,9 @@ extra  = [
 ]
 
 [dependency-groups]
+sharding = [
+    "zarrs",
+]
 dev = [
     "bump2version",
 ]

src/spatialdata/__init__.py

@@ -131,6 +131,10 @@
     "settings",
 ]
 
+import zarr
+
+zarr.config.set({"codec_pipeline.path": "zarrs.ZarrsCodecPipeline"})
+
 
 def __getattr__(name: str) -> Any:
     if name in _submodules:

src/spatialdata/_core/_utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
+from typing import Any
 
 from anndata import AnnData
 
@@ -164,3 +165,26 @@ 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, dict[str, Any] | list[dict[str, Any]]] | list[dict[str, Any]],
+    element_name: str,
+) -> dict[str, Any] | list[dict[str, Any]]:
+
+    if isinstance(raster_write_kwargs, dict) and (kwargs := raster_write_kwargs.get(element_name)):
+        element_raster_write_kwargs = 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

@@ -1108,6 +1108,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, dict[str, Any] | list[dict[str, Any]]] | list[dict[str, Any]] | None = None,
     ) -> None:
         """
         Write the `SpatialData` object to a Zarr store.
@@ -1155,7 +1156,27 @@ 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
+        """
+        from spatialdata._core._utils import create_raster_element_kwargs
         from spatialdata._io._utils import _resolve_zarr_store
         from spatialdata._io.format import _parse_formats
 
@@ -1173,6 +1194,10 @@ 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_raster_write_kwargs = create_raster_element_kwargs(raster_write_kwargs, element_name)
+
             self._write_element(
                 element=element,
                 zarr_container_path=file_path,
@@ -1181,6 +1206,7 @@ def write(
                 overwrite=False,
                 parsed_formats=parsed,
                 shapes_geometry_encoding=shapes_geometry_encoding,
+                element_raster_write_kwargs=element_raster_write_kwargs,
             )
 
         if self.path != file_path and update_sdata_path:
@@ -1198,6 +1224,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: dict[str, Any] | list[dict[str, Any]] | None = None,
     ) -> None:
         from spatialdata._io.io_zarr import _get_groups_for_element
 
@@ -1231,13 +1258,15 @@ def _write_element(
                 group=element_group,
                 name=element_name,
                 element_format=parsed_formats["raster"],
+                storage_options=element_raster_write_kwargs,
             )
         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,
             )
         elif element_type == "points":
             write_points(
@@ -1268,6 +1297,9 @@ 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, dict[str, Any] | list[dict[str, Any]] | Any]
+        | list[dict[str, Any]]
+        | None = None,
     ) -> None:
         """
         Write a single element, or a list of elements, to the Zarr store used for backing.
@@ -1286,12 +1318,32 @@ 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
 
         Notes
         -----
         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 +1383,10 @@ 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_raster_write_kwargs = create_raster_element_kwargs(raster_write_kwargs, element_name)
+
         self._write_element(
             element=element,
             zarr_container_path=self.path,
@@ -1339,6 +1395,7 @@ def write_element(
             overwrite=overwrite,
             parsed_formats=parsed_formats,
             shapes_geometry_encoding=shapes_geometry_encoding,
+            element_raster_write_kwargs=element_raster_write_kwargs,
         )
         # After every write, metadata should be consolidated, otherwise this can lead to IO problems like when deleting.
         if self.has_consolidated_metadata():

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
 
@@ -283,12 +283,27 @@ def _write_raster(
     raster_format
         The format used to write the raster data.
     storage_options
-        Additional options for writing the raster data, like chunks and compression.
+        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 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
     label_metadata
         Label metadata which can only be defined when writing 'labels'.
     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
@@ -305,6 +320,24 @@ def _write_raster(
         for c in channels:
             metadata["metadata"]["omero"]["channels"].append({"label": c})  # type: ignore[union-attr, index, call-overload]
 
+    if isinstance(storage_options, dict):
+        storage_options = {
+            **{k.split("_")[1]: v for k, v in asdict(settings).items() if k in ("raster_chunks", "raster_shards")},
+            **storage_options,
+        }
+    elif isinstance(storage_options, list):
+        storage_options = [
+            {
+                **{k.split("_")[1]: v for k, v in asdict(settings).items() if k in ("raster_chunks", "raster_shards")},
+                **x,
+            }
+            for x in storage_options
+        ]
+    elif not storage_options:
+        storage_options = {
+            k.split("_")[1]: v for k, v in asdict(settings).items() if k in ("raster_chunks", "raster_shards")
+        }
+
     if isinstance(raster_data, DataArray):
         _write_raster_dataarray(
             raster_type,