Fix zarr ZipStore write failure for large trees (v0.2.6)

.gitignore

@@ -30,3 +30,7 @@ __pycache__/
 
 # Environment
 environment.yml
+
+# Development
+/dev/
+CLAUDE.md

CHANGELOG.md

@@ -14,11 +14,16 @@ and this project adheres to [Semantic Versioning][].
 
 ### Changed
 
-- `dtype` parameter in `TreeData.__init__` is now deprecated; passing it raises a `FutureWarning` and has no effect, matching anndata's removal of this argument
-- `obsm`, `varm`, `layers`, `raw`, `shape`, `filename`, `filemode`, `asview`, and related parameters in `TreeData.__init__` are now keyword-only, matching the anndata 0.13.x API
+- Zarr writes now store `obst`/`vart` in a columnar layout (one array per node/edge attribute key) instead of a single JSON scalar. Dense numeric/boolean columns use native dtypes; ragged or complex columns fall back to per-element JSON. This fixes a `TypeError: string too large to store inside array` crash when writing large trees to zarr v3 (#86)
 
 ### Fixed
 
+- Fixed `UserWarning: Duplicate name: 'zarr.json'` spam when writing to a zarr `ZipStore` by staging the write in a `MemoryStore` and copying to the `ZipStore` once (#86)
+- `read_zarr` now auto-detects `.zip` paths and opens them as a `ZipStore` (#86)
+- Node/edge attributes explicitly set to `None` are now preserved through a zarr round-trip and kept distinct from absent attributes (#86)
+- Rewriting into an existing zarr group now removes trees that are no longer present, so the persisted tree set matches the `TreeData` being written (#86)
+- `dtype` parameter in `TreeData.__init__` is now deprecated; passing it raises a `FutureWarning` and has no effect, matching anndata's removal of this argument
+- `obsm`, `varm`, `layers`, `raw`, `shape`, `filename`, `filemode`, `asview`, and related parameters in `TreeData.__init__` are now keyword-only, matching the anndata 0.13.x API
 - Fixed compatibility with anndata 0.13.x: removed `dtype` from internal `_init_as_actual` call, which was removed from AnnData in 0.13.x
 - Fixed compatibility with anndata 0.13.x "Unify X and layers" change: `layers.copy()` and `dict(layers)` no longer include the internal `None` key (used to store X) in `_mutated_copy` and write routines
 - Fixed repr incorrectly showing `layers: None` when anndata 0.13.x stores X as `layers[None]`

docs/fileformat.md

@@ -0,0 +1,73 @@
+# On-disk format
+
+TreeData is written to `.h5td` (HDF5) and `.zarr` stores using the same
+hierarchical, self-describing layout as AnnData. A TreeData file **is** an
+AnnData file with a few extra top-level elements, so almost everything is
+covered by the [AnnData on-disk format][anndata-format] — read that first. This
+page documents only what TreeData adds.
+
+[anndata-format]: https://anndata.scverse.org/en/stable/fileformat-prose.html
+
+## Top-level
+
+The root group carries `encoding-type: "treedata"` and, in addition to the
+standard AnnData elements (`X`, `obs`, `var`, `obsm`, `uns`, …), stores:
+
+| Element | Type | Description |
+| --- | --- | --- |
+| `label` | scalar (string) or absent | Column in `obs`/`var` that holds each node's tree key |
+| `allow_overlap` | scalar (bool) | Whether trees may share nodes |
+| `alignment` | scalar (string) | One of `leaves`, `nodes`, `subset` |
+| `obst` | see below | Trees aligned to observations |
+| `vart` | see below | Trees aligned to variables |
+
+`obst` and `vart` are mappings of a key to a :class:`networkx.DiGraph`. They are
+the only elements stored differently from AnnData, and the layout depends on the
+backend.
+
+## Trees in HDF5 (`.h5td`)
+
+Each of `obst`/`vart` is stored as a **single JSON string** dataset mapping tree
+key → `{"nodes": {...}, "edges": {...}}`, with node and edge attributes inline:
+
+```json
+{"tree1": {"nodes": {"root": {"depth": 0}, "0": {}},
+           "edges": {"root": {"0": {"length": 1.5}}}}}
+```
+
+## Trees in zarr
+
+zarr cannot store a single arbitrarily large string, so each of `obst`/`vart` is
+a **group** with one subgroup per tree, stored columnar:
+
+```
+obst/
+  tree1/
+    nodes                 # 1-D array of node IDs
+    edges                 # (n_edges, 2) array of (source, target) pairs
+    node_attrs/           # one array per node-attribute key (optional)
+      depth
+      characters
+    edge_attrs/           # one array per edge-attribute key (optional)
+      length
+```
+
+The number of arrays scales with the number of **attribute keys**, not the
+number of nodes.
+
+Each attribute column has one entry per node (or edge), in the order of `nodes`
+(or `edges`), and is encoded one of two ways:
+
+- **Native** — a dense column whose values are all a single numeric or boolean
+  type is stored with the corresponding native array dtype (e.g. `int64`,
+  `float64`, `bool`).
+- **JSON** — any other column (strings, ragged/nested values, mixed types,
+  columns missing from some nodes) is stored as a string array of per-element
+  JSON. The empty string `""` marks an absent attribute, so a value explicitly
+  set to `null` is preserved and kept distinct from a missing one.
+
+:::{note}
+zarr files written before treedata 0.3.0 stored `obst`/`vart` as a single JSON
+string (as in HDF5). Such files are still read correctly; the format is detected
+at read time.
+:::

docs/index.md

@@ -14,6 +14,8 @@ notebooks/concatenation
 
 notebooks/alignment
 
+fileformat.md
+
 changelog.md
 contributing.md
 references.md

pyproject.toml

@@ -4,7 +4,7 @@ requires = [ "hatchling" ]
 
 [project]
 name = "treedata"
-version = "0.2.5"
+version = "0.2.6"
 description = "anndata with trees"
 readme = "README.md"
 license = { file = "LICENSE" }

src/treedata/_core/read.py

@@ -4,7 +4,8 @@
 import warnings
 from collections.abc import MutableMapping
 from importlib.metadata import version as get_version
-from typing import TYPE_CHECKING, Literal
+from os import PathLike
+from typing import Literal
 
 import anndata as ad
 import h5py
@@ -14,12 +15,12 @@
 
 from treedata._core.treedata import TreeData
 
-if TYPE_CHECKING:
-    from os import PathLike
-
 ANNDATA_VERSION = version.parse(get_version("anndata"))
 USE_EXPERIMENTAL = ANNDATA_VERSION < version.parse("0.11.0")
 
+# Sentinel distinguishing an absent attribute from an attribute explicitly set to ``None``.
+_MISSING = object()
+
 
 def _read_elem(elem):
     """Read an element from a store."""
@@ -43,10 +44,49 @@ def _dict_to_digraph(graph_dict: dict) -> nx.DiGraph:
 
 
 def _parse_axis_trees(data: str) -> dict:
-    """Parse AxisTrees from a string."""
+    """Parse AxisTrees from a JSON string (HDF5 format)."""
     return {k: _dict_to_digraph(v) for k, v in json.loads(data).items()}
 
 
+def _decode_attr_columns(cols: dict) -> dict[str, list]:
+    """Decode columnar attribute arrays into per-element Python lists.
+
+    Native numeric/bool arrays are returned as-is (every element present).
+    Object arrays hold per-element JSON strings, where the empty string ``""``
+    marks an absent attribute; it is decoded to the :data:`_MISSING` sentinel so
+    that an attribute explicitly set to ``None`` (JSON ``"null"``) is preserved
+    and kept distinct from a missing attribute.
+    """
+    decoded: dict[str, list] = {}
+    for k, arr in cols.items():
+        dtype = getattr(arr, "dtype", None)
+        if dtype is not None and dtype.kind in "iufb":
+            decoded[k] = list(arr.tolist())
+        else:
+            decoded[k] = [_MISSING if s == "" else json.loads(s) for s in arr]
+    return decoded
+
+
+def _read_axis_trees_zarr(g: zarr.Group) -> dict:
+    """Read AxisTrees from a zarr group written in the columnar format."""
+    trees = {}
+    for name in g.keys():
+        tg = g[name]
+        if not isinstance(tg, zarr.Group):
+            continue
+        G = nx.DiGraph()
+        nodes = list(_read_elem(tg["nodes"]))
+        node_attrs = _decode_attr_columns(_read_elem(tg["node_attrs"])) if "node_attrs" in tg else {}
+        for i, node in enumerate(nodes):
+            G.add_node(node, **{k: col[i] for k, col in node_attrs.items() if col[i] is not _MISSING})
+        edges = _read_elem(tg["edges"])
+        edge_attrs = _decode_attr_columns(_read_elem(tg["edge_attrs"])) if "edge_attrs" in tg else {}
+        for i, (src, dst) in enumerate(edges):
+            G.add_edge(src, dst, **{k: col[i] for k, col in edge_attrs.items() if col[i] is not _MISSING})
+        trees[name] = G
+    return trees
+
+
 def _parse_legacy(treedata_attrs: dict) -> dict:
     """Parse tree attributes from AnnData uns field."""
     if treedata_attrs is not None:
@@ -92,7 +132,11 @@ def _read_tdata(f, filename, backed) -> dict:
     # Read axis tree elements
     for k in ["obst", "vart"]:
         if k in f:
-            d[k] = _parse_axis_trees(_read_elem(f[k]))
+            elem = f[k]
+            if isinstance(elem, zarr.Group):
+                d[k] = _read_axis_trees_zarr(elem)
+            else:
+                d[k] = _parse_axis_trees(_read_elem(elem))
     # Read legacy treedata format
     if "raw.treedata" in f:
         d.update(_parse_legacy(json.loads(_read_elem(f["raw.treedata"]))))
@@ -148,6 +192,9 @@ def _open_zarr_group(store, *, mode: str = "r") -> tuple[zarr.Group, bool]:
     """Open a Zarr group and signal whether it should be closed."""
     if isinstance(store, zarr.Group):
         return store, False
+    # zarr v3 does not auto-detect zip files from string paths; open explicitly
+    if isinstance(store, (str, PathLike)) and str(store).endswith(".zip"):
+        store = zarr.storage.ZipStore(store, mode=mode)
     return zarr.open(store, mode=mode), True
 
 

src/treedata/_core/write.py

@@ -23,6 +23,9 @@
 USE_EXPERIMENTAL = ANNDATA_VERSION < version.parse("0.11.0")
 ZARR_V2 = version.parse(get_version("zarr")) < version.parse("3.0.0")
 
+# Sentinel distinguishing an absent attribute from an attribute explicitly set to ``None``.
+_MISSING = object()
+
 
 def _make_serializable(data: Any) -> Any:
     """Make a dictionary serializable."""
@@ -60,11 +63,70 @@ def _digraph_to_dict(G: nx.DiGraph) -> dict:
 
 
 def _serialize_axis_trees(trees: AxisTrees) -> str:
-    """Serialize AxisTrees."""
+    """Serialize AxisTrees to a JSON string (used for HDF5)."""
     d = {k: _digraph_to_dict(v) for k, v in trees.items()}
     return json.dumps(_make_serializable(d))
 
 
+def _encode_attr_column(values: list) -> np.ndarray:
+    """Encode one attribute column (values in node/edge order) as a 1-D array.
+
+    A column is stored with a native numpy dtype when it is *dense* (every
+    node/edge has the attribute) and every value is a scalar of a single
+    numeric or boolean Python type. This keeps the common case (e.g. ``depth``,
+    ``time``, branch ``length``) compact and typed.
+
+    Otherwise — ragged values (``list``/``set``/``np.ndarray``/nested), mixed
+    types, string values, sparse columns, or explicit ``None`` — the column
+    falls back to an object array of per-element JSON strings. The empty string
+    ``""`` is reserved as the "attribute absent" marker, so an attribute
+    explicitly set to ``None`` (stored as JSON ``"null"``) is preserved and kept
+    distinct from a missing attribute.
+
+    The number of attribute keys is expected to be small, so this stores only a
+    handful of arrays per tree — object count scales with the number of keys,
+    never with the number of nodes/edges.
+    """
+    dense = all(v is not _MISSING for v in values)
+    serialized = [v if v is _MISSING else _make_serializable(v) for v in values]
+    present_types = {type(v) for v in serialized if v is not _MISSING}
+    # Native path: dense column of a single numeric/bool Python type.
+    if dense and len(present_types) == 1 and next(iter(present_types)) in (int, float, bool):
+        return np.array(serialized)
+    return np.array(["" if v is _MISSING else json.dumps(v) for v in serialized], dtype=object)
+
+
+def _write_axis_trees_zarr(f: zarr.Group, key: str, trees: AxisTrees) -> None:
+    """Write AxisTrees to zarr in a columnar layout (one array per attribute key).
+
+    Each ``DiGraph`` is stored as a subgroup with separate arrays for node IDs,
+    edge pairs, and per-attribute columns (see :func:`_encode_attr_column`).
+    This avoids the previous single-JSON-scalar format, which zarr v3 could not
+    store above a size limit and which crashed on large trees.
+    """
+    # Remove any previously written trees so the persisted set matches ``trees``
+    # exactly (e.g. when rewriting into an existing group).
+    if key in f:
+        del f[key]
+    g = f.require_group(key)
+    for name, G in trees.items():
+        tg = g.require_group(str(name))
+        nodes = list(G.nodes())
+        edges = list(G.edges())
+        node_ids = np.array([str(n) for n in nodes], dtype=object)
+        edge_ids = np.array([(str(u), str(v)) for u, v in edges], dtype=object) if edges else np.zeros((0, 2), object)
+        _write_elem(tg, "nodes", node_ids, dataset_kwargs={})
+        _write_elem(tg, "edges", edge_ids, dataset_kwargs={})
+        node_attr_keys = {k for n in nodes for k in G.nodes[n]}
+        if node_attr_keys:
+            cols = {k: _encode_attr_column([G.nodes[n].get(k, _MISSING) for n in nodes]) for k in node_attr_keys}
+            _write_elem(tg, "node_attrs", cols, dataset_kwargs={})
+        edge_attr_keys = {k for u, v in edges for k in G[u][v]}
+        if edge_attr_keys:
+            cols = {k: _encode_attr_column([G[u][v].get(k, _MISSING) for u, v in edges]) for k in edge_attr_keys}
+            _write_elem(tg, "edge_attrs", cols, dataset_kwargs={})
+
+
 def _write_tdata(f, tdata, filename, chunks=None, **kwargs) -> None:
     """Write TreeData to file."""
     # Add encoding type and version
@@ -86,7 +148,10 @@ def _write_tdata(f, tdata, filename, chunks=None, **kwargs) -> None:
     _write_elem(f, "layers", {k: v for k, v in tdata.layers.items() if k is not None}, dataset_kwargs=kwargs)
     # Write axis tree elements
     for key in ["obst", "vart"]:
-        _write_elem(f, key, _serialize_axis_trees(getattr(tdata, key)), dataset_kwargs=kwargs)
+        if isinstance(f, zarr.Group):
+            _write_axis_trees_zarr(f, key, getattr(tdata, key))
+        else:
+            _write_elem(f, key, _serialize_axis_trees(getattr(tdata, key)), dataset_kwargs=kwargs)
     # Write raw
     if tdata.raw is not None:
         tdata.strings_to_categoricals(tdata.raw.var)
@@ -179,6 +244,52 @@ def _close_zarr_group(group: zarr.Group) -> None:
             close()
 
 
+def _is_zip_store(store: Any) -> bool:
+    """Return True if store is or resolves to a zarr ZipStore."""
+    if ZARR_V2:
+        return False
+    return isinstance(store, zarr.storage.ZipStore) or (
+        isinstance(store, (str, PathLike)) and str(store).endswith(".zip")
+    )
+
+
+def _write_zarr_via_memory(
+    zip_target: zarr.storage.ZipStore | PathLike | str,
+    tdata: TreeData,
+    chunks: tuple[int, ...] | None,
+    **kwargs: Any,
+) -> None:
+    """Write to a ZipStore via an intermediate MemoryStore.
+
+    zarr v3's ZipStore is append-only: every metadata update appends a new
+    zarr.json entry, producing duplicate-name warnings. Writing to MemoryStore
+    first (which supports overwriting) and then bulk-copying the final state to
+    the ZipStore ensures each key is written exactly once.
+    """
+    mem_store = zarr.storage.MemoryStore()
+    f = zarr.open_group(mem_store, mode="w", zarr_format=3)
+    _write_tdata(f, tdata, zip_target, chunks, **kwargs)
+
+    if isinstance(zip_target, zarr.storage.ZipStore):
+        zip_store = zip_target
+        should_close = False
+    else:
+        zip_store = zarr.storage.ZipStore(str(zip_target), mode="w")
+        should_close = True
+
+    # ZipStore has no __setitem__; writes go through the async set() method.
+    from zarr.core.sync import sync as _zarr_sync
+
+    async def _bulk_copy() -> None:
+        for key, value in mem_store._store_dict.items():
+            await zip_store.set(key, value)
+
+    _zarr_sync(_bulk_copy())
+
+    if should_close:
+        zip_store.close()
+
+
 def write_zarr(
     filename: MutableMapping | PathLike | zarr.Group, tdata: TreeData, chunks: tuple[int, ...] | None = None, **kwargs
 ) -> None:
@@ -193,11 +304,10 @@ def write_zarr(
     kwargs
         Additional keyword arguments passed to :func:`zarr.save`.
     """
-    # if not ZARR_V2 and "zarr_format" not in kwargs:
-    #    f = zarr.open(store=filename, mode="w", zarr_format=3)
-    # else:
-    print(kwargs)
-    f, should_close = _open_zarr_group(filename)
-    _write_tdata(f, tdata, filename, chunks, **kwargs)
-    if should_close:
-        _close_zarr_group(f)
+    if _is_zip_store(filename):
+        _write_zarr_via_memory(filename, tdata, chunks, **kwargs)
+    else:
+        f, should_close = _open_zarr_group(filename)
+        _write_tdata(f, tdata, filename, chunks, **kwargs)
+        if should_close:
+            _close_zarr_group(f)