diff --git a/.gitignore b/.gitignore index 7019b90..93565e5 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,7 @@ __pycache__/ # Environment environment.yml + +# Development +/dev/ +CLAUDE.md diff --git a/CHANGELOG.md b/CHANGELOG.md index ea076e6..d4d8ad4 100644 --- a/CHANGELOG.md +++ b/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]` diff --git a/docs/fileformat.md b/docs/fileformat.md new file mode 100644 index 0000000..0f38eaa --- /dev/null +++ b/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. +::: diff --git a/docs/index.md b/docs/index.md index 34eb5c9..9979153 100644 --- a/docs/index.md +++ b/docs/index.md @@ -14,6 +14,8 @@ notebooks/concatenation notebooks/alignment +fileformat.md + changelog.md contributing.md references.md diff --git a/pyproject.toml b/pyproject.toml index 68b3994..8f8ff38 100644 --- a/pyproject.toml +++ b/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" } diff --git a/src/treedata/_core/read.py b/src/treedata/_core/read.py index 738437e..a90ef2d 100755 --- a/src/treedata/_core/read.py +++ b/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 diff --git a/src/treedata/_core/write.py b/src/treedata/_core/write.py index 919000f..e224397 100755 --- a/src/treedata/_core/write.py +++ b/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) diff --git a/tests/test_readwrite.py b/tests/test_readwrite.py index f432298..3896103 100755 --- a/tests/test_readwrite.py +++ b/tests/test_readwrite.py @@ -93,6 +93,164 @@ def test_zarr_readwrite(tdata, tmp_path): assert tdata2.obs.loc["0", "tree"] == "1,2" +def test_zarr_dtypes(tdata, tmp_path): + """Columnar zarr storage round-trips all supported attribute types.""" + tdata.obst["1"].nodes["root"]["list"] = [1, 2, 3] + tdata.obst["1"].nodes["root"]["tuple"] = (1, 2, 3) + tdata.obst["1"].nodes["root"]["set"] = {1, 2, 3} + tdata.obst["1"].nodes["root"]["np_float"] = np.float64(1.0) + tdata.obst["1"].nodes["root"]["np_array"] = np.array([[1, 2], [3, 4]]) + tdata.obst["1"].nodes["root"]["pd_series"] = pd.Series(["a", "b", "c"]) + tdata.obst["1"].nodes["root"]["nested_list"] = [[1, 2], [3, 4], [5, 6]] + tdata.obst["1"].nodes["root"]["str_val"] = "hello" + tdata.obst["1"].nodes["root"]["int_val"] = 42 + tdata.obst["1"].nodes["root"]["float_val"] = 3.14 + tdata.obst["1"].nodes["root"]["bool_val"] = True + tdata.write_zarr(tmp_path / "test.zarr") + tdata2 = td.read_zarr(tmp_path / "test.zarr") + root = tdata2.obst["1"].nodes["root"] + assert root["list"] == [1, 2, 3] and isinstance(root["list"], list) + assert root["tuple"] == [1, 2, 3] and isinstance(root["tuple"], list) + assert root["set"] == [1, 2, 3] and isinstance(root["set"], list) + assert root["np_float"] == 1.0 and isinstance(root["np_float"], float) + assert root["np_array"] == [[1, 2], [3, 4]] and isinstance(root["np_array"], list) + assert root["pd_series"] == ["a", "b", "c"] and isinstance(root["pd_series"], list) + assert root["nested_list"] == [[1, 2], [3, 4], [5, 6]] + assert root["str_val"] == "hello" + assert root["int_val"] == 42 + assert root["float_val"] == pytest.approx(3.14) + assert root["bool_val"] is True + + +def test_zarr_edge_attrs(tmp_path): + """Edge attributes with complex types survive zarr round-trip.""" + tree = nx.DiGraph() + tree.add_edges_from([("root", "A"), ("root", "B"), ("A", "C")]) + tree["root"]["A"]["length"] = 1.5 + tree["root"]["A"]["tags"] = ["x", "y"] + tree["root"]["B"]["length"] = 2.0 + tree["root"]["B"]["tags"] = ["z"] + tree["A"]["C"]["length"] = 0.5 + tree["A"]["C"]["tags"] = [] + tdata = td.TreeData(np.eye(3), obst={"t": tree}, alignment="subset") + tdata.write_zarr(tmp_path / "test.zarr") + tdata2 = td.read_zarr(tmp_path / "test.zarr") + G = tdata2.obst["t"] + assert G["root"]["A"]["length"] == pytest.approx(1.5) + assert G["root"]["A"]["tags"] == ["x", "y"] + assert G["root"]["B"]["length"] == pytest.approx(2.0) + assert G["root"]["B"]["tags"] == ["z"] + assert G["A"]["C"]["tags"] == [] + + +def test_zarr_missing_node_attrs(tmp_path): + """Nodes without a particular attribute don't gain it as None after round-trip.""" + tree = nx.DiGraph() + tree.add_nodes_from(["root", "A", "B"]) + tree.add_edges_from([("root", "A"), ("root", "B")]) + tree.nodes["root"]["depth"] = 0 + tree.nodes["A"]["depth"] = 1 + tree.nodes["A"]["label"] = "leaf_a" + # B has no attributes at all + tdata = td.TreeData(np.eye(3), obst={"t": tree}, alignment="subset") + tdata.write_zarr(tmp_path / "test.zarr") + tdata2 = td.read_zarr(tmp_path / "test.zarr") + G = tdata2.obst["t"] + assert G.nodes["root"]["depth"] == 0 + assert "label" not in G.nodes["root"] + assert G.nodes["A"]["depth"] == 1 + assert G.nodes["A"]["label"] == "leaf_a" + assert dict(G.nodes["B"]) == {} + + +def test_zarr_zip_store(tdata, tmp_path): + """Writing to a ZipStore (the original bug report) works end-to-end.""" + import zarr + + zip_path = tmp_path / "test.zarr.zip" + store = zarr.storage.ZipStore(str(zip_path), mode="w") + tdata.write_zarr(store=store) + store.close() + tdata2 = td.read_zarr(str(zip_path)) + assert np.array_equal(tdata2.X, tdata.X) + check_graph_equality(tdata2.obst["1"], tdata.obst["1"]) + assert tdata2.obst["1"].nodes["root"]["depth"] == 0 + assert tdata2.obst["1"].nodes["root"]["characters"] == ["-1", "1"] + + +def test_zarr_large_tree(tmp_path): + """A tree with many nodes and multiple attributes per node round-trips correctly.""" + # balanced_tree(2, 10) has 2047 nodes — large enough to exercise the code path + # that failed with scalar string storage, without being slow + tree = nx.balanced_tree(2, 10, create_using=nx.DiGraph) + for node in tree.nodes(): + tree.nodes[node]["tags"] = [str(node), f"node_{node}"] + tree.nodes[node]["weight"] = float(node) / tree.number_of_nodes() + for u, v in tree.edges(): + tree[u][v]["dist"] = abs(u - v) + tree[u][v]["meta"] = [u, v] + tdata = td.TreeData(np.eye(3), obst={"big": tree}, alignment="subset") + tdata.write_zarr(tmp_path / "big.zarr") + tdata2 = td.read_zarr(tmp_path / "big.zarr") + G = tdata2.obst["big"] + assert G.number_of_nodes() == tree.number_of_nodes() + assert G.number_of_edges() == tree.number_of_edges() + assert isinstance(G.nodes["0"]["tags"], list) + assert G.nodes["0"]["tags"] == ["0", "node_0"] + assert G["0"]["1"]["meta"] == [0, 1] + + +def test_zarr_explicit_none(tmp_path): + """An attribute explicitly set to None is preserved and kept distinct from a missing one.""" + tree = nx.DiGraph() + tree.add_nodes_from(["root", "A", "B"]) + tree.add_edges_from([("root", "A"), ("root", "B")]) + tree.nodes["root"]["value"] = None # explicit None + tree.nodes["A"]["value"] = 1 + # B has no "value" attribute at all + tree["root"]["A"]["weight"] = None # explicit None on an edge + tdata = td.TreeData(np.eye(3), obst={"t": tree}, alignment="subset") + tdata.write_zarr(tmp_path / "test.zarr") + G = td.read_zarr(tmp_path / "test.zarr").obst["t"] + assert "value" in G.nodes["root"] and G.nodes["root"]["value"] is None + assert G.nodes["A"]["value"] == 1 + assert "value" not in G.nodes["B"] + assert "weight" in G["root"]["A"] and G["root"]["A"]["weight"] is None + + +def test_zarr_native_dtype_storage(tmp_path): + """Dense scalar numeric/bool columns are stored with a native (non-JSON) dtype.""" + import zarr + + tree = nx.DiGraph() + tree.add_edges_from([("root", "A"), ("root", "B")]) + for i, node in enumerate(tree.nodes()): + tree.nodes[node]["depth"] = i # dense int column -> native + tree.nodes[node]["ragged"] = [i, i] # ragged column -> JSON + tdata = td.TreeData(np.eye(3), obst={"t": tree}, alignment="subset") + tdata.write_zarr(tmp_path / "test.zarr") + g = zarr.open(str(tmp_path / "test.zarr"), mode="r") + assert g["obst/t/node_attrs/depth"].dtype.kind in "iu" # stored as native integers + assert g["obst/t/node_attrs/ragged"].dtype.kind not in "iufb" # stored as JSON strings + G = td.read_zarr(tmp_path / "test.zarr").obst["t"] + assert G.nodes["root"]["depth"] == 0 and isinstance(G.nodes["root"]["depth"], int) + assert G.nodes["A"]["ragged"] == [1, 1] + + +def test_zarr_stale_tree_removed(tdata, tmp_path): + """Rewriting into an existing group drops trees no longer present in the object.""" + import zarr + + store_path = str(tmp_path / "test.zarr") + tdata.write_zarr(store_path) + assert set(td.read_zarr(store_path).obst.keys()) == {"1", "2"} + # Rewrite into the same store with one obst tree removed + del tdata.obst["2"] + group = zarr.open(store_path, mode="a") + tdata.write_zarr(group) + assert set(td.read_zarr(store_path).obst.keys()) == {"1"} + + def test_read_anndata(X, tmp_path): adata = ad.AnnData(X) file_path = tmp_path / "test.h5ad"