From 8740b504254fd8345207aa264254a32b90a4f0dd Mon Sep 17 00:00:00 2001 From: colganwi Date: Tue, 21 Apr 2026 18:15:05 -0400 Subject: [PATCH 1/3] Fix zarr ZipStore write failure for large trees Stores tree node/edge attributes as columnar awkward arrays (one ak.Array per attribute key) using anndata's existing write_elem/read_elem codec. This avoids zarr's scalar string size limit that caused TypeError on large trees, stores typed data natively (float64, var * string, etc.), and eliminates the duplicate zarr.json warnings from ZipStore's append-only behavior by routing writes through an intermediate MemoryStore. Adds awkward>=2 as a dependency. Co-Authored-By: Claude Sonnet 4.6 --- pyproject.toml | 3 +- src/treedata/_core/read.py | 56 +++++++++++++++++-- src/treedata/_core/write.py | 103 +++++++++++++++++++++++++++++++--- tests/test_readwrite.py | 107 ++++++++++++++++++++++++++++++++++++ 4 files changed, 253 insertions(+), 16 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 265ea68..00def54 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" } @@ -24,6 +24,7 @@ classifiers = [ ] dependencies = [ "anndata>=0.11", + "awkward>=2", "h5py", "networkx", "numpy", diff --git a/src/treedata/_core/read.py b/src/treedata/_core/read.py index 738437e..b244fc6 100755 --- a/src/treedata/_core/read.py +++ b/src/treedata/_core/read.py @@ -4,9 +4,11 @@ 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 awkward as ak import h5py import networkx as nx import zarr @@ -14,9 +16,6 @@ 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") @@ -43,10 +42,48 @@ 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_column(arr) -> list: + """Decode an attribute column to a Python list. + + Handles both awkward arrays (new format) and JSON-string numpy arrays (legacy format). + """ + if isinstance(arr, ak.Array): + return arr.to_list() + return [json.loads(v) for v in arr] + + +def _read_axis_trees_zarr(g: zarr.Group) -> dict: + """Read AxisTrees from a zarr group written in columnar format. + + None values are dropped when reconstructing node/edge attribute dicts, + matching networkx semantics where an absent attribute differs from None. + """ + 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: dict[str, list] = {} + if "node_attrs" in tg: + node_attrs = {k: _decode_attr_column(v) for k, v in _read_elem(tg["node_attrs"]).items()} + for i, node in enumerate(nodes): + G.add_node(node, **{k: v[i] for k, v in node_attrs.items() if v[i] is not None}) + edges = _read_elem(tg["edges"]) + edge_attrs: dict[str, list] = {} + if "edge_attrs" in tg: + edge_attrs = {k: _decode_attr_column(v) for k, v in _read_elem(tg["edge_attrs"]).items()} + for i, (src, dst) in enumerate(edges): + G.add_edge(src, dst, **{k: v[i] for k, v in edge_attrs.items() if v[i] is not None}) + 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 +129,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 +189,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 9df1164..d8874fd 100755 --- a/src/treedata/_core/write.py +++ b/src/treedata/_core/write.py @@ -9,6 +9,7 @@ from typing import Any, Literal import anndata as ad +import awkward as ak import h5py import networkx as nx import numpy as np @@ -60,11 +61,45 @@ 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 _write_axis_trees_zarr(f: zarr.Group, key: str, trees: AxisTrees) -> None: + """Write AxisTrees to zarr using columnar awkward arrays (one array per attribute key). + + Stores each DiGraph as a subgroup with separate arrays for node IDs, edge + pairs, and per-attribute columns. Each attribute column is an awkward array + built with ak.from_iter, which preserves native types (float64, int64, + var * string, etc.) without JSON encoding. + """ + g = f.require_group(key) + for name, G in trees.items(): + tg = g.require_group(str(name)) + nodes = [str(n) for n in G.nodes()] + edges = [(str(u), str(v)) for u, v in G.edges()] + _write_elem(tg, "nodes", np.array(nodes, dtype=object), dataset_kwargs={}) + edge_arr = np.array(edges, dtype=object) if edges else np.zeros((0, 2), dtype=object) + _write_elem(tg, "edges", edge_arr, dataset_kwargs={}) + node_attr_keys = {k for n in G.nodes() for k in G.nodes[n]} + if node_attr_keys: + _write_elem( + tg, + "node_attrs", + {k: ak.from_iter([_make_serializable(G.nodes[n].get(k)) for n in G.nodes()]) for k in node_attr_keys}, + dataset_kwargs={}, + ) + edge_attr_keys = {k for u, v in G.edges() for k in G[u][v]} + if edge_attr_keys: + _write_elem( + tg, + "edge_attrs", + {k: ak.from_iter([_make_serializable(G[u][v].get(k)) for u, v in G.edges()]) for k in edge_attr_keys}, + dataset_kwargs={}, + ) + + def _write_tdata(f, tdata, filename, chunks=None, **kwargs) -> None: """Write TreeData to file.""" # Add encoding type and version @@ -84,7 +119,10 @@ def _write_tdata(f, tdata, filename, chunks=None, **kwargs) -> None: _write_elem(f, key, dict(getattr(tdata, key)), 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) @@ -177,6 +215,15 @@ 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( filename: MutableMapping | PathLike | zarr.Group, tdata: TreeData, chunks: tuple[int, ...] | None = None, **kwargs ) -> None: @@ -191,11 +238,49 @@ 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 _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) + + +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 attribute update rewrites zarr.json, + 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 + + # Copy every key from MemoryStore to ZipStore exactly once. + # ZipStore has no __setitem__; writes go through the async set() method. + # zarr.core.sync.sync is the same utility zarr uses for its own sync API. + 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: - _close_zarr_group(f) + zip_store.close() diff --git a/tests/test_readwrite.py b/tests/test_readwrite.py index f432298..54276ec 100755 --- a/tests/test_readwrite.py +++ b/tests/test_readwrite.py @@ -93,6 +93,113 @@ 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_read_anndata(X, tmp_path): adata = ad.AnnData(X) file_path = tmp_path / "test.h5ad" From 0566fc52e60dcb8f9b497fb0d6eee1650479e626 Mon Sep 17 00:00:00 2001 From: colganwi Date: Tue, 21 Apr 2026 18:20:26 -0400 Subject: [PATCH 2/3] Move _write_zarr_via_memory before write_zarr Co-Authored-By: Claude Sonnet 4.6 --- src/treedata/_core/write.py | 56 ++++++++++++++++++------------------- 1 file changed, 27 insertions(+), 29 deletions(-) diff --git a/src/treedata/_core/write.py b/src/treedata/_core/write.py index d8874fd..c881ba5 100755 --- a/src/treedata/_core/write.py +++ b/src/treedata/_core/write.py @@ -224,29 +224,6 @@ def _is_zip_store(store: Any) -> bool: ) -def write_zarr( - filename: MutableMapping | PathLike | zarr.Group, tdata: TreeData, chunks: tuple[int, ...] | None = None, **kwargs -) -> None: - """Write `.zarr`-formatted zarr file. - - Parameters - ---------- - filename - Filename of data file. Defaults to backing file. - tdata - TreeData object to write. - kwargs - Additional keyword arguments passed to :func:`zarr.save`. - """ - 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) - - def _write_zarr_via_memory( zip_target: zarr.storage.ZipStore | PathLike | str, tdata: TreeData, @@ -255,10 +232,10 @@ def _write_zarr_via_memory( ) -> None: """Write to a ZipStore via an intermediate MemoryStore. - zarr v3's ZipStore is append-only: every attribute update rewrites zarr.json, - 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. + 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) @@ -271,9 +248,7 @@ def _write_zarr_via_memory( zip_store = zarr.storage.ZipStore(str(zip_target), mode="w") should_close = True - # Copy every key from MemoryStore to ZipStore exactly once. # ZipStore has no __setitem__; writes go through the async set() method. - # zarr.core.sync.sync is the same utility zarr uses for its own sync API. from zarr.core.sync import sync as _zarr_sync async def _bulk_copy() -> None: @@ -284,3 +259,26 @@ async def _bulk_copy() -> None: if should_close: zip_store.close() + + +def write_zarr( + filename: MutableMapping | PathLike | zarr.Group, tdata: TreeData, chunks: tuple[int, ...] | None = None, **kwargs +) -> None: + """Write `.zarr`-formatted zarr file. + + Parameters + ---------- + filename + Filename of data file. Defaults to backing file. + tdata + TreeData object to write. + kwargs + Additional keyword arguments passed to :func:`zarr.save`. + """ + 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) From 8f7a195983a113fcefa8f5b4f998f07706f8c16c Mon Sep 17 00:00:00 2001 From: colganwi Date: Tue, 7 Jul 2026 17:53:22 -0400 Subject: [PATCH 3/3] zarr format without awkward arrays --- .gitignore | 4 ++ CHANGELOG.md | 7 ++++ docs/fileformat.md | 73 +++++++++++++++++++++++++++++++++++ docs/index.md | 2 + pyproject.toml | 1 - src/treedata/_core/read.py | 43 +++++++++++---------- src/treedata/_core/write.py | 77 +++++++++++++++++++++++++------------ tests/test_readwrite.py | 51 ++++++++++++++++++++++++ 8 files changed, 212 insertions(+), 46 deletions(-) create mode 100644 docs/fileformat.md 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 399055c..4a69643 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,8 +14,15 @@ and this project adheres to [Semantic Versioning][]. ### Changed +- 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) + ## [0.2.4] - 2025-11-05 ### Added 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 00def54..bd4e0e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,6 @@ classifiers = [ ] dependencies = [ "anndata>=0.11", - "awkward>=2", "h5py", "networkx", "numpy", diff --git a/src/treedata/_core/read.py b/src/treedata/_core/read.py index b244fc6..a90ef2d 100755 --- a/src/treedata/_core/read.py +++ b/src/treedata/_core/read.py @@ -8,7 +8,6 @@ from typing import Literal import anndata as ad -import awkward as ak import h5py import networkx as nx import zarr @@ -19,6 +18,9 @@ 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.""" @@ -46,22 +48,27 @@ def _parse_axis_trees(data: str) -> dict: return {k: _dict_to_digraph(v) for k, v in json.loads(data).items()} -def _decode_attr_column(arr) -> list: - """Decode an attribute column to a Python list. +def _decode_attr_columns(cols: dict) -> dict[str, list]: + """Decode columnar attribute arrays into per-element Python lists. - Handles both awkward arrays (new format) and JSON-string numpy arrays (legacy format). + 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. """ - if isinstance(arr, ak.Array): - return arr.to_list() - return [json.loads(v) for v in arr] + 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 columnar format. - - None values are dropped when reconstructing node/edge attribute dicts, - matching networkx semantics where an absent attribute differs from None. - """ + """Read AxisTrees from a zarr group written in the columnar format.""" trees = {} for name in g.keys(): tg = g[name] @@ -69,17 +76,13 @@ def _read_axis_trees_zarr(g: zarr.Group) -> dict: continue G = nx.DiGraph() nodes = list(_read_elem(tg["nodes"])) - node_attrs: dict[str, list] = {} - if "node_attrs" in tg: - node_attrs = {k: _decode_attr_column(v) for k, v in _read_elem(tg["node_attrs"]).items()} + 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: v[i] for k, v in node_attrs.items() if v[i] is not None}) + 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: dict[str, list] = {} - if "edge_attrs" in tg: - edge_attrs = {k: _decode_attr_column(v) for k, v in _read_elem(tg["edge_attrs"]).items()} + 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: v[i] for k, v in edge_attrs.items() if v[i] is not None}) + 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 diff --git a/src/treedata/_core/write.py b/src/treedata/_core/write.py index c881ba5..fa8a76b 100755 --- a/src/treedata/_core/write.py +++ b/src/treedata/_core/write.py @@ -9,7 +9,6 @@ from typing import Any, Literal import anndata as ad -import awkward as ak import h5py import networkx as nx import numpy as np @@ -24,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.""" @@ -66,38 +68,63 @@ def _serialize_axis_trees(trees: AxisTrees) -> str: 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 using columnar awkward arrays (one array per attribute key). + """Write AxisTrees to zarr in a columnar layout (one array per attribute key). - Stores each DiGraph as a subgroup with separate arrays for node IDs, edge - pairs, and per-attribute columns. Each attribute column is an awkward array - built with ak.from_iter, which preserves native types (float64, int64, - var * string, etc.) without JSON encoding. + 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 = [str(n) for n in G.nodes()] - edges = [(str(u), str(v)) for u, v in G.edges()] - _write_elem(tg, "nodes", np.array(nodes, dtype=object), dataset_kwargs={}) - edge_arr = np.array(edges, dtype=object) if edges else np.zeros((0, 2), dtype=object) - _write_elem(tg, "edges", edge_arr, dataset_kwargs={}) - node_attr_keys = {k for n in G.nodes() for k in G.nodes[n]} + 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: - _write_elem( - tg, - "node_attrs", - {k: ak.from_iter([_make_serializable(G.nodes[n].get(k)) for n in G.nodes()]) for k in node_attr_keys}, - dataset_kwargs={}, - ) - edge_attr_keys = {k for u, v in G.edges() for k in G[u][v]} + 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: - _write_elem( - tg, - "edge_attrs", - {k: ak.from_iter([_make_serializable(G[u][v].get(k)) for u, v in G.edges()]) for k in edge_attr_keys}, - dataset_kwargs={}, - ) + 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: diff --git a/tests/test_readwrite.py b/tests/test_readwrite.py index 54276ec..3896103 100755 --- a/tests/test_readwrite.py +++ b/tests/test_readwrite.py @@ -200,6 +200,57 @@ def test_zarr_large_tree(tmp_path): 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"