spatialdata-js-util 0.1.0__py3-none-any.whl

spatialdata_js_util/store.py

@@ -0,0 +1,305 @@
+"""Filesystem-level helpers for SpatialData Zarr stores.
+
+These operate on the store layout directly rather than through `spatialdata`, so
+that copy-then-rewrite workflows can touch one element without materialising the
+whole object — and so the reader-facing parts of this package stay importable
+without the heavy write dependencies.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Sequence
+
+import pandas as pd
+import pyarrow.dataset as ds
+
+
+def read_json(path: str | Path) -> dict[str, Any]:
+    return json.loads(Path(path).read_text())
+
+
+def json_bytes(value: dict[str, Any]) -> bytes:
+    return json.dumps(value, indent=2, sort_keys=True).encode("utf-8")
+
+
+def write_json(path: str | Path, value: dict[str, Any]) -> None:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_bytes(json_bytes(value))
+
+
+def drop_consolidated_metadata(group_path: str | Path) -> bool:
+    """Remove a group's own inline consolidated metadata, if it has any.
+
+    A nested block is a cache of the metadata below it. Rewriting anything under
+    the group leaves that cache describing the old encoding, and readers trust
+    the cache over the files — so an array can be rewritten correctly and still
+    fail to load. The store root's listing is refreshed separately and remains
+    the authoritative one.
+
+    Returns whether a block was actually removed.
+    """
+    meta_path = Path(group_path) / "zarr.json"
+    if not meta_path.is_file():
+        return False
+    doc = read_json(meta_path)
+    if "consolidated_metadata" not in doc:
+        return False
+    doc.pop("consolidated_metadata")
+    write_json(meta_path, doc)
+    return True
+
+
+def _implicit_group(zarr_format: int) -> dict[str, Any]:
+    return {"attributes": {}, "node_type": "group", "zarr_format": zarr_format}
+
+
+def _collect_consolidated_metadata(store_path: Path, *, zarr_format: int = 3) -> dict[str, Any]:
+    """Map every node under *store_path* to its metadata document.
+
+    Entries are filled in for intermediate groups that have no `zarr.json` of
+    their own. Some writers create element directories without writing metadata
+    for the collection that contains them — `points/transcripts/zarr.json` but no
+    `points/zarr.json`. Listing the child without its parent leaves an orphan
+    that zarr cannot attach to the hierarchy, and it rejects the *whole* store:
+    `zarr.open_group` then raises `GroupNotFoundError` at the root, so the store
+    stops opening entirely rather than merely missing one element.
+    """
+    metadata: dict[str, Any] = {}
+    for meta_path in store_path.rglob("zarr.json"):
+        if meta_path == store_path / "zarr.json":
+            continue
+        rel = meta_path.parent.relative_to(store_path).as_posix()
+        metadata[rel] = read_json(meta_path)
+
+    _fill_implicit_parents(metadata, zarr_format)
+    return metadata
+
+
+def _fill_implicit_parents(metadata: dict[str, Any], zarr_format: int) -> None:
+    """Add a group entry for every ancestor path that has none, in place.
+
+    Listing `points/transcripts` without `points` leaves an orphan that zarr
+    cannot attach to the hierarchy, and it rejects the *whole* store rather than
+    the one element — so every writer of this map has to close the gap.
+    """
+    for rel in list(metadata):
+        parts = rel.split("/")
+        for depth in range(1, len(parts)):
+            parent = "/".join(parts[:depth])
+            if parent not in metadata:
+                metadata[parent] = _implicit_group(zarr_format)
+
+
+def consolidated_metadata_orphans(metadata: dict[str, Any]) -> list[str]:
+    """Entries whose parent path is absent — these make a store unreadable."""
+    return sorted(
+        key for key in metadata if "/" in key and key.rsplit("/", 1)[0] not in metadata
+    )
+
+
+def refresh_consolidated_metadata(store_path: str | Path) -> None:
+    """Rewrite the store root's inline consolidated metadata from disk.
+
+    Required after any structural rewrite: readers trust the consolidated copy,
+    so a stale one hides newly written arrays or reports old codecs.
+    """
+    root = Path(store_path)
+    root_path = root / "zarr.json"
+    root_meta = read_json(root_path)
+    metadata = _collect_consolidated_metadata(
+        root, zarr_format=int(root_meta.get("zarr_format", 3))
+    )
+
+    orphans = consolidated_metadata_orphans(metadata)
+    if orphans:
+        # Never write metadata that would make the store fail to open.
+        raise ValueError(
+            f"Refusing to write consolidated metadata with orphaned entries for {root}: "
+            f"{', '.join(orphans)}"
+        )
+
+    root_meta["consolidated_metadata"] = {
+        "kind": "inline",
+        "must_understand": False,
+        "metadata": metadata,
+    }
+    write_json(root_path, root_meta)
+
+
+def list_element_keys(zarr_path: str | Path, collection: str) -> list[str]:
+    """List element names under a top-level store collection (images, tables, …)."""
+    root = Path(zarr_path) / collection
+    if not root.is_dir():
+        return []
+    return sorted(
+        child.name
+        for child in root.iterdir()
+        if child.is_dir() and ((child / "zarr.json").is_file() or (child / ".zgroup").is_file())
+    )
+
+
+def zarr_format_of(path: str | Path) -> int:
+    """Return the zarr format (2 or 3) a group on disk was written in.
+
+    AnnData groups inside a store must be rewritten in the format they were read
+    in; mixing formats within one store breaks readers.
+    """
+    node = Path(path)
+    if (node / "zarr.json").is_file():
+        return int(read_json(node / "zarr.json").get("zarr_format", 3))
+    if (node / ".zgroup").is_file():
+        return int(read_json(node / ".zgroup").get("zarr_format", 2))
+    raise FileNotFoundError(f"No zarr group metadata found at {node}")
+
+
+def _points_root(zarr_path: Path) -> Path:
+    return zarr_path / "points"
+
+
+def validate_points_key(points_key: str) -> str:
+    key = points_key.strip()
+    if not key:
+        raise ValueError("Points element name cannot be empty.")
+    if Path(key).name != key or key in {".", ".."}:
+        raise ValueError(
+            "Points element name must be a single name under points/, not a path."
+        )
+    return key
+
+
+def list_points_keys(zarr_path: str | Path) -> list[str]:
+    root = _points_root(Path(zarr_path))
+    if not root.is_dir():
+        return []
+    return sorted(
+        child.name
+        for child in root.iterdir()
+        if child.is_dir() and (child / "zarr.json").is_file()
+    )
+
+
+def read_points_element_attrs(zarr_path: str | Path, points_key: str) -> dict[str, Any]:
+    element_json = _points_root(Path(zarr_path)) / points_key / "zarr.json"
+    if not element_json.is_file():
+        raise FileNotFoundError(f"Points element not found: points/{points_key}")
+    attrs = read_json(element_json).get("attributes", {})
+    spatialdata_attrs = attrs.get("spatialdata_attrs", {})
+    if not isinstance(spatialdata_attrs, dict):
+        spatialdata_attrs = {}
+    return {
+        "axes": attrs.get("axes", []),
+        "feature_key": spatialdata_attrs.get("feature_key"),
+        "instance_key": spatialdata_attrs.get("instance_key"),
+        "version": spatialdata_attrs.get("version"),
+    }
+
+
+def points_parquet_path(zarr_path: str | Path, points_key: str) -> Path:
+    return _points_root(Path(zarr_path)) / validate_points_key(points_key) / "points.parquet"
+
+
+def experimental_points_output_path(zarr_path: str | Path, points_key: str) -> Path:
+    return Path(zarr_path) / "points.experimental" / validate_points_key(points_key) / "points.parquet"
+
+
+def copy_points_element_metadata(
+    zarr_path: str | Path,
+    *,
+    source_key: str,
+    dest_key: str,
+) -> None:
+    source = _points_root(Path(zarr_path)) / validate_points_key(source_key) / "zarr.json"
+    dest = _points_root(Path(zarr_path)) / validate_points_key(dest_key) / "zarr.json"
+    if not source.is_file():
+        raise FileNotFoundError(f"Points element metadata not found: {source}")
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(source, dest)
+
+
+def _points_element_consolidated_entry(
+    zarr_path: Path,
+    points_key: str,
+    *,
+    metadata: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    if metadata is not None:
+        entry = metadata.get(f"points/{points_key}")
+        if isinstance(entry, dict):
+            return deepcopy(entry)
+    element_json = _points_root(zarr_path) / points_key / "zarr.json"
+    element_doc = read_json(element_json)
+    return {
+        "attributes": element_doc.get("attributes", {}),
+        "node_type": element_doc.get("node_type", "group"),
+        "zarr_format": element_doc.get("zarr_format", 3),
+    }
+
+
+def read_store_consolidated_metadata(zarr_path: str | Path) -> dict[str, Any]:
+    root_json = Path(zarr_path) / "zarr.json"
+    if not root_json.is_file():
+        raise FileNotFoundError(f"Missing store metadata: {root_json}")
+    doc = read_json(root_json)
+    consolidated = doc.get("consolidated_metadata")
+    if not isinstance(consolidated, dict):
+        raise ValueError(f"Store has no consolidated metadata: {root_json}")
+    metadata = consolidated.get("metadata")
+    if not isinstance(metadata, dict):
+        raise ValueError(f"Store consolidated metadata has no metadata map: {root_json}")
+    return metadata
+
+
+def register_points_elements_in_consolidated_metadata(
+    zarr_path: str | Path,
+    element_keys: Sequence[str],
+    *,
+    template_key: str,
+) -> None:
+    """Register sibling points elements in the store root consolidated metadata."""
+    store_path = Path(zarr_path)
+    root_json = store_path / "zarr.json"
+    doc = read_json(root_json)
+    consolidated = doc.get("consolidated_metadata")
+    if not isinstance(consolidated, dict):
+        consolidated = {"kind": "inline", "metadata": {}}
+        doc["consolidated_metadata"] = consolidated
+    metadata = consolidated.get("metadata")
+    if not isinstance(metadata, dict):
+        metadata = {}
+        consolidated["metadata"] = metadata
+
+    template_entry = _points_element_consolidated_entry(
+        store_path,
+        template_key,
+        metadata=metadata,
+    )
+    for key in element_keys:
+        metadata[f"points/{key}"] = deepcopy(template_entry)
+
+    # Adding `points/<key>` to a listing with no `points` entry orphans it, and
+    # zarr then refuses the whole store — the same failure `refresh_consolidated_
+    # metadata` guards against, reachable here because this writes the map
+    # directly. A source store missing the parent is exactly the case this module
+    # exists to handle, so it cannot be assumed away.
+    _fill_implicit_parents(metadata, int(doc.get("zarr_format", 3)))
+    orphans = consolidated_metadata_orphans(metadata)
+    if orphans:
+        raise ValueError(
+            f"Refusing to write consolidated metadata with orphaned entries for {store_path}: "
+            f"{', '.join(orphans)}"
+        )
+    write_json(root_json, doc)
+
+
+def read_points_dataframe(parquet_path: str | Path) -> pd.DataFrame:
+    path = Path(parquet_path)
+    if not path.exists():
+        raise FileNotFoundError(f"Points Parquet not found: {path}")
+    # `ds.dataset` handles a single file and a partitioned directory alike.
+    table = ds.dataset(path, format="parquet").to_table()
+    return table.to_pandas()

spatialdata_js_util/tables.py

@@ -0,0 +1,310 @@
+"""Convert SpatialData table matrices to CSC.
+
+Browser clients read expression data one *variable* at a time — colour cells by
+one gene, then another. In CSR (AnnData's default on disk) a single gene's values
+are scattered across every row's slice, so one gene costs a read of essentially
+the whole matrix. In CSC that gene is one contiguous run of `indptr[j]:indptr[j+1]`,
+so it is a single ranged read.
+
+The conversion is lossless and reversible: only the sparse layout changes, not
+the values, and `scipy`/`anndata` read either layout transparently. It is
+therefore safe to apply to a store that other tools also read.
+
+Only the matrices themselves are rewritten. An earlier version re-serialised the
+whole AnnData, which quietly re-encoded `obs` and `var` with the installed
+AnnData's current conventions — turning the index into a `nullable-string-array`
+group and sharding the arrays — so readers that had been fine with the store
+suddenly could not find the variable names. Everything written here is therefore
+scoped to the matrix being converted, and pinned to the older encodings.
+"""
+
+from __future__ import annotations
+
+import shutil
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Iterator, Sequence
+
+from .errors import WriterCommandError
+from .provenance import package_version
+from .store import (
+    drop_consolidated_metadata,
+    list_element_keys,
+    refresh_consolidated_metadata,
+    write_json,
+    zarr_format_of,
+)
+
+if TYPE_CHECKING:
+    from anndata import AnnData
+
+#: Matrix keys converted by default: `X` plus every entry of `layers`.
+DEFAULT_MATRIX_KEYS = ("X", "layers")
+
+
+@dataclass(frozen=True)
+class CscConversion:
+    store_path: Path
+    manifest_path: Path | None
+    manifest: dict[str, Any]
+
+
+def _sparse_format(matrix: Any) -> str | None:
+    """Return the scipy sparse format string ('csr', 'csc', …) or None if dense."""
+    return getattr(matrix, "format", None)
+
+
+def to_csc(matrix: Any, *, densify: bool = False) -> Any:
+    """Return *matrix* in CSC layout.
+
+    Already-CSC input is returned unchanged. Dense input is only converted when
+    *densify* is set, because making a dense matrix sparse can multiply its size
+    rather than shrink it.
+    """
+    import scipy.sparse as sp
+
+    fmt = _sparse_format(matrix)
+    if fmt == "csc":
+        return matrix
+    if fmt is not None:
+        return matrix.tocsc()
+    if densify:
+        return sp.csc_matrix(matrix)
+    return matrix
+
+
+def anndata_to_csc(
+    adata: "AnnData",
+    *,
+    layers: bool = True,
+    densify: bool = False,
+) -> dict[str, str]:
+    """Convert `adata.X` (and optionally `adata.layers`) to CSC in place.
+
+    Returns a per-matrix report of ``{key: "csr->csc" | "already-csc" | …}``.
+    """
+    report: dict[str, str] = {}
+
+    before = _sparse_format(adata.X)
+    converted = to_csc(adata.X, densify=densify)
+    if converted is not adata.X:
+        adata.X = converted
+    report["X"] = _describe(before, _sparse_format(adata.X))
+
+    if layers:
+        for key in list(adata.layers.keys()):
+            # anndata >=0.12 exposes `X` as a `None`-keyed layer alias; converting
+            # it again would double-report (and rewrite) the matrix handled above.
+            if key is None:
+                continue
+            layer_before = _sparse_format(adata.layers[key])
+            layer_converted = to_csc(adata.layers[key], densify=densify)
+            if layer_converted is not adata.layers[key]:
+                adata.layers[key] = layer_converted
+            report[f"layers/{key}"] = _describe(
+                layer_before, _sparse_format(adata.layers[key])
+            )
+    return report
+
+
+def _describe(before: str | None, after: str | None) -> str:
+    if before is None and after is None:
+        return "dense (unchanged)"
+    if before == after:
+        return f"already-{after}"
+    return f"{before or 'dense'}->{after}"
+
+
+#: AnnData settings pinned while writing, so a matrix we rewrite is encoded the
+#: way the rest of the store already is. Names absent from the installed AnnData
+#: are skipped.
+#:
+#: `auto_shard_zarr_v3` wraps arrays in `sharding_indexed`, and
+#: `allow_write_nullable_strings` stores a string column as a
+#: `nullable-string-array` *group* of `values`/`mask` instead of a plain string
+#: array. Both are newer AnnData defaults that older readers — including the
+#: JS runtime this package targets — do not understand.
+_PINNED_WRITE_SETTINGS: dict[str, Any] = {
+    "auto_shard_zarr_v3": False,
+    "allow_write_nullable_strings": False,
+}
+
+
+@contextmanager
+def _anndata_write_settings(*, zarr_format: int) -> Iterator[None]:
+    import anndata as ad
+
+    settings = ad.settings
+    pinned = {**_PINNED_WRITE_SETTINGS, "zarr_write_format": zarr_format}
+    previous = {
+        name: getattr(settings, name) for name in pinned if hasattr(settings, name)
+    }
+    try:
+        for name, value in pinned.items():
+            if name in previous:
+                setattr(settings, name, value)
+        yield
+    finally:
+        for name, value in previous.items():
+            setattr(settings, name, value)
+
+
+def _convert_group_matrix(group: Any, name: str, *, densify: bool) -> str:
+    """Convert one matrix child of a zarr group in place, reporting what changed."""
+    from anndata.io import read_elem, write_elem
+
+    matrix = read_elem(group[name])
+    before = _sparse_format(matrix)
+    converted = to_csc(matrix, densify=densify)
+    if converted is not matrix:
+        write_elem(group, name, converted)
+    return _describe(before, _sparse_format(converted))
+
+
+def _convert_table_matrices(
+    table_path: Path, *, layers: bool, densify: bool
+) -> dict[str, str]:
+    """Rewrite a table's matrices in place, touching nothing else.
+
+    Only the matrix children are rewritten. Re-serialising the whole AnnData
+    would re-encode `obs`, `var`, `uns` and friends with whatever conventions the
+    installed AnnData currently prefers, silently changing parts of the store the
+    conversion has no business touching.
+    """
+    import zarr
+
+    # Consolidated metadata makes the group read-only; the store root's copy is
+    # refreshed once at the end of the conversion instead.
+    root = zarr.open_group(str(table_path), mode="a", use_consolidated=False)
+    report: dict[str, str] = {}
+
+    with _anndata_write_settings(zarr_format=zarr_format_of(table_path)):
+        # `X` is optional in AnnData — a table may carry only `layers`, or only
+        # annotations. Converting what is there beats refusing the whole table.
+        if "X" in root:
+            report["X"] = _convert_group_matrix(root, "X", densify=densify)
+        else:
+            report["X"] = "absent"
+        if layers and "layers" in root:
+            layers_group = root["layers"]
+            for name in sorted(layers_group.keys()):
+                report[f"layers/{name}"] = _convert_group_matrix(
+                    layers_group, name, densify=densify
+                )
+
+    # The table group may carry its own consolidated listing, which now describes
+    # the pre-conversion encoding. Readers trust it over the files, so a matrix
+    # rewritten correctly would still fail to load.
+    drop_consolidated_metadata(table_path)
+    return report
+
+
+def _table_dimensions(table_path: Path) -> tuple[int, int]:
+    """Report `(n_obs, n_vars)`, or `(0, 0)` for a table with no `X`.
+
+    Only used for the manifest, so an absent matrix is worth reporting as zero
+    rather than failing a conversion that has nothing to do with dimensions.
+    """
+    import zarr
+
+    root = zarr.open_group(str(table_path), mode="r", use_consolidated=False)
+    if "X" not in root:
+        return 0, 0
+    shape = root["X"].attrs.get("shape")
+    if shape is None:  # dense X is a plain array
+        shape = root["X"].shape
+    return int(shape[0]), int(shape[1])
+
+
+def convert_store_tables_to_csc(
+    source: str | Path,
+    dest: str | Path | None = None,
+    *,
+    tables: Sequence[str] | None = None,
+    layers: bool = True,
+    densify: bool = False,
+    overwrite: bool = False,
+    manifest: bool = True,
+) -> CscConversion:
+    """Rewrite a store's table matrices in CSC layout.
+
+    With *dest*, the store is copied first and the source left untouched; without
+    it, tables are rewritten in place. Pass *tables* to restrict which table
+    elements are converted (default: all of them).
+    """
+    source_path = Path(source)
+    if not (source_path / "zarr.json").is_file():
+        raise WriterCommandError(f"Not a zarr store (no zarr.json): {source_path}")
+
+    if dest is None:
+        store_path = source_path
+        in_place = True
+    else:
+        store_path = Path(dest)
+        in_place = False
+        # The overwrite branch below deletes the destination before copying, so a
+        # destination that *is* the source — or contains it — would delete the
+        # store we are about to read. Checked against resolved paths so a symlink
+        # or `.` in the argument cannot slip past.
+        resolved_source = source_path.resolve()
+        resolved_dest = store_path.resolve()
+        if resolved_dest == resolved_source or resolved_dest in resolved_source.parents:
+            raise WriterCommandError(
+                f"Destination would delete the source store: {store_path}\n"
+                "Pass no destination to convert in place."
+            )
+        if store_path.exists():
+            if not overwrite:
+                raise WriterCommandError(
+                    f"Destination already exists: {store_path}\nPass overwrite to replace it."
+                )
+            shutil.rmtree(store_path)
+        shutil.copytree(source_path, store_path)
+
+    available = list_element_keys(store_path, "tables")
+    if not available:
+        raise WriterCommandError(f"No table elements found under {store_path / 'tables'}")
+
+    selected = list(tables) if tables else available
+    unknown = [key for key in selected if key not in available]
+    if unknown:
+        raise WriterCommandError(
+            f"Unknown table element(s): {', '.join(unknown)}. Available: {', '.join(available)}"
+        )
+
+    reports: list[dict[str, Any]] = []
+    for key in selected:
+        table_path = store_path / "tables" / key
+        n_obs, n_vars = _table_dimensions(table_path)
+        matrices = _convert_table_matrices(table_path, layers=layers, densify=densify)
+        reports.append(
+            {
+                "path": f"tables/{key}",
+                "n_obs": n_obs,
+                "n_vars": n_vars,
+                "zarr_format": zarr_format_of(table_path),
+                "matrices": matrices,
+            }
+        )
+
+    refresh_consolidated_metadata(store_path)
+
+    manifest_data = {
+        "format": "spatialdata-csc-tables/v1",
+        "source": str(source_path),
+        "output": str(store_path),
+        "in_place": in_place,
+        "tables": reports,
+        "packages": {
+            "anndata": package_version("anndata"),
+            "scipy": package_version("scipy"),
+            "spatialdata": package_version("spatialdata"),
+            "spatialdata-js-util": package_version("spatialdata-js-util"),
+            "zarr": package_version("zarr"),
+        },
+    }
+    manifest_path = store_path.with_suffix(".csc-manifest.json") if manifest else None
+    if manifest_path is not None:
+        write_json(manifest_path, manifest_data)
+    return CscConversion(store_path, manifest_path, manifest_data)