Merge pull request #11 from cellgeni/fix/zarr-attr-write-batching

Batch Zarr attribute writes: fixes the split hang in CI

Author: Claptar

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})

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(