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,