diff --git a/src/adata/elements/spec.py b/src/adata/elements/spec.py index 541dd1a..1466717 100644 --- a/src/adata/elements/spec.py +++ b/src/adata/elements/spec.py @@ -10,7 +10,7 @@ from __future__ import annotations -from typing import Any, Optional, Tuple +from typing import Any, Dict, Optional, Tuple ANNDATA = "anndata" RAW = "raw" @@ -95,6 +95,25 @@ def encoding_type(obj: Any) -> Optional[str]: return encoding_of(obj)[0] +def set_attrs(obj: Any, values: Dict[str, Any]) -> None: + """Set several attributes in as few writes as the backend allows. + + On Zarr each assignment persists the whole metadata document through its + sync-over-async bridge, so setting attributes one at a time costs a round + trip each. HDF5 has no such batch API and falls back to assignment. + """ + if not values: + return + + put = getattr(obj.attrs, "put", None) + if callable(put): + put({**dict(obj.attrs), **values}) + return + + for key, value in values.items(): + obj.attrs[key] = value + + def set_encoding(obj: Any, enc_type: str, version: Optional[str] = None) -> None: """Stamp ``encoding-type``/``encoding-version`` on a group or dataset. @@ -103,5 +122,4 @@ def set_encoding(obj: Any, enc_type: str, version: Optional[str] = None) -> None """ if version is None: version = CURRENT_VERSION[enc_type] - obj.attrs["encoding-type"] = enc_type - obj.attrs["encoding-version"] = version + set_attrs(obj, {"encoding-type": enc_type, "encoding-version": version}) diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py index 96d351d..58e7bad 100644 --- a/src/adata/storage/__init__.py +++ b/src/adata/storage/__init__.py @@ -229,8 +229,27 @@ def _normalize_attr_value(value: Any, target_backend: str) -> Any: def copy_attrs(src_attrs: Any, dst_attrs: Any, *, target_backend: str) -> None: - for k, v in src_attrs.items(): - dst_attrs[k] = _normalize_attr_value(v, target_backend) + """Copy attributes across, normalising values for the target backend. + + Written in one go on Zarr. Each `attrs[k] = v` there persists the whole + metadata document through zarr's sync-over-async bridge, so setting them + one at a time means a separate round trip per attribute -- slow, and the + place a `split` over many groups was observed to wedge in CI. + """ + normalized = { + str(k): _normalize_attr_value(v, target_backend) + for k, v in src_attrs.items() + } + if not normalized: + return + + put = getattr(dst_attrs, "put", None) + if target_backend == "zarr" and callable(put): + put({**dict(dst_attrs), **normalized}) + return + + for k, v in normalized.items(): + dst_attrs[k] = v def dataset_create_kwargs(