diff --git a/docs/tutorials/cohorts.md b/docs/tutorials/cohorts.md index b25b648..3edba0f 100644 --- a/docs/tutorials/cohorts.md +++ b/docs/tutorials/cohorts.md @@ -57,65 +57,48 @@ Only URLs are replayed, so a store republished at the same URL is not detected. Pass `pin_versions=True` to record the catalog version; downloading into a directory pinned to a different version then raises `ManifestVersionMismatch` unless `allow_version_change=True`. -## Freeze (network, not executed here) +## Freeze (offline, not executed here) -`freeze()` pins the selection to its exact members, reads each store's fingerprint, and returns a -`FrozenCohort` with a deterministic `hash`. `write` saves it as JSON (large member lists go to a parquet -sidecar); `FrozenCohort.load` restores it offline, without a catalog. +`freeze()` pins the selection to its exact members (full uids) and returns a `CohortManifest` with a +deterministic `hash`. `write` saves it as JSON (large member lists go to a parquet sidecar); +`CohortManifest.load` restores it offline, and `SpatialDataCohort.from_manifest` rebuilds the cohort +without a catalog query. ```python # not executed -from sddb import CohortDriftError, FrozenCohort +from sddb import CohortManifest, SpatialDataCohort -frozen = cohort.freeze() -frozen.write("cohort.json") -frozen = FrozenCohort.load("cohort.json") # offline reload -frozen.hash, frozen.members +manifest = cohort.freeze() +manifest.write("cohort.json") +manifest = CohortManifest.load("cohort.json") # offline reload +cohort = SpatialDataCohort.from_manifest(manifest) +manifest.hash ``` -### Verify-on-read +### Liveness and updates -A `zarr_url` is a "latest" pointer, not a version-pinned one. So reading a frozen cohort's data -(`frozen.to_anndata()`, `frozen.build(...)`) re-checks each object's fingerprint first and raises -`CohortDriftError` if a store changed since the freeze. Pass `verify=False` to skip the check. -A live `SpatialDataCohort.to_anndata()` is unfrozen and performs no such check. - -```python -# not executed -adata = frozen.to_anndata() # table=None: default table of each member -``` +A `zarr_url` is a "latest" pointer, not a version-pinned one. `cohort.verify()` checks that every pinned +member still resolves, and `cohort.check_updates(catalog)` reports members that a newer catalog version +supersedes. `download(strict=...)` controls whether a missing pinned store is an error. ## Split `split` assigns whole studies (`by="study_id"`) to one fold, so no study leaks across train, val and test. It is deterministic for a given `seed`, independent of row order, and refuses to create empty folds on -tiny cohorts. The `SplitManifest` records the parent cohort hash, and can be written and loaded. +tiny cohorts. It returns `{fold: SpatialDataCohort}`; each fold carries a `SplitManifest` (pinned to the +parent cohort hash) as `.split_manifest`, which can be written and loaded. ```python # not executed from sddb import SplitManifest -split = frozen.split(by="study_id", train=0.8, val=0.1, test=0.1, seed=42) -split.write("split.json") +folds = cohort.split(by="study_id", train=0.8, val=0.1, test=0.1, seed=42) +folds["train"].split_manifest.write("split.json") split = SplitManifest.load("split.json") ``` -## Check and build with a task adapter - -A `TaskAdapter` (protocol in `sddb.adapter`) states what a task needs (`requirements`), validates metadata -(`validate_meta`) and data (`validate`), and turns a store into a model input (`build`). Concrete adapters -live in method repos, not in this client; `my_adapter` below stands for one of them. +## Task adapters -`check` reports compatibility from metadata (add `deep=True` to also open the stores). `build` returns a lazy -task view; with `materialize=True` each member is built once and provenance is written to `out`. A `split` -must derive from the same cohort, otherwise `ValueError` is raised. - -```python -# not executed -report = frozen.check(my_adapter) -print(report) - -view = frozen.build(my_adapter, split=split, materialize=True, out="./task") -for uid, load in view: - sdata = load() -``` +A `TaskAdapter` (protocol in `sddb.adapter`) states what a task needs and turns a store into a model input. +Concrete adapters and their execution live in method repos, not in this client. Use `cohort.sdatas` or +`cohort.load(elements=[...])` to get the data to hand to an adapter. diff --git a/src/sddb/__init__.py b/src/sddb/__init__.py index e7a28b7..f0dc72e 100644 --- a/src/sddb/__init__.py +++ b/src/sddb/__init__.py @@ -7,23 +7,21 @@ except PackageNotFoundError: # pragma: no cover __version__ = "0.0.0" -from sddb.adapter import PreflightReport, TaskAdapter +from sddb.adapter import TaskAdapter from sddb.catalog import Catalog from sddb.cohort import SpatialDataCohort from sddb.dataset import Dataset -from sddb.freeze import CohortDriftError, FrozenCohort -from sddb.manifest import Manifest, ManifestVersionMismatch +from sddb.manifest import CohortDriftError, CohortManifest, Manifest, ManifestVersionMismatch from sddb.remote import open_sdata from sddb.split import SplitManifest __all__ = [ "Catalog", "CohortDriftError", + "CohortManifest", "Dataset", - "FrozenCohort", "Manifest", "ManifestVersionMismatch", - "PreflightReport", "SpatialDataCohort", "SplitManifest", "TaskAdapter", diff --git a/src/sddb/_catalog_schema.py b/src/sddb/_catalog_schema.py index 687578d..70d2529 100644 --- a/src/sddb/_catalog_schema.py +++ b/src/sddb/_catalog_schema.py @@ -114,6 +114,8 @@ COLLECTION_REQUIRED: tuple[str, ...] = ("slug", "name", "members") FACETS: tuple[str, ...] = ( + "study_id", + "doi", "organism", "tissue", "disease", diff --git a/src/sddb/adapter.py b/src/sddb/adapter.py index 9eebab8..c4e4ca8 100644 --- a/src/sddb/adapter.py +++ b/src/sddb/adapter.py @@ -1,21 +1,14 @@ -"""Task-adapter interface: turn a frozen cohort into a task-specific view with provenance.""" +"""Task-adapter interface: the contract a downstream task repo implements. + +The client defines only this Protocol. Building/materializing a task view from a cohort lives +downstream of the cohort, not in the client. +""" from __future__ import annotations -import json -import re -from collections.abc import Callable, Iterator -from dataclasses import dataclass -from datetime import UTC, datetime -from pathlib import Path from typing import Any, Protocol, runtime_checkable -import pandas as pd - -import sddb -from sddb import remote -from sddb.concat import _table_name -from sddb.objectmeta import ObjectMeta, fetch_many +from sddb.objectmeta import ObjectMeta Issue = str @@ -43,178 +36,3 @@ def validate(self, sdata: Any) -> list[Issue]: def build(self, sdata: Any) -> Any: """Produce the task-specific item for one object.""" ... - - -@dataclass -class PreflightReport: - """Counts of compatible members and excluded uids grouped by reason.""" - - total: int - compatible: int - excluded: dict[Issue, list[str]] # reason -> uids - tier: str # "advisory" | "authoritative" - - def __repr__(self) -> str: - lines = [f"{self.total} candidate / {self.compatible} compatible [{self.tier}]"] - lines += [f" {len(uids)} {reason}: {', '.join(uids)}" for reason, uids in sorted(self.excluded.items())] - return "\n".join(lines) - - -_ELEMENT_KINDS = remote._ELEMENT_KINDS - - -def _metas(members: pd.DataFrame, **kw: Any) -> list[ObjectMeta]: - return fetch_many([(str(r.uid), str(r.zarr_url)) for r in members.itertuples(index=False)], **kw) - - -def _match(name: str, select: str, how: str) -> bool: - """exact/suffix/substring are case-insensitive; regex uses the pattern's own flags.""" - n, s = name.lower(), select.lower() - if how == "exact": - return n == s - if how == "suffix": # anchored: `square_016um` must not catch `square_016um_for_vitessce` - return n == s or (n.endswith(s) and not n[len(n) - len(s) - 1].isalnum()) - if how == "substring": - return s in n - if how == "regex": - return re.search(select, name) is not None - raise ValueError(f"unknown match mode {how!r}") - - -def _sdata_elements(sdata: Any) -> list[str]: - """Sorted ``"kind/name"`` of every element present in an opened SpatialData.""" - return sorted(f"{kind}/{name}" for kind in _ELEMENT_KINDS for name in getattr(sdata, kind)) - - -def _resolve_elements(url: str, roles: list[dict[str, Any]]) -> list[str]: - from sddb.remote import elements as _elements - - inv = _elements(url) # {"kind/name": {...}} - resolved: list[str] = [] - for role in roles: - kind = role["role"] - kind = kind if kind in _ELEMENT_KINDS else kind + "s" # image -> images, shape -> shapes, ... - select, how = role["select"], role.get("match", "suffix") - hits = [p for p in inv if p.split("/", 1)[0] == kind and _match(p.split("/", 1)[1], select, how)] - if len(hits) == 1: - resolved.append(hits[0]) - elif not hits: - raise ValueError(f"role {role!r}: no element matched; present: {sorted(inv)}") - else: - raise ValueError(f"role {role!r}: matched {len(hits)} elements {hits}; tighten `select`/`match`") - return resolved - - -def _requirements(adapter_obj: Any) -> dict[str, Any]: - req = getattr(adapter_obj, "requirements", None) - return (req() if callable(req) else None) or {} - - -def _open(url: str, adapter_obj: TaskAdapter) -> Any: - from sddb.remote import open_sdata - - roles = _requirements(adapter_obj).get("elements") - if not roles: - return open_sdata(url, lazy=False) - return open_sdata(url, lazy=False, elements=_resolve_elements(url, roles)) - - -def check(members: pd.DataFrame, adapter: TaskAdapter, *, deep: bool = False) -> PreflightReport: - """Compatibility report. ``deep=False`` uses metadata only (advisory); ``deep=True`` opens each object.""" - excluded: dict[Issue, list[str]] = {} - compatible = 0 - if deep: - pairs = [(str(r.uid), adapter.validate(_open(str(r.zarr_url), adapter))) for r in members.itertuples(index=False)] - tier = "authoritative" - else: - pairs = [(m.uid, adapter.validate_meta(m)) for m in _metas(members)] - tier = "advisory" - for uid, issues in pairs: - if issues: - for i in issues: - excluded.setdefault(i, []).append(uid) - else: - compatible += 1 - return PreflightReport(len(members), compatible, excluded, tier) - - -class TaskView: - """Lazy iterable of ``(uid, thunk)``; calling a thunk opens the object and runs ``adapter.build``.""" - - def __init__(self, items: list[tuple[str, Callable[[], Any]]]) -> None: - self._items = items - - def __iter__(self) -> Iterator[tuple[str, Callable[[], Any]]]: - return iter(self._items) - - def __len__(self) -> int: - return len(self._items) - - def __getitem__(self, i: int) -> tuple[str, Callable[[], Any]]: - return self._items[i] - - -def build( - members: pd.DataFrame, - adapter: TaskAdapter, - *, - split: Any = None, - materialize: bool = False, - out: str | Path | None = None, - verify: bool = True, - cohort_hash: str = "", -) -> TaskView: - """Task view over ``members``. Lazy by default; ``materialize=True`` builds all and writes provenance. - - With ``materialize=True`` each ``adapter.build(sdata)`` runs once for its side effects and provenance; the - returned ``TaskView`` stays lazy (re-invoking a thunk re-opens the store). - """ - if split is not None and split.parent_hash != cohort_hash: - raise ValueError(f"split.parent_hash {split.parent_hash!r} does not match cohort hash {cohort_hash!r}") - if verify: - from sddb.freeze import verify_members - - verify_members(members) - - def thunk(url: str) -> Callable[[], Any]: - return lambda: adapter.build(_open(url, adapter)) - - view = TaskView([(str(r.uid), thunk(str(r.zarr_url))) for r in members.itertuples(index=False)]) - if not materialize: - return view - out_dir = Path(out or ".") - out_dir.mkdir(parents=True, exist_ok=True) - reqs = _requirements(adapter) - roles = reqs.get("elements") - produced = [] - for r in members.itertuples(index=False): - sdata = _open(str(r.zarr_url), adapter) - out_item = adapter.build(sdata) - produced.append( - { - "uid": str(r.uid), - "table": _table_name(r), - "output": type(out_item).__name__, - "elements": _sdata_elements(sdata) if roles else None, - } - ) - split_id = ( - None - if split is None - else {"by": split.by, "seed": split.seed, "ratios": split.ratios, "mode": split.mode} - ) - (out_dir / "provenance.json").write_text( - json.dumps( - { - "cohort_hash": cohort_hash, - "created_at": datetime.now(UTC).isoformat(), - "split": split_id, - "adapter": {"name": adapter.name, "version": adapter.version, "config": adapter.config}, - "sddb_version": sddb.__version__, - "requirements": reqs, - "members": produced, - }, - indent=2, - ) - ) - return view diff --git a/src/sddb/catalog.py b/src/sddb/catalog.py index 21be451..bdb7769 100644 --- a/src/sddb/catalog.py +++ b/src/sddb/catalog.py @@ -272,7 +272,7 @@ def search(self, text: str) -> SpatialDataCohort: tokens = re.findall(r"\w+", text.lower()) matched: dict[str, list[str]] = {} for col in schema.FACETS: - if col not in self._df.columns: + if col == "doi" or col not in self._df.columns: # DOIs tokenize to "10" etc. continue values = sorted(str(v) for v in self._df[col].dropna().unique()) hits = [v for v in values if any(_token_matches(t, v) for t in tokens)] @@ -289,3 +289,23 @@ def search(self, text: str) -> SpatialDataCohort: return SpatialDataCohort( self._df[mask.fillna(False).astype(bool)], matched=matched, source=self._source(), filter={"search": text} ) + + def cohort(self, *, uids: list[str] | None = None, study_ids: list[str] | None = None) -> SpatialDataCohort: + """Build a cohort from an explicit set of ``uids`` or ``study_ids`` (exactly one). + + Raises ``KeyError`` naming any id not in the catalog; a valid but empty result warns. + """ + if (uids is None) == (study_ids is None): + raise ValueError("pass exactly one of uids= or study_ids=") + col, want = ("uid", uids) if uids is not None else ("study_id", study_ids) + if col not in self._df.columns: + raise ValueError(f"catalog has no {col!r} column") + want = [str(x) for x in want] + sub = self._df[self._df[col].astype("string").isin(want)] + found = set(sub[col].astype("string")) + missing = [x for x in want if x not in found] + if missing: + raise KeyError(f"unknown {col}(s) not in catalog: {missing}") + if sub.empty: + warnings.warn(f"no catalog rows for {col} in {want}", stacklevel=2) + return SpatialDataCohort(sub, source=self._source(), filter={col: sorted(want)}) diff --git a/src/sddb/cohort.py b/src/sddb/cohort.py index 93ccaa9..6cf7011 100644 --- a/src/sddb/cohort.py +++ b/src/sddb/cohort.py @@ -2,19 +2,23 @@ from __future__ import annotations +import warnings from collections.abc import Iterator, Sequence +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import TYPE_CHECKING, Any, overload import pandas as pd from sddb.dataset import Dataset, Source +from sddb.remote import _open_fs if TYPE_CHECKING: from anndata import AnnData + from spatialdata import SpatialData - from sddb.freeze import FrozenCohort - from sddb.manifest import Manifest + from sddb.catalog import Catalog + from sddb.manifest import CohortManifest, Manifest _MODALITY_FLAGS = ("hne_image", "if_image", "ftu_annotation") # catalog boolean columns @@ -72,38 +76,140 @@ def __init__( self.matched: dict[str, list[str]] = matched or {} self._source = source or Source() self._filter: dict[str, Any] = filter or {} + self.split_manifest: Any = None # set when this cohort is a fold of a split() - def object_metadata(self, *, workers: int = 8) -> pd.DataFrame: - """Fetch metadata-only fingerprints for every member (all-or-nothing; raises ``ObjectMetaError``).""" - from sddb.objectmeta import fetch_many + def freeze(self) -> CohortManifest: + """Pin this selection to its exact members and a deterministic hash. Offline/cheap. - pairs = [(str(r.uid), str(r.zarr_url)) for r in self._df.itertuples(index=False)] - metas = fetch_many(pairs, workers=workers) - return pd.DataFrame( - [ - {"uid": m.uid, "lamin_version": m.lamin_version, "fingerprint": m.fingerprint, "tables": list(m.tables)} - for m in metas - ], - columns=["uid", "lamin_version", "fingerprint", "tables"], - ) - - def freeze(self, *, workers: int = 8) -> FrozenCohort: - """Pin this selection to its exact members and a deterministic content hash. - - The recorded filter captures effective defaults (e.g. ``validation="pass"``). Cohorts built by - paths other than ``query``/``search`` may carry an empty filter. + The recorded filter captures effective defaults (e.g. ``validation="pass"``). Cohorts built + by paths other than ``query``/``search`` may carry an empty filter. """ from datetime import UTC, datetime - from sddb.freeze import SCHEMA_VERSION, FrozenCohort, cohort_hash + from sddb.manifest import SCHEMA_VERSION, CohortManifest, cohort_hash - meta = self.object_metadata(workers=workers) - members = self._df.merge(meta, on="uid", how="left") flt = dict(self._filter) - return FrozenCohort( - SCHEMA_VERSION, datetime.now(UTC).isoformat(), self._source, flt, cohort_hash(flt, members), members + return CohortManifest( + SCHEMA_VERSION, + datetime.now(UTC).isoformat(), + self._source, + flt, + cohort_hash(flt, self._df["uid"]), + self._df.copy(), ) + @classmethod + def from_manifest(cls, m: CohortManifest, *, catalog: Catalog | None = None) -> SpatialDataCohort: + """Rebuild a cohort pinned to exactly the manifest's members (offline — no catalog re-query). + + If ``catalog`` is given, warn once about any members a newer version supersedes (see + :meth:`check_updates`). + """ + coh = cls(m.members, source=m.source, filter=m.filter) + if catalog is not None: + coh.check_updates(catalog, warn=True) + return coh + + def verify(self, *, workers: int = 8) -> None: + """Liveness check: raise ``CohortDriftError`` if any pinned store no longer resolves. + + Immutable version-exact stores cannot drift, so existence is the only failure mode. Not called by + :meth:`download`, which runs its own inline liveness loop; not run by cheap ops. + """ + from sddb.manifest import CohortDriftError + + pairs = [(str(r.uid), str(r.zarr_url)) for r in self._df.itertuples(index=False)] + + def resolves(uid: str, url: str) -> tuple[str, bool]: + try: + fs, root = _open_fs(url) + return uid, bool(fs.exists(root)) + except Exception: + return uid, False + + with ThreadPoolExecutor(max_workers=max(1, workers)) as pool: + results = list(pool.map(lambda p: resolves(*p), pairs)) + missing = {uid: "store does not resolve" for uid, ok in results if not ok} + if missing: + raise CohortDriftError(missing) + + def check_updates(self, catalog: Catalog, *, warn: bool = True) -> dict[str, str]: + """Return ``{pinned_uid: current_uid}`` for members a newer catalog version supersedes. + + Matches each pinned ``uid[:16]`` (stem) against the catalog's current uid for that stem (the + catalog is latest-only). Warns once listing the superseded members when ``warn``. + """ + current = catalog.to_df() + latest_by_stem = {str(u)[:16]: str(u) for u in current["uid"]} + superseded: dict[str, str] = {} + for u in self._df["uid"]: + u = str(u) + cur = latest_by_stem.get(u[:16]) + if cur is not None and cur != u: + superseded[u] = cur + if superseded and warn: + warnings.warn( + f"{len(superseded)} member(s) superseded by a newer version: {sorted(superseded)}", + stacklevel=2, + ) + return superseded + + def _derive(self, df: pd.DataFrame, filter: dict[str, Any]) -> SpatialDataCohort: + """A new cohort over ``df`` carrying this cohort's source (used by subsetting/splits).""" + return SpatialDataCohort(df, matched=self.matched, source=self._source, filter=filter) + + def split( + self, *, by: str = "study_id", train: float, val: float, test: float, seed: int + ) -> dict[str, SpatialDataCohort]: + """Group-safe deterministic split into ``{fold: SpatialDataCohort}``. + + Each returned cohort carries the derived ``SplitManifest`` as ``.split_manifest`` (pinned to + this cohort's hash). Splitting is always group-safe (``mode="group"``). + """ + from sddb.manifest import cohort_hash + from sddb.split import make_split + + parent = cohort_hash(self._filter, self._df["uid"]) + sm = make_split(self._df, parent, by=by, train=train, val=val, test=test, seed=seed, mode="group") + uid_str = self._df["uid"].astype("string") + folds: dict[str, SpatialDataCohort] = {} + for fold in ("train", "val", "test"): + want = [u for u, f in sm.assignments.items() if f == fold] + sub = self._derive(self._df[uid_str.isin(want)], self._filter) + sub.split_manifest = sm + folds[fold] = sub + return folds + + @property + def members(self) -> pd.DataFrame: + """The underlying catalog rows (a view; treat as read-only — use methods to derive new cohorts).""" + return self._df + + def filter(self, **facets: Any) -> SpatialDataCohort: + """Return a new cohort keeping rows matching every ``column=value`` (a list/tuple/set means isin).""" + df = self._df + mask = pd.Series(True, index=df.index) + norm: dict[str, Any] = {} + for col, val in facets.items(): + if col not in df.columns: + raise ValueError(f"unknown column {col!r}; present: {list(df.columns)}") + mask &= df[col].isin(list(val) if isinstance(val, (list, tuple, set)) else [val]) + norm[col] = sorted(val, key=str) if isinstance(val, (list, tuple, set)) else val + return self._derive(df[mask.fillna(False).astype(bool)], {**self._filter, **norm}) + + def where(self, mask: pd.Series) -> SpatialDataCohort: + """Return a new cohort keeping rows where the boolean ``mask`` (aligned to ``members``) is True.""" + return self._derive(self._df[mask.fillna(False).astype(bool)], self._filter) + + def __or__(self, other: SpatialDataCohort) -> SpatialDataCohort: + """Union, deduped by uid (first occurrence wins). Provenance filter is dropped (no longer a pure facet set).""" + df = pd.concat([self._df, other._df]).drop_duplicates("uid", keep="first") + return self._derive(df, {}) + + def __sub__(self, other: SpatialDataCohort) -> SpatialDataCohort: + """Difference: drop any member whose uid is in ``other``.""" + return self._derive(self._df[~self._df["uid"].isin(other._df["uid"])], self._filter) + def citations(self, path: str | Path, *, bib_url: str | None = None) -> Path: """Write the BibTeX entries for the studies in this set to ``path``.""" from sddb.citations import write_citations @@ -117,24 +223,36 @@ def citations_text(self, *, bib_url: str | None = None) -> tuple[str, int]: return citations_text(self, bib_url=bib_url) def download( - self, dest: str | Path, *, workers: int = 4, pin_versions: bool = False, allow_version_change: bool = False + self, + dest: str | Path, + *, + workers: int = 4, + pin_versions: bool = False, + allow_version_change: bool = False, + strict: bool = False, ) -> Manifest: """Download every store under ``dest`` (resumable) and return the manifest. Raises ``ManifestVersionMismatch`` if ``dest`` holds a manifest pinned to another catalog version - (unless ``allow_version_change``). + (unless ``allow_version_change``). A missing pinned store falls back to the latest version with a + loud warning unless ``strict=True`` (then it raises ``CohortDriftError``). """ from sddb.manifest import download return download( - self, dest, workers=workers, pin_versions=pin_versions, allow_version_change=allow_version_change + self, + dest, + workers=workers, + pin_versions=pin_versions, + allow_version_change=allow_version_change, + strict=strict, ) def to_anndata(self, *, table: str | None = None) -> AnnData: - """Concatenate each member's table into one AnnData (no drift check: a live cohort is unfrozen).""" + """Concatenate each member's table into one AnnData (outer join, obs[uid]+obs[study_id]).""" from sddb.concat import concat_tables - return concat_tables(self._df, table=table, verify=False) + return concat_tables(self._df, table=table) def to_df(self) -> pd.DataFrame: """Return the underlying rows as a DataFrame copy.""" @@ -160,6 +278,15 @@ def summary(self) -> dict[str, Any]: """Headline numbers for this set.""" return summary(self._df) + @property + def sdatas(self) -> list[SpatialData]: + """Open every member as a full SpatialData (download-then-open; explicit and potentially large).""" + return [self[i].load(lazy=False) for i in range(len(self))] + + def load(self, *, elements: list[str] | None = None) -> list[SpatialData]: + """Open every member, localizing only ``elements`` (``"kind/name"``) when given.""" + return [self[i].load(lazy=False, elements=elements) for i in range(len(self))] + def __len__(self) -> int: return len(self._df) diff --git a/src/sddb/concat.py b/src/sddb/concat.py index 1ec0dbd..96e0d8c 100644 --- a/src/sddb/concat.py +++ b/src/sddb/concat.py @@ -15,30 +15,17 @@ def _table_name(row: Any) -> str: return str(dt) if isinstance(dt, str) and dt else "table" -def concat_tables(members: pd.DataFrame, *, table: str | None = None, verify: bool = True) -> ad.AnnData: - """Concatenate one table per member into a single AnnData (outer join, ``obs["uid"]`` labels the member). +def concat_tables(members: pd.DataFrame, *, table: str | None = None) -> ad.AnnData: + """Concatenate one table per member into one AnnData (outer join). - Only ``/tables/`` is opened, so images/shapes are never touched. Works for local - paths and (anonymous) remote URLs. - - Parameters - ---------- - members - DataFrame with ``uid`` and ``zarr_url`` (optionally ``default_table``, ``fingerprint``). - table - Table name to read from every member; defaults to each member's ``default_table`` else ``"table"``. - verify - If True, run :func:`sddb.freeze.verify_members` first (raises ``CohortDriftError`` on drift). + ``obs['uid']`` labels the member; ``obs['study_id']`` is added from the members frame (all-NA if + the column is absent). Only ``/tables/`` is opened — images/shapes are untouched. Raises ------ KeyError If a member's table is absent or unreadable; the message names the member uid. """ - if verify: - from sddb.freeze import verify_members - - verify_members(members) parts: list[ad.AnnData] = [] keys: list[str] = [] for row in members.itertuples(index=False): @@ -49,4 +36,11 @@ def concat_tables(members: pd.DataFrame, *, table: str | None = None, verify: bo except Exception as e: raise KeyError(f"member {row.uid!r}: table {name!r} not readable at {url}: {e}") from e keys.append(str(row.uid)) - return ad.concat(parts, join="outer", label="uid", keys=keys, index_unique="-") + out = ad.concat(parts, join="outer", label="uid", keys=keys, index_unique="-") + study_by_uid = ( + dict(zip(members["uid"].astype("string"), members["study_id"].astype("string"), strict=True)) + if "study_id" in members.columns + else {} + ) + out.obs["study_id"] = out.obs["uid"].astype("string").map(study_by_uid).astype("string") + return out diff --git a/src/sddb/dataset.py b/src/sddb/dataset.py index ddffb0a..b731399 100644 --- a/src/sddb/dataset.py +++ b/src/sddb/dataset.py @@ -63,7 +63,7 @@ def __getattr__(self, name: str) -> Any: return None raise AttributeError(f"{type(self).__name__!r} object has no attribute {name!r}") - def load(self, *, lazy: bool = True, version: str | None = None) -> SpatialData: + def load(self, *, lazy: bool = True, elements: list[str] | None = None, version: str | None = None) -> SpatialData: """Open this dataset as SpatialData. Parameters @@ -72,6 +72,9 @@ def load(self, *, lazy: bool = True, version: str | None = None) -> SpatialData: Open a local store in place (dask-backed). Remote stores cannot be opened lazily yet (spatialdata upstream limitation); for a remote dataset pass ``lazy=False`` to copy the store to the cache and open it, or use :meth:`elements` to inspect without downloading. + elements + Localize only the selected ``"kind/name"`` elements (e.g. ``["images/he", + "tables/square_016um"]``). version Reserved; version pinning is not implemented yet. @@ -83,7 +86,7 @@ def load(self, *, lazy: bool = True, version: str | None = None) -> SpatialData: if version is not None: raise NotImplementedError("version pinning is not implemented yet") try: - return remote.open_sdata(self.zarr_url, lazy=lazy) + return remote.open_sdata(self.zarr_url, lazy=lazy, elements=elements) except NotImplementedError as e: if "remote" in str(e).lower() and "lazy" in str(e).lower(): raise NotImplementedError( diff --git a/src/sddb/freeze.py b/src/sddb/freeze.py deleted file mode 100644 index 6eb81a9..0000000 --- a/src/sddb/freeze.py +++ /dev/null @@ -1,144 +0,0 @@ -"""FrozenCohort: a citable, hash-pinned snapshot of a cohort selection.""" - -from __future__ import annotations - -import hashlib -import json -from dataclasses import dataclass -from pathlib import Path -from typing import TYPE_CHECKING, Any - -import pandas as pd - -from sddb.dataset import Source - -if TYPE_CHECKING: - from anndata import AnnData - from spatialdata import SpatialData - - from sddb.split import SplitManifest - -SCHEMA_VERSION = "1" -_SIDECAR_THRESHOLD = 200 # members beyond this go to a parquet sidecar - - -def cohort_hash(filter: dict[str, Any], members: pd.DataFrame) -> str: - """sha256 over the filter and ``(uid, lamin_version, fingerprint)`` of each member, sorted by uid.""" - rows = [ - { - "uid": str(r.uid), - "lamin_version": (None if pd.isna(r.lamin_version) else str(r.lamin_version)), - "fingerprint": str(r.fingerprint), - } - for r in members.sort_values("uid").itertuples(index=False) - ] - payload = json.dumps({"filter": filter, "members": rows}, sort_keys=True, separators=(",", ":"), default=str) - return "sha256:" + hashlib.sha256(payload.encode()).hexdigest() - - -@dataclass(frozen=True) -class FrozenCohort: - """Immutable record of a cohort: filter, per-member fingerprints and a deterministic ``hash``.""" - - schema_version: str - created_at: str - source: Source - filter: dict[str, Any] - hash: str - members: pd.DataFrame # read-only by contract; frozen=True does NOT protect in-place mutation - - def to_anndata(self, *, table: str | None = None, verify: bool = True) -> AnnData: - """Concatenate each member's table into one AnnData, verifying fingerprints first by default.""" - from sddb.concat import concat_tables - - return concat_tables(self.members, table=table, verify=verify) - - def open_member(self, uid: str, *, elements: list[str] | None = None, verify: bool = True) -> SpatialData: - """Localize+open one member's store. Verifies the member's fingerprint first (verify-on-read).""" - from sddb.remote import open_sdata - - row = self.members[self.members["uid"].astype(str) == str(uid)] - if row.empty: - raise KeyError(f"uid {uid!r} not in this cohort") - if verify: - verify_members(row) - return open_sdata(str(row.iloc[0]["zarr_url"]), lazy=False, elements=elements) - - def split( - self, *, by: str = "study_id", train: float, val: float, test: float, seed: int, mode: str = "group" - ) -> SplitManifest: - """Group-safe deterministic split of this cohort; the manifest records ``self.hash`` as ``parent_hash``.""" - from sddb.split import make_split - - return make_split(self.members, self.hash, by=by, train=train, val=val, test=test, seed=seed, mode=mode) - - def check(self, adapter: Any, *, deep: bool = False) -> Any: - """Compatibility report of this cohort against a TaskAdapter (advisory unless ``deep``).""" - from sddb.adapter import check - - return check(self.members, adapter, deep=deep) - - def build( - self, adapter: Any, *, split: Any = None, materialize: bool = False, out: str | Path | None = None, verify: bool = True - ) -> Any: - """Task view via ``adapter``; ``split`` must derive from this cohort. Lazy unless ``materialize``.""" - from sddb.adapter import build - - return build( - self.members, adapter, split=split, materialize=materialize, out=out, verify=verify, cohort_hash=self.hash - ) - - def write(self, path: str | Path) -> Path: - """Write JSON to ``path`` (members go to a parquet sidecar when large); return the path.""" - p = Path(path) - head = { - "schema_version": self.schema_version, - "created_at": self.created_at, - "source": {"url": self.source.url, "version": self.source.version}, - "filter": self.filter, - "hash": self.hash, - "n_members": len(self.members), - } - if len(self.members) > _SIDECAR_THRESHOLD: - side = p.with_suffix(".members.parquet") - self.members.to_parquet(side) - head["members_ref"] = side.name - else: - head["members"] = json.loads(self.members.to_json(orient="records")) - p.write_text(json.dumps(head, indent=2)) - return p - - @classmethod - def load(cls, path: str | Path) -> FrozenCohort: - """Read a cohort written by :meth:`write`.""" - p = Path(path) - head = json.loads(p.read_text()) - if "members_ref" in head: - members = pd.read_parquet(p.parent / head["members_ref"]) - else: - members = pd.DataFrame(head["members"]) - src = Source(url=head["source"]["url"], version=head["source"]["version"]) - return cls(head["schema_version"], head["created_at"], src, head["filter"], head["hash"], members) - - -class CohortDriftError(RuntimeError): - """A frozen cohort's objects changed since freeze; `drifted` maps uid -> reason.""" - - def __init__(self, drifted: dict[str, str]) -> None: - self.drifted = drifted - super().__init__(f"{len(drifted)} object(s) drifted since freeze: {sorted(drifted)}") - - -def verify_members(members: pd.DataFrame, *, workers: int = 8) -> None: - """Re-fetch each member's metadata and raise CohortDriftError on fingerprint mismatch.""" - from sddb.objectmeta import fetch_many - - pairs = [(str(r.uid), str(r.zarr_url)) for r in members.itertuples(index=False)] - now = {m.uid: m for m in fetch_many(pairs, workers=workers)} - drifted: dict[str, str] = {} - for r in members.itertuples(index=False): - cur = now[str(r.uid)] - if cur.fingerprint != str(r.fingerprint): - drifted[str(r.uid)] = f"fingerprint {r.fingerprint} -> {cur.fingerprint}" - if drifted: - raise CohortDriftError(drifted) diff --git a/src/sddb/manifest.py b/src/sddb/manifest.py index ef79bc4..839932e 100644 --- a/src/sddb/manifest.py +++ b/src/sddb/manifest.py @@ -1,4 +1,4 @@ -"""Resumable cohort download with a JSON manifest.""" +"""Cohort identity record (CohortManifest, cohort_hash, CohortDriftError) and resumable cohort download with a JSON manifest.""" from __future__ import annotations @@ -6,20 +6,92 @@ import json import shutil import threading +import warnings +from collections.abc import Iterable from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict, dataclass from datetime import UTC, datetime from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import fsspec +import pandas as pd +from sddb.dataset import Source from sddb.remote import storage_options if TYPE_CHECKING: from sddb.cohort import SpatialDataCohort +SCHEMA_VERSION = "2" +_SIDECAR_THRESHOLD = 200 # members beyond this go to a parquet sidecar + + +def cohort_hash(filter: dict[str, Any], uids: Iterable[Any]) -> str: + """sha256 over the canonical filter and the sorted full lamin uids (version is baked into the uid).""" + payload = json.dumps( + {"filter": filter, "uids": sorted(str(u) for u in uids)}, + sort_keys=True, + separators=(",", ":"), + default=str, + ) + return "sha256:" + hashlib.sha256(payload.encode()).hexdigest() + + +class CohortDriftError(RuntimeError): + """A pinned cohort member no longer resolves; `drifted` maps uid -> reason.""" + + def __init__(self, drifted: dict[str, str]) -> None: + self.drifted = drifted + super().__init__(f"{len(drifted)} member(s) no longer resolve: {sorted(drifted)}") + + +@dataclass(frozen=True, eq=False) +class CohortManifest: + """Serializable cohort identity: filter, member snapshot (pinned full uids) and a deterministic hash.""" + + schema_version: str + created_at: str + source: Source + filter: dict[str, Any] + hash: str + members: pd.DataFrame # read-only by contract + + def write(self, path: str | Path) -> Path: + """Write JSON to ``path``; members spill to a parquet sidecar when large. Returns the path.""" + p = Path(path) + head: dict[str, Any] = { + "schema_version": self.schema_version, + "created_at": self.created_at, + "source": {"url": self.source.url, "version": self.source.version}, + "filter": self.filter, + "hash": self.hash, + "n_members": len(self.members), + } + if len(self.members) > _SIDECAR_THRESHOLD: + side = p.with_suffix(".members.parquet") + self.members.to_parquet(side) + head["members_ref"] = side.name + else: + head["members_columns"] = list(self.members.columns) + head["members"] = json.loads(self.members.to_json(orient="records")) + p.write_text(json.dumps(head, indent=2)) + return p + + @classmethod + def load(cls, path: str | Path) -> CohortManifest: + """Read a manifest written by :meth:`write`.""" + p = Path(path) + head = json.loads(p.read_text()) + if "members_ref" in head: + members = pd.read_parquet(p.parent / head["members_ref"]) + else: + members = pd.DataFrame(head["members"], columns=head.get("members_columns")) + src = Source(url=head["source"]["url"], version=head["source"]["version"]) + return cls(head["schema_version"], head["created_at"], src, head["filter"], head["hash"], members) + + @dataclass class ManifestEntry: """One downloaded store.""" @@ -139,6 +211,40 @@ def plan_sizes(results: SpatialDataCohort) -> list[tuple[str, int]]: return out +def _store_exists(zarr_url: str) -> bool: + """Whether a store resolves on its filesystem (anon S3 or local). + + Only a confirmed non-existence returns False; transient/auth/network errors propagate. + """ + from sddb.remote import _open_fs + + try: + fs, root = _open_fs(zarr_url) + return bool(fs.exists(root)) + except FileNotFoundError: + return False + + +def _latest_for_stem(source: Source, stem: str) -> tuple[str, str] | None: + """Resolve the current ``(uid, zarr_url)`` for a uid stem from the cohort's source catalog. + + Returns None if the catalog is unreachable or the stem is absent. + """ + from sddb.catalog import Catalog + + if not source.url: + return None + try: + df = Catalog(source.url).to_df() + except Exception: + return None + hit = df[df["uid"].astype("string").str[:16] == stem] + if hit.empty: + return None + row = hit.iloc[0] + return str(row["uid"]), str(row["zarr_url"]) + + def download( results: SpatialDataCohort, dest: str | Path, @@ -146,6 +252,7 @@ def download( workers: int = 4, pin_versions: bool = False, allow_version_change: bool = False, + strict: bool = False, ) -> Manifest: """Download every dataset's store to ``dest/.zarr`` and write ``dest/manifest.json``. @@ -156,6 +263,9 @@ def download( ``allow_version_change`` is True (then it is overwritten). Only the catalog version is pinned: per-dataset byte-level pinning across republishes is not implemented. The sha256 covers store file paths and sizes, not contents. + + Each store's liveness is checked first: a missing pinned store raises :class:`CohortDriftError` + when ``strict``, else falls back to the latest version (via the source catalog) with a warning. """ root = Path(dest) current = results._source.version @@ -172,5 +282,37 @@ def download( ) elif recorded is not None: pinned = True - pairs = [(str(d.uid), str(d.zarr_url)) for d in results] + pairs: list[tuple[str, str]] = [] + missing: dict[str, str] = {} + complete: dict[str, ManifestEntry] = {} + if mpath.exists(): + complete = {e.uid: e for e in Manifest.read(mpath).entries if e.status == "complete"} + for d in results: + uid, url = str(d.uid), str(d.zarr_url) + prev = complete.get(uid) + if prev is not None and prev.zarr_url == url and (root / f"{uid}.zarr").is_dir(): + pairs.append((uid, url)) # already complete locally: no liveness check, no fetch + continue + try: + alive = _store_exists(url) + except Exception as err: + raise RuntimeError(f"could not determine liveness of {uid} store at {url}: {err}") from err + if alive: + pairs.append((uid, url)) + continue + if strict: + missing[uid] = f"pinned store missing: {url}" + continue + latest = _latest_for_stem(results._source, uid[:16]) + if latest is None: + missing[uid] = f"pinned store missing and no fallback resolvable: {url}" + continue + warnings.warn( + f"member {uid}: pinned store missing; FALLING BACK to latest {latest[0]} " + f"({latest[1]}) - bytes DIFFER from the pinned version", + stacklevel=2, + ) + pairs.append((uid, latest[1])) + if missing: + raise CohortDriftError(missing) return _fetch_all(pairs, root, current if pinned else None, workers) diff --git a/tests/_fixtures.py b/tests/_fixtures.py index a45ab37..f703401 100644 --- a/tests/_fixtures.py +++ b/tests/_fixtures.py @@ -93,11 +93,40 @@ def make_fixture_catalog() -> pd.DataFrame: def make_fixture_collections() -> pd.DataFrame: """Build a valid synthetic collections table (one row per collection).""" rows = [ - ("c-smith", "Smith2023", "Smith 2023 atlas", "publication", "EpCAM atlas", - "10.1/smith", "DOI", "Xenium", "human", "lung", 2, 2, 180_000, ["uid0001", "uid0002"], - *[None] * 7), # abstract/study_* columns are optional: null here - ("c-lee", "Lee2022", "Lee 2022 brain", "replicates", None, None, None, - "Visium", "mouse", "brain", 1, 1, 4_000, ["uid0003"], *[None] * 7), + ( + "c-smith", + "Smith2023", + "Smith 2023 atlas", + "publication", + "EpCAM atlas", + "10.1/smith", + "DOI", + "Xenium", + "human", + "lung", + 2, + 2, + 180_000, + ["uid0001", "uid0002"], + *[None] * 7, + ), # abstract/study_* columns are optional: null here + ( + "c-lee", + "Lee2022", + "Lee 2022 brain", + "replicates", + None, + None, + None, + "Visium", + "mouse", + "brain", + 1, + 1, + 4_000, + ["uid0003"], + *[None] * 7, + ), ] df = pd.DataFrame(rows, columns=list(COLLECTION_COLUMNS)) return cast_scalars(df, COLLECTION_COLUMNS) diff --git a/tests/test_adapter.py b/tests/test_adapter.py deleted file mode 100644 index 0240de0..0000000 --- a/tests/test_adapter.py +++ /dev/null @@ -1,110 +0,0 @@ -from __future__ import annotations - -import json -from types import SimpleNamespace - -import pandas as pd -import pytest - -from sddb.adapter import build, check -from sddb.objectmeta import ObjectMeta - - -class FakeAdapter: - name = "fake" - version = "0.1" - config = {"k": 1} - - def requirements(self): - return {"tables": ["table"]} - - def validate_meta(self, meta): - return [] if "tables/table" in meta.elements else ["missing table"] - - def validate(self, sdata): - return [] - - def build(self, sdata): - return {"n": 1} - - -def _meta(uid, has_table): - el = {"tables/table": {}} if has_table else {"images/img": {}} - return ObjectMeta(uid, f"z/{uid}", None, el, ("table",) if has_table else (), "sha256:x") - - -def test_check_advisory_tier_counts_and_reasons(monkeypatch): - members = pd.DataFrame({"uid": ["a", "b"], "zarr_url": ["za", "zb"], "fingerprint": ["x", "y"]}) - monkeypatch.setattr("sddb.adapter._metas", lambda m, **k: [_meta("a", True), _meta("b", False)]) - rep = check(members, FakeAdapter(), deep=False) - assert rep.total == 2 - assert rep.compatible == 1 - assert rep.tier == "advisory" - assert "b" in rep.excluded["missing table"] - - -def test_check_deep_uses_authoritative_validate(monkeypatch): - members = pd.DataFrame({"uid": ["a"], "zarr_url": ["za"], "fingerprint": ["x"]}) - monkeypatch.setattr("sddb.adapter._open", lambda url, a: object()) - rep = check(members, FakeAdapter(), deep=True) - assert rep.tier == "authoritative" - assert rep.compatible == 1 - - -def _members(): - return pd.DataFrame({"uid": ["a"], "zarr_url": ["za"], "fingerprint": ["x"], "default_table": ["table"]}) - - -def test_build_materialize_writes_provenance(tmp_path, monkeypatch): - monkeypatch.setattr("sddb.adapter._open", lambda url, a: object()) # skip real sdata open - build(_members(), FakeAdapter(), materialize=True, out=tmp_path, verify=False, cohort_hash="h") - prov = json.loads((tmp_path / "provenance.json").read_text()) - assert prov["cohort_hash"] == "h" - assert prov["split"] is None - assert prov["created_at"] - assert prov["adapter"] == {"name": "fake", "version": "0.1", "config": {"k": 1}} - assert "sddb_version" in prov - assert prov["members"][0] == {"uid": "a", "table": "table", "output": "dict", "elements": None} - - -def test_build_lazy_opens_nothing(monkeypatch): - def boom(url, a): - raise AssertionError("opened") - - monkeypatch.setattr("sddb.adapter._open", boom) - view = build(_members(), FakeAdapter(), verify=False, cohort_hash="h") - assert len(view) == 1 - - -def test_build_split_parent_hash_mismatch_raises(): - split = SimpleNamespace(parent_hash="other") - with pytest.raises(ValueError, match="parent_hash"): - build(_members(), FakeAdapter(), split=split, verify=False, cohort_hash="h") - - -def test_build_split_match_records_split_identity(tmp_path, monkeypatch): - monkeypatch.setattr("sddb.adapter._open", lambda url, a: object()) - split = SimpleNamespace(parent_hash="h", by="study_id", seed=7, ratios={"train": 0.8}, mode="group") - build(_members(), FakeAdapter(), split=split, materialize=True, out=tmp_path, verify=False, cohort_hash="h") - prov = json.loads((tmp_path / "provenance.json").read_text()) - assert prov["cohort_hash"] == "h" - assert prov["split"]["seed"] == 7 - assert prov["split"] == {"by": "study_id", "seed": 7, "ratios": {"train": 0.8}, "mode": "group"} - assert "split_hash" not in prov - - -def test_build_verify_runs_before_open(tmp_path, monkeypatch): - calls = [] - monkeypatch.setattr("sddb.freeze.verify_members", lambda m, **k: calls.append("verify")) - monkeypatch.setattr("sddb.adapter._open", lambda url, a: calls.append("open") or object()) - build(_members(), FakeAdapter(), materialize=True, out=tmp_path, verify=True, cohort_hash="h") - assert calls == ["verify", "open"] - - -def test_build_verify_false_skips_verify(tmp_path, monkeypatch): - calls = [] - monkeypatch.setattr("sddb.freeze.verify_members", lambda m, **k: calls.append("verify")) - monkeypatch.setattr("sddb.adapter._open", lambda url, a: calls.append("open") or object()) - build(_members(), FakeAdapter(), materialize=True, out=tmp_path, verify=False, cohort_hash="h") - assert "verify" not in calls - assert calls == ["open"] diff --git a/tests/test_adapter_protocol.py b/tests/test_adapter_protocol.py new file mode 100644 index 0000000..f202220 --- /dev/null +++ b/tests/test_adapter_protocol.py @@ -0,0 +1,22 @@ +from sddb.adapter import TaskAdapter + + +def test_protocol_is_runtime_checkable(): + class Fake: + name = "fake" + version = "1" + config: dict = {} + + def requirements(self): + return {} + + def validate_meta(self, meta): + return [] + + def validate(self, sdata): + return [] + + def build(self, sdata): + return sdata + + assert isinstance(Fake(), TaskAdapter) diff --git a/tests/test_catalog.py b/tests/test_catalog.py index 01ada85..87a4e9a 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -41,7 +41,8 @@ def test_query_by_license_and_noncommercial(cat): assert {d.uid for d in cat.query(license="CC-BY-NC-SA-4.0")} == {"uid0002"} # the headline filter: exclude NonCommercial (uid0002 is CC-BY-NC-SA) comm = {d.uid for d in cat.query(validation=None, noncommercial=False)} - assert "uid0002" not in comm and "uid0001" in comm + assert "uid0002" not in comm + assert "uid0001" in comm # keep only NonCommercial assert {d.uid for d in cat.query(noncommercial=True)} == {"uid0002"} @@ -323,3 +324,12 @@ def test_query_range_column_absent_from_catalog(tmp_path): df.to_parquet(p) with pytest.raises(ValueError, match="not present in this catalog"): Catalog(p.as_uri(), cache_dir=tmp_path / "cache").query(total_counts__gte=1) + + +def test_search_ignores_doi_fragments(cat): + df = cat._df.copy() + df["doi"] = "10.1038/s41592-000" + cat._df = df + with pytest.warns(UserWarning, match="nothing matched"): + res = cat.search("10") + assert "doi" not in res.matched diff --git a/tests/test_catalog_cohort.py b/tests/test_catalog_cohort.py new file mode 100644 index 0000000..5903c8c --- /dev/null +++ b/tests/test_catalog_cohort.py @@ -0,0 +1,44 @@ +import pandas as pd +import pytest + +from sddb.catalog import Catalog + + +def _cat(tmp_path): + df = pd.DataFrame( + { + "uid": ["u1", "u2", "u3"], + "study_id": ["s_a", "s_a", "s_b"], + "technology": ["Visium HD"] * 3, + "assay": ["spatial"] * 3, + "organism": ["human"] * 3, + "validation_status": ["pass"] * 3, + "zarr_url": ["s3://b/1.zarr", "s3://b/2.zarr", "s3://b/3.zarr"], + } + ) + p = tmp_path / "catalog.parquet" + df.to_parquet(p) + return Catalog.from_file(p) + + +def test_cohort_by_uids(tmp_path): + coh = _cat(tmp_path).cohort(uids=["u1", "u3"]) + assert sorted(coh.to_df()["uid"]) == ["u1", "u3"] + + +def test_cohort_by_study_ids(tmp_path): + coh = _cat(tmp_path).cohort(study_ids=["s_a"]) + assert sorted(coh.to_df()["uid"]) == ["u1", "u2"] + + +def test_cohort_unknown_uid_raises(tmp_path): + with pytest.raises(KeyError, match="u9"): + _cat(tmp_path).cohort(uids=["u1", "u9"]) + + +def test_cohort_requires_exactly_one_selector(tmp_path): + cat = _cat(tmp_path) + with pytest.raises(ValueError, match="exactly one"): + cat.cohort() + with pytest.raises(ValueError, match="exactly one"): + cat.cohort(uids=["u1"], study_ids=["s_a"]) diff --git a/tests/test_catalog_schema.py b/tests/test_catalog_schema.py index b51000b..169e6d9 100644 --- a/tests/test_catalog_schema.py +++ b/tests/test_catalog_schema.py @@ -156,8 +156,13 @@ def test_collection_columns_include_study_fields(): import sddb._catalog_schema as cs expected = { - "abstract", "abstract_license", "abstract_source", - "study_title", "study_authors", "study_year", "study_journal", + "abstract", + "abstract_license", + "abstract_source", + "study_title", + "study_authors", + "study_year", + "study_journal", } assert expected <= set(cs.COLLECTION_COLUMNS) # all are nullable string columns diff --git a/tests/test_cli.py b/tests/test_cli.py index 8c0a066..285d3c9 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -111,7 +111,8 @@ def test_genes_multi_gene_and_validation(tmp_path): write_fixture_gene_index(tmp_path / "gene_index.parquet") r = runner.invoke(app, ["genes", "EPCAM", "PTPRC", "--catalog", cat.as_uri()]) assert r.exit_code == 0, r.output - assert "uid0002" in r.output and "uid0001" not in r.output + assert "uid0002" in r.output + assert "uid0001" not in r.output r = runner.invoke(app, ["genes", "Alb", "--catalog", cat.as_uri()]) assert "uid0004" in r.output # bronze included by default r = runner.invoke(app, ["genes", "Alb", "--validation", "pass", "--catalog", cat.as_uri()]) @@ -124,12 +125,14 @@ def test_genes_ensembl_id_matches_symbol_and_flags_bronze(tmp_path): by_sym = runner.invoke(app, ["genes", "EPCAM", "--catalog", cat.as_uri()]) by_id = runner.invoke(app, ["genes", "ENSG00000119888", "--catalog", cat.as_uri()]) assert by_id.exit_code == 0, by_id.output - assert by_id.output == by_sym.output and "uid0001" in by_id.output + assert by_id.output == by_sym.output + assert "uid0001" in by_id.output assert "validation_status" in by_id.output bronze = runner.invoke(app, ["genes", "Alb", "--catalog", cat.as_uri()]) assert "fail" in bronze.output multi = runner.invoke(app, ["genes", "EPCAM", "KRT8", "--catalog", cat.as_uri()]) - assert "validation_status" in multi.output and "tier" in multi.output + assert "validation_status" in multi.output + assert "tier" in multi.output def test_query_min_obs_excludes_smaller(cat_url): diff --git a/tests/test_cohort_freeze.py b/tests/test_cohort_freeze.py new file mode 100644 index 0000000..764bbd3 --- /dev/null +++ b/tests/test_cohort_freeze.py @@ -0,0 +1,32 @@ +import pandas as pd + +from sddb.cohort import SpatialDataCohort +from sddb.dataset import Source +from sddb.manifest import CohortManifest + + +def _cohort(n=3): + df = pd.DataFrame( + { + "uid": [f"{'u' * 12}{i:04d}0000" for i in range(n)], + "zarr_url": [f"s3://b/.lamindb/u{i}.zarr" for i in range(n)], + "study_id": ["s"] * n, + "technology": ["Visium HD"] * n, + } + ) + return SpatialDataCohort(df, source=Source("cat.parquet", version="v1"), filter={"organism": "human"}) + + +def test_freeze_is_offline_and_roundtrips(tmp_path): + coh = _cohort(3) + m = coh.freeze() + assert isinstance(m, CohortManifest) + assert m.filter == {"organism": "human"} + p = m.write(tmp_path / "c.json") + coh2 = SpatialDataCohort.from_manifest(CohortManifest.load(p)) + assert coh2.freeze().hash == m.hash + assert list(coh2.to_df()["uid"]) == list(coh.to_df()["uid"]) + + +def test_freeze_hash_tracks_membership(): + assert _cohort(3).freeze().hash != _cohort(2).freeze().hash diff --git a/tests/test_cohort_load.py b/tests/test_cohort_load.py new file mode 100644 index 0000000..f669b5f --- /dev/null +++ b/tests/test_cohort_load.py @@ -0,0 +1,38 @@ +import inspect + +import pandas as pd + +from sddb import remote +from sddb.cohort import SpatialDataCohort +from sddb.dataset import Dataset + + +def test_dataset_load_accepts_elements(monkeypatch): + seen = {} + monkeypatch.setattr( + remote, + "open_sdata", + lambda url, lazy=True, elements=None: seen.update(url=url, lazy=lazy, elements=elements) or "SD", + ) + d = Dataset(pd.Series({"uid": "u1", "zarr_url": "s3://b/1.zarr"})) + assert d.load(lazy=False, elements=["images/he"]) == "SD" + assert seen == {"url": "s3://b/1.zarr", "lazy": False, "elements": ["images/he"]} + + +def test_cohort_load_and_sdatas(monkeypatch): + calls = [] + monkeypatch.setattr( + remote, "open_sdata", lambda url, lazy=True, elements=None: calls.append((url, lazy, elements)) or url + ) + coh = SpatialDataCohort(pd.DataFrame({"uid": ["u1", "u2"], "zarr_url": ["s3://b/1.zarr", "s3://b/2.zarr"]})) + out = coh.load(elements=["tables/t"]) + assert out == ["s3://b/1.zarr", "s3://b/2.zarr"] + assert all(lazy is False for _, lazy, _ in calls) + assert all(el == ["tables/t"] for *_, el in calls) + calls.clear() + assert coh.sdatas == ["s3://b/1.zarr", "s3://b/2.zarr"] + assert all(lazy is False and el is None for _, lazy, el in calls) + + +def test_load_elements_param_present(): + assert "elements" in inspect.signature(Dataset.load).parameters diff --git a/tests/test_cohort_manifest.py b/tests/test_cohort_manifest.py new file mode 100644 index 0000000..98dc33b --- /dev/null +++ b/tests/test_cohort_manifest.py @@ -0,0 +1,84 @@ +import pandas as pd + +from sddb.dataset import Source +from sddb.manifest import SCHEMA_VERSION, CohortManifest, cohort_hash + + +def _members(n=3): + return pd.DataFrame( + { + "uid": [f"uid{i:016d}0000"[-20:] for i in range(n)], + "zarr_url": [f"s3://b/.lamindb/uid{i}.zarr" for i in range(n)], + "study_id": ["s"] * n, + "database_version": ["0.0.14"] * n, + } + ) + + +def test_hash_is_order_independent_and_version_sensitive(): + m = _members(3) + h1 = cohort_hash({"organism": "human"}, m["uid"]) + h2 = cohort_hash({"organism": "human"}, m["uid"][::-1]) # reversed order + assert h1 == h2 + assert h1.startswith("sha256:") + changed = m["uid"].tolist() + changed[0] = changed[0][:16] + "0001" # same stem, newer version + assert cohort_hash({"organism": "human"}, pd.Series(changed)) != h1 + + +def test_manifest_roundtrip_small(tmp_path): + m = _members(3) + man = CohortManifest( + SCHEMA_VERSION, + "2026-10-15T00:00:00+00:00", + Source("cat.parquet", version="v1"), + {"organism": "human"}, + cohort_hash({"organism": "human"}, m["uid"]), + m, + ) + p = man.write(tmp_path / "cohort.json") + back = CohortManifest.load(p) + assert back.hash == man.hash + assert back.filter == {"organism": "human"} + assert sorted(back.members["uid"]) == sorted(m["uid"]) + assert back.source.version == "v1" + + +def test_manifest_roundtrip_large_uses_parquet_sidecar(tmp_path): + m = _members(250) # over the 200 sidecar threshold + man = CohortManifest( + SCHEMA_VERSION, + "2026-10-15T00:00:00+00:00", + Source("cat.parquet"), + {}, + cohort_hash({}, m["uid"]), + m, + ) + p = man.write(tmp_path / "big.json") + assert (tmp_path / "big.members.parquet").exists() + back = CohortManifest.load(p) + assert len(back.members) == 250 + assert back.hash == man.hash + + +def test_empty_cohort_roundtrip_keeps_columns(tmp_path): + from sddb.cohort import SpatialDataCohort + + coh = SpatialDataCohort(_members(3).iloc[0:0], source=Source("cat.parquet")) + man = coh.freeze() + back = CohortManifest.load(man.write(tmp_path / "e.json")) + assert list(back.members.columns) == list(man.members.columns) + assert SpatialDataCohort.from_manifest(back).freeze().hash == man.hash + + +def test_set_filter_serializes_and_hash_is_stable(tmp_path): + from sddb.cohort import SpatialDataCohort + + df = _members(3).assign(tissue=["a", "b", "c"]) + hashes = { + SpatialDataCohort(df, source=Source("c.parquet")).filter(tissue={"b", "a"}).freeze().hash for _ in range(5) + } + assert len(hashes) == 1 + man = SpatialDataCohort(df, source=Source("c.parquet")).filter(tissue={"b", "a"}).freeze() + assert man.filter == {"tissue": ["a", "b"]} + man.write(tmp_path / "s.json") diff --git a/tests/test_cohort_split.py b/tests/test_cohort_split.py new file mode 100644 index 0000000..260690d --- /dev/null +++ b/tests/test_cohort_split.py @@ -0,0 +1,35 @@ +import pandas as pd + +from sddb.cohort import SpatialDataCohort +from sddb.split import SplitManifest + + +def _coh(n_studies=6): + uids = [f"u{i}" for i in range(n_studies)] + return SpatialDataCohort( + pd.DataFrame({"uid": uids, "study_id": [f"s{i}" for i in range(n_studies)], "zarr_url": ["x"] * n_studies}) + ) + + +def test_split_returns_sub_cohorts_partitioning_the_parent(): + coh = _coh(6) + folds = coh.split(by="study_id", train=0.5, val=0.25, test=0.25, seed=0) + assert set(folds) == {"train", "val", "test"} + assert all(isinstance(f, SpatialDataCohort) for f in folds.values()) + union = set().union(*[set(f.to_df()["uid"]) for f in folds.values()]) + assert union == set(coh.to_df()["uid"]) + total = sum(len(f) for f in folds.values()) + assert total == len(coh) # disjoint + covering + + +def test_split_is_deterministic(): + a = _coh(6).split(by="study_id", train=0.5, val=0.25, test=0.25, seed=7) + b = _coh(6).split(by="study_id", train=0.5, val=0.25, test=0.25, seed=7) + assert {k: sorted(v.to_df()["uid"]) for k, v in a.items()} == {k: sorted(v.to_df()["uid"]) for k, v in b.items()} + + +def test_fold_carries_split_manifest(tmp_path): + folds = _coh(6).split(by="study_id", train=0.5, val=0.25, test=0.25, seed=0) + sm = folds["train"].split_manifest + assert isinstance(sm, SplitManifest) + assert sm.write(tmp_path / "split.json").exists() diff --git a/tests/test_cohort_subset.py b/tests/test_cohort_subset.py new file mode 100644 index 0000000..c045377 --- /dev/null +++ b/tests/test_cohort_subset.py @@ -0,0 +1,57 @@ +import pandas as pd +import pytest + +from sddb.cohort import SpatialDataCohort + + +def _coh(uids, tissue): + return SpatialDataCohort(pd.DataFrame({"uid": uids, "tissue": tissue, "zarr_url": ["x"] * len(uids)})) + + +def test_filter_returns_new_cohort_and_leaves_source_unchanged(): + coh = _coh(["a", "b", "c"], ["breast", "lung", "breast"]) + sub = coh.filter(tissue="breast") + assert sorted(sub.to_df()["uid"]) == ["a", "c"] + assert len(coh) == 3 # unchanged + + +def test_filter_list_value_is_isin(): + coh = _coh(["a", "b", "c"], ["breast", "lung", "brain"]) + assert sorted(coh.filter(tissue=["breast", "brain"]).to_df()["uid"]) == ["a", "c"] + + +def test_filter_unknown_column_raises(): + with pytest.raises(ValueError, match="nope"): + _coh(["a"], ["breast"]).filter(nope="x") + + +def test_where_mask(): + coh = _coh(["a", "b"], ["breast", "lung"]) + out = coh.where(coh.members["tissue"] == "lung") + assert list(out.to_df()["uid"]) == ["b"] + + +def test_union_dedupes_by_uid(): + a = _coh(["a", "b"], ["breast", "lung"]) + b = _coh(["b", "c"], ["lung", "brain"]) + u = a | b + assert sorted(u.to_df()["uid"]) == ["a", "b", "c"] # b not doubled + assert len(u) == 3 + + +def test_difference(): + a = _coh(["a", "b", "c"], ["t"] * 3) + b = _coh(["b"], ["t"]) + assert sorted((a - b).to_df()["uid"]) == ["a", "c"] + + +def test_empty_filter_is_empty_cohort_not_error(): + coh = _coh(["a"], ["breast"]) + out = coh.filter(tissue="nonexistent") + assert len(out) == 0 + out.summary() # must not raise on an empty cohort + + +def test_slice_preserves_filter(): + coh = SpatialDataCohort(pd.DataFrame({"uid": ["a", "b", "c"], "zarr_url": ["x"] * 3}), filter={"tissue": "breast"}) + assert coh[:2]._filter == {"tissue": "breast"} diff --git a/tests/test_cohort_verify.py b/tests/test_cohort_verify.py new file mode 100644 index 0000000..fd52f42 --- /dev/null +++ b/tests/test_cohort_verify.py @@ -0,0 +1,47 @@ +import warnings + +import pandas as pd +import pytest + +from sddb import cohort as cohort_mod +from sddb.cohort import SpatialDataCohort +from sddb.manifest import CohortDriftError + + +class _FS: + def __init__(self, present): + self._present = present + + def exists(self, root): + return root in self._present + + +def test_verify_passes_when_all_present(monkeypatch): + present = {"root1", "root2"} + monkeypatch.setattr(cohort_mod, "_open_fs", lambda url: (_FS(present), url.split("//")[-1])) + coh = SpatialDataCohort(pd.DataFrame({"uid": ["a", "b"], "zarr_url": ["s3://root1", "s3://root2"]})) + coh.verify() # must not raise + + +def test_verify_raises_naming_missing(monkeypatch): + monkeypatch.setattr(cohort_mod, "_open_fs", lambda url: (_FS({"root1"}), url.split("//")[-1])) + coh = SpatialDataCohort(pd.DataFrame({"uid": ["a", "b"], "zarr_url": ["s3://root1", "s3://gone"]})) + with pytest.raises(CohortDriftError) as ei: + coh.verify() + assert "b" in ei.value.drifted + + +def test_check_updates_detects_superseded(): + # pinned uid has version 0000; catalog's current row for the same stem is 0002 + stem = "s" * 16 + pinned = SpatialDataCohort(pd.DataFrame({"uid": [stem + "0000"], "zarr_url": ["s3://x"]})) + + class _Cat: + def to_df(self_): + return pd.DataFrame({"uid": [stem + "0002"], "zarr_url": ["s3://y"]}) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + superseded = pinned.check_updates(_Cat()) + assert superseded == {stem + "0000": stem + "0002"} + assert any("supersed" in str(x.message).lower() for x in w) diff --git a/tests/test_concat.py b/tests/test_concat.py index e932f33..27b01ef 100644 --- a/tests/test_concat.py +++ b/tests/test_concat.py @@ -20,7 +20,7 @@ def test_concat_reads_only_tables(tmp_path): members = pd.DataFrame( {"uid": ["a", "b"], "zarr_url": [a, b], "fingerprint": ["x", "y"], "default_table": ["table", "table"]} ) - adata = concat_tables(members, verify=False) + adata = concat_tables(members) assert adata.n_obs == 10 assert set(adata.obs["uid"]) == {"a", "b"} @@ -29,26 +29,22 @@ def test_concat_missing_table_names_member(tmp_path): a = _store(tmp_path / "a.zarr", 4) members = pd.DataFrame({"uid": ["a"], "zarr_url": [a], "fingerprint": ["x"], "default_table": ["nope"]}) with pytest.raises(KeyError) as ei: - concat_tables(members, verify=False) + concat_tables(members) assert "'a'" in str(ei.value) def test_concat_without_default_table_column(tmp_path): a = _store(tmp_path / "a.zarr", 2) members = pd.DataFrame({"uid": ["a"], "zarr_url": [a]}) - assert concat_tables(members, verify=False).n_obs == 2 + assert concat_tables(members).n_obs == 2 def test_cohort_delegators(tmp_path): from sddb.cohort import SpatialDataCohort - from sddb.dataset import Source - from sddb.freeze import FrozenCohort a = _store(tmp_path / "a.zarr", 3) df = pd.DataFrame({"uid": ["a"], "zarr_url": [a], "fingerprint": ["x"]}) assert SpatialDataCohort(df).to_anndata().n_obs == 3 - fc = FrozenCohort("1", "now", Source(), {}, "sha256:h", df) - assert fc.to_anndata(verify=False).n_obs == 3 @pytest.mark.network diff --git a/tests/test_concat_study_id.py b/tests/test_concat_study_id.py new file mode 100644 index 0000000..a5ff5ec --- /dev/null +++ b/tests/test_concat_study_id.py @@ -0,0 +1,45 @@ +import anndata as ad +import numpy as np +import pandas as pd + +from sddb import concat + + +def _fake_group(adatas): + """Patch concat so each member returns a prebuilt AnnData instead of reading zarr.""" + it = iter(adatas) + + def read_zarr(_group): + return next(it) + + return read_zarr + + +def test_to_anndata_adds_uid_and_study_id(monkeypatch): + a = ad.AnnData(np.ones((2, 2))) + b = ad.AnnData(np.ones((3, 2))) + monkeypatch.setattr(concat.ad, "read_zarr", _fake_group([a, b])) + monkeypatch.setattr(concat, "_open_group", lambda url: url) + members = pd.DataFrame( + {"uid": ["u1", "u2"], "zarr_url": ["s3://b/1.zarr", "s3://b/2.zarr"], "study_id": ["s_a", "s_b"]} + ) + out = concat.concat_tables(members) + assert "uid" in out.obs + assert "study_id" in out.obs + assert set(out.obs["study_id"].dropna()) == {"s_a", "s_b"} + + +def test_to_anndata_without_study_id_column(monkeypatch): + a = ad.AnnData(np.ones((2, 2))) + monkeypatch.setattr(concat.ad, "read_zarr", _fake_group([a])) + monkeypatch.setattr(concat, "_open_group", lambda url: url) + members = pd.DataFrame({"uid": ["u1"], "zarr_url": ["s3://b/1.zarr"]}) + out = concat.concat_tables(members) # must not raise + assert "study_id" in out.obs + assert out.obs["study_id"].isna().all() + + +def test_verify_param_removed(): + import inspect + + assert "verify" not in inspect.signature(concat.concat_tables).parameters diff --git a/tests/test_download_fallback.py b/tests/test_download_fallback.py new file mode 100644 index 0000000..2e0bf88 --- /dev/null +++ b/tests/test_download_fallback.py @@ -0,0 +1,86 @@ +import warnings + +import pandas as pd +import pytest + +from sddb import manifest as m +from sddb.cohort import SpatialDataCohort +from sddb.dataset import Source + + +def _coh(uid, url, source_url="cat.parquet"): + df = pd.DataFrame({"uid": [uid], "zarr_url": [url]}) + return SpatialDataCohort(df, source=Source(source_url)) + + +def test_strict_raises_on_missing_store(monkeypatch, tmp_path): + monkeypatch.setattr(m, "_store_exists", lambda url: False) + coh = _coh("u1", "s3://b/missing.zarr") + with pytest.raises(m.CohortDriftError, match="u1"): + coh.download(tmp_path, strict=True) + + +def test_fallback_resolves_latest_with_warning(monkeypatch, tmp_path): + stem = "s" * 16 + present = {"s3://b/latest.zarr"} + monkeypatch.setattr(m, "_store_exists", lambda url: url in present) + + def _latest_for(source, stem_arg): + assert stem_arg == stem + return stem + "0002", "s3://b/latest.zarr" + + monkeypatch.setattr(m, "_latest_for_stem", _latest_for) + captured = {} + monkeypatch.setattr(m, "_fetch_all", lambda pairs, root, cv, workers: captured.setdefault("pairs", pairs)) + coh = _coh(stem + "0000", "s3://b/missing.zarr") + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + coh.download(tmp_path) + assert captured["pairs"] == [(stem + "0000", "s3://b/latest.zarr")] + assert any("fall" in str(x.message).lower() or "latest" in str(x.message).lower() for x in w) + + +def test_missing_store_no_fallback_raises_nonstrict(monkeypatch, tmp_path): + monkeypatch.setattr(m, "_store_exists", lambda url: False) + coh = SpatialDataCohort(pd.DataFrame({"uid": ["u1"], "zarr_url": [str(tmp_path / "nope.zarr")]})) + with pytest.raises(m.CohortDriftError, match="u1"): + coh.download(tmp_path / "out") + + +def test_latest_for_stem_no_url_returns_none(): + assert m._latest_for_stem(Source(), "s" * 16) is None + + +def test_transient_liveness_error_does_not_fall_back(monkeypatch, tmp_path): + def boom(url): + raise OSError("connection reset") + + monkeypatch.setattr(m, "_store_exists", boom) + monkeypatch.setattr(m, "_latest_for_stem", lambda s, st: pytest.fail("must not fall back")) + coh = _coh("u1", "s3://b/x.zarr") + with pytest.raises(RuntimeError, match="could not determine liveness of u1"): + coh.download(tmp_path) + + +def test_complete_member_skips_liveness_on_resume(monkeypatch, tmp_path): + from tests._fixtures import make_tiny_sdata_zarr + + z = make_tiny_sdata_zarr(tmp_path) + coh = _coh("u1", str(z)) + coh.download(tmp_path / "out") + monkeypatch.setattr(m, "_store_exists", lambda url: False) + monkeypatch.setattr(m, "_latest_for_stem", lambda s, st: pytest.fail("must not fall back")) + res = coh.download(tmp_path / "out", strict=True) + assert [e.status for e in res.entries] == ["complete"] + + +def test_real_missing_local_store_strict_raises_drift(tmp_path): + coh = _coh("u1", str(tmp_path / "missing.zarr"), source_url=None) + with pytest.raises(m.CohortDriftError, match="u1"): + coh.download(tmp_path / "out", strict=True) + + +def test_real_missing_local_store_nonstrict_no_fallback_raises_drift(tmp_path): + coh = SpatialDataCohort(pd.DataFrame({"uid": ["u1"], "zarr_url": [str(tmp_path / "missing.zarr")]})) + with pytest.raises(m.CohortDriftError, match="u1"): + coh.download(tmp_path / "out", strict=False) diff --git a/tests/test_facets_study_doi.py b/tests/test_facets_study_doi.py new file mode 100644 index 0000000..d125551 --- /dev/null +++ b/tests/test_facets_study_doi.py @@ -0,0 +1,29 @@ +import pandas as pd + +from sddb import _catalog_schema as schema +from sddb.catalog import Catalog + + +def test_study_id_and_doi_are_facets(): + assert "study_id" in schema.FACETS + assert "doi" in schema.FACETS + + +def test_query_by_study_id(tmp_path): + df = pd.DataFrame( + { + "uid": ["u1", "u2"], + "study_id": ["oliveira_2025", "other_2024"], + "doi": ["10.1/x", "10.2/y"], + "technology": ["Visium HD"] * 2, + "assay": ["spatial"] * 2, + "organism": ["human"] * 2, + "validation_status": ["pass"] * 2, + "zarr_url": ["s3://b/1.zarr", "s3://b/2.zarr"], + } + ) + p = tmp_path / "catalog.parquet" + df.to_parquet(p) + cat = Catalog.from_file(p) + assert list(cat.query(study_id="oliveira_2025").to_df()["uid"]) == ["u1"] + assert list(cat.query(doi="10.2/y").to_df()["uid"]) == ["u2"] diff --git a/tests/test_freeze.py b/tests/test_freeze.py deleted file mode 100644 index c533176..0000000 --- a/tests/test_freeze.py +++ /dev/null @@ -1,95 +0,0 @@ -from __future__ import annotations - -import pandas as pd -from tests._fixtures import make_tiny_sdata_zarr - -from sddb.cohort import SpatialDataCohort -from sddb.dataset import Source -from sddb.freeze import FrozenCohort, cohort_hash - - -def _members(): - return pd.DataFrame( - { - "uid": ["u2", "u1"], - "lamin_version": [None, None], - "fingerprint": ["sha256:b", "sha256:a"], - "study_id": ["S2", "S1"], - "zarr_url": ["z2", "z1"], - } - ) - - -def test_hash_is_order_independent_and_excludes_timestamp(): - m = _members() - h1 = cohort_hash({"organism": "human"}, m) - h2 = cohort_hash({"organism": "human"}, m.iloc[::-1].reset_index(drop=True)) - assert h1 == h2 - assert h1.startswith("sha256:") - - -def test_hash_changes_with_membership(): - m = _members() - assert cohort_hash({}, m) != cohort_hash({}, m.iloc[:1]) - - -def test_write_load_roundtrip(tmp_path): - fc = FrozenCohort( - "1", - "2026-10-07T00:00:00Z", - Source(url="cat", version="v"), - {"organism": "human"}, - cohort_hash({"organism": "human"}, _members()), - _members(), - ) - p = fc.write(tmp_path / "cohort.json") - back = FrozenCohort.load(p) - assert back.hash == fc.hash - assert list(back.members["uid"]) == list(fc.members["uid"]) - assert back.filter == {"organism": "human"} - - -def test_write_load_sidecar_roundtrip(tmp_path): - n = 250 - m = pd.DataFrame({"uid": [f"u{i}" for i in range(n)], "lamin_version": [None] * n, "fingerprint": ["sha256:x"] * n}) - fc = FrozenCohort("1", "t", Source(url="c", version="v"), {}, cohort_hash({}, m), m) - p = fc.write(tmp_path / "c.json") - assert (tmp_path / "c.members.parquet").exists() - back = FrozenCohort.load(p) - assert len(back.members) == n - assert back.hash == fc.hash - - -def test_query_records_deterministic_filter(tmp_path): - import json - - import pytest - from tests._fixtures import write_fixture_catalog - - from sddb.catalog import Catalog - - p = write_fixture_catalog(tmp_path / "cat.parquet") - cat = Catalog(p.as_uri(), cache_dir=tmp_path / "cache") - a = cat.query(organism={"mouse", "human"})._filter - b = cat.query(organism=("human", "mouse"))._filter - assert a == b == {"validation": "pass", "organism": ["human", "mouse"]} - json.dumps(a) - assert cat.query(validation=None)._filter == {} - with pytest.warns(UserWarning, match="nothing matched"): - assert cat.search("zzzzqqq")._filter == {"search": "zzzzqqq"} - - -def test_slice_keeps_filter(): - df = pd.DataFrame({"uid": ["u1", "u2", "u3"], "zarr_url": ["a", "b", "c"]}) - parent = SpatialDataCohort(df, filter={"organism": "human"}) - assert parent[:2]._filter == {"organism": "human"} - - -def test_freeze_local_store(tmp_path): - z = make_tiny_sdata_zarr(tmp_path) - df = pd.DataFrame({"uid": ["u1"], "zarr_url": [str(z)], "study_id": ["S1"]}) - fc = SpatialDataCohort(df, filter={"organism": "human"}).freeze(workers=1) - assert fc.members.loc[0, "fingerprint"].startswith("sha256:") - assert fc.hash == cohort_hash({"organism": "human"}, fc.members) - again = SpatialDataCohort(df, filter={"organism": "human"}).freeze(workers=1) - assert again.hash == fc.hash diff --git a/tests/test_freeze_verify.py b/tests/test_freeze_verify.py deleted file mode 100644 index 66e753a..0000000 --- a/tests/test_freeze_verify.py +++ /dev/null @@ -1,34 +0,0 @@ -from __future__ import annotations - -import shutil - -import anndata as ad -import numpy as np -import pandas as pd -import pytest -from spatialdata import SpatialData -from spatialdata.models import Image2DModel - -from sddb.cohort import SpatialDataCohort -from sddb.freeze import CohortDriftError, verify_members - - -def _write(path, n_obs): - img = Image2DModel.parse(np.zeros((3, 8, 8), "uint8"), dims=("c", "y", "x")) - SpatialData( - images={"img": img}, - tables={"table": ad.AnnData(X=np.ones((n_obs, 3), "float32"))}, - ).write(path) - return path - - -def test_verify_passes_then_detects_drift(tmp_path): - z = _write(tmp_path / "o.zarr", 4) - df = pd.DataFrame({"uid": ["o"], "zarr_url": [str(z)], "study_id": ["S"]}) - fc = SpatialDataCohort(df, filter={}).freeze() - verify_members(fc.members) # no raise: unchanged - shutil.rmtree(z) - _write(z, 9) # republish different data at the same path - with pytest.raises(CohortDriftError) as ei: - verify_members(fc.members) - assert "o" in ei.value.drifted diff --git a/tests/test_genes.py b/tests/test_genes.py index 2b2386c..3b25c00 100644 --- a/tests/test_genes.py +++ b/tests/test_genes.py @@ -74,8 +74,13 @@ def _stub_genes(tmp_path): ] df = pd.DataFrame(rows, columns=["symbol", "feature_id", "uid", "total_counts", "fraction_obs_detected"]) df = df.astype( - {"symbol": "string", "feature_id": "string", "uid": "string", "total_counts": "Float64", - "fraction_obs_detected": "Float64"} + { + "symbol": "string", + "feature_id": "string", + "uid": "string", + "total_counts": "Float64", + "fraction_obs_detected": "Float64", + } ) assert df["feature_id"].isna().sum() == 1 df.to_parquet(tmp_path / "gene_index.parquet") @@ -118,7 +123,8 @@ def test_legacy_symbol_only_no_crash(tmp_path): def test_backcompat_single_shape(tmp_path): g = _stub_genes(tmp_path) a, b = g.where_expressed("EPCAM").to_df(), g.datasets_with("EPCAM").to_df() - assert list(a.columns) == list(b.columns) and list(a["uid"]) == list(b["uid"]) + assert list(a.columns) == list(b.columns) + assert list(a["uid"]) == list(b["uid"]) # --------------------------------------------------------------------------- @@ -166,8 +172,13 @@ def _canonical_gene_index(): ] df = pd.DataFrame(rows, columns=["symbol", "feature_id", "uid", "total_counts", "fraction_obs_detected"]) return df.astype( - {"symbol": "string", "feature_id": "string", "uid": "string", "total_counts": "Float64", - "fraction_obs_detected": "Float64"} + { + "symbol": "string", + "feature_id": "string", + "uid": "string", + "total_counts": "Float64", + "fraction_obs_detected": "Float64", + } ) diff --git a/tests/test_load_remote_fallback.py b/tests/test_load_remote_fallback.py index 836a6bc..47d2210 100644 --- a/tests/test_load_remote_fallback.py +++ b/tests/test_load_remote_fallback.py @@ -12,7 +12,7 @@ def test_load_remote_lazy_gives_helpful_error(monkeypatch): """Test that Dataset.load(lazy=True) on a remote URL gives a helpful error mentioning .download() and lazy=False.""" # Monkeypatch remote.open_sdata to raise NotImplementedError like upstream spatialdata would - def mock_open_sdata(zarr_url: str, *, lazy: bool = True, cache_dir=None): + def mock_open_sdata(zarr_url: str, *, lazy: bool = True, cache_dir=None, elements=None): if lazy and zarr_url.startswith(("s3://", "http://", "https://")): raise NotImplementedError("Remote lazy/partial open is not supported") raise AssertionError(f"Unexpected call to mock_open_sdata: lazy={lazy}, zarr_url={zarr_url}") diff --git a/tests/test_localization.py b/tests/test_localization.py index aa56762..b895357 100644 --- a/tests/test_localization.py +++ b/tests/test_localization.py @@ -10,7 +10,7 @@ from shapely.geometry import Point from spatialdata.models import Image2DModel, ShapesModel, TableModel -from sddb import adapter, remote +from sddb import remote def _shapes(n): @@ -162,109 +162,6 @@ def test_region_closure_autoincludes(vhd_store, tmp_path): assert out.tables["square_008um"].obs["region"].cat.categories.tolist() == ["s_square_008um"] -def test_resolve_elements_one_each(vhd_store): - roles = [ - {"role": "image", "select": "full_image"}, - {"role": "table", "select": "square_008um"}, - {"role": "shapes", "select": "square_008um"}, - ] - got = adapter._resolve_elements(str(vhd_store), roles) - assert sorted(got) == ["images/s_full_image", "shapes/s_square_008um", "tables/square_008um"] - - -def test_resolve_elements_suffix_rejects_sibling(tmp_path): - # suffix "square_016um" must NOT also match "square_016um_for_vitessce" - store = tmp_path / "s.zarr" - sd.SpatialData( - shapes={ - "a_square_016um": _shapes(5), - "a_square_016um_for_vitessce": _shapes(5), - } - ).write(store) - got = adapter._resolve_elements(str(store), [{"role": "shapes", "select": "square_016um"}]) - assert got == ["shapes/a_square_016um"] - - -def test_resolve_elements_ambiguous_raises(vhd_store): - with pytest.raises(ValueError, match="matched 3"): - adapter._resolve_elements(str(vhd_store), [{"role": "table", "select": "um", "match": "substring"}]) - - -def test_resolve_elements_missing_raises(vhd_store): - with pytest.raises(ValueError, match="no element"): - adapter._resolve_elements(str(vhd_store), [{"role": "image", "select": "nope"}]) - - -class _FakeAdapter: - name, version, config = "fake", "0", {} - - def requirements(self): - return {"elements": [ - {"role": "image", "select": "full_image"}, - {"role": "table", "select": "square_008um"}, - {"role": "shapes", "select": "square_008um"}, - ]} - - def validate_meta(self, meta): - return [] - - def validate(self, sdata): - return [] if set(sdata.tables) == {"square_008um"} else ["wrong tables"] - - def build(self, sdata): - return sorted(sdata.tables) - - -def _members_df(vhd_store): - return pd.DataFrame([{ - "uid": "u1", "zarr_url": str(vhd_store), "lamin_version": "1", "fingerprint": "fp", - }]) - - -def test_build_localizes_from_requirements(vhd_store, tmp_path, monkeypatch): - monkeypatch.setenv("SDDB_CACHE_DIR", str(tmp_path / "cache")) - from sddb import adapter - view = adapter.build(_members_df(vhd_store), _FakeAdapter(), verify=False) - uid, thunk = view[0] - assert thunk() == ["square_008um"] - - -def test_build_provenance_records_elements(vhd_store, tmp_path, monkeypatch): - monkeypatch.setenv("SDDB_CACHE_DIR", str(tmp_path / "cache")) - out = tmp_path / "prov" - adapter.build(_members_df(vhd_store), _FakeAdapter(), materialize=True, out=out, verify=False) - prov = json.loads((out / "provenance.json").read_text()) - assert prov["requirements"]["elements"][0]["role"] == "image" - member = prov["members"][0] - assert member["elements"] == ["images/s_full_image", "shapes/s_square_008um", "tables/square_008um"] - - -class _Fake002Adapter(_FakeAdapter): - def requirements(self): - return {"elements": [ - {"role": "image", "select": "full_image"}, - {"role": "table", "select": "square_002um"}, - {"role": "shapes", "select": "square_002um"}, - ]} - - def validate(self, sdata): - return [] - - -def test_provenance_differs_by_element_selector(vhd_store, tmp_path, monkeypatch): - monkeypatch.setenv("SDDB_CACHE_DIR", str(tmp_path / "cache")) - members = _members_df(vhd_store) - out_a, out_b = tmp_path / "prov_a", tmp_path / "prov_b" - adapter.build(members, _FakeAdapter(), materialize=True, out=out_a, verify=False) - adapter.build(members, _Fake002Adapter(), materialize=True, out=out_b, verify=False) - a = json.loads((out_a / "provenance.json").read_text()) - b = json.loads((out_b / "provenance.json").read_text()) - assert a["members"][0]["elements"] == ["images/s_full_image", "shapes/s_square_008um", "tables/square_008um"] - assert b["members"][0]["elements"] == ["images/s_full_image", "shapes/s_square_002um", "tables/square_002um"] - assert a["members"][0]["elements"] != b["members"][0]["elements"] - assert a["requirements"]["elements"] != b["requirements"]["elements"] - - def test_close_region_tolerates_corrupt_table_metadata(tmp_path): import fsspec @@ -274,51 +171,6 @@ def test_close_region_tolerates_corrupt_table_metadata(tmp_path): assert remote._close_region(fs, str(tmp_path), ["tables/t"]) == ["tables/t"] -def test_open_member_verifies_then_opens(vhd_store, tmp_path, monkeypatch): - monkeypatch.setenv("SDDB_CACHE_DIR", str(tmp_path / "cache")) - import sddb.freeze as fz - import sddb.remote as rem - from sddb.dataset import Source - from sddb.freeze import FrozenCohort - - members = _members_df(vhd_store) - fc = FrozenCohort("1", "2026-10-08T00:00:00Z", Source(url=None, version=None), {}, "sha256:x", members) - - calls = [] - monkeypatch.setattr(fz, "verify_members", lambda m, **k: calls.append("verify")) - _orig = rem.open_sdata - - def spy(*a, **k): - calls.append("open") - return _orig(*a, **k) - - monkeypatch.setattr(rem, "open_sdata", spy) - - sdata = fc.open_member("u1", elements=["tables/square_008um"]) - assert calls == ["verify", "open"] - assert set(sdata.tables) == {"square_008um"} - - -@pytest.mark.network -def test_localize_oliveira_excludes_2um(tmp_path): - # Real Oliveira Visium HD store: localizing 8um + H&E must not download the 8.1M-row 2um table. - import sddb - - cat = sddb.Catalog() - url = cat.to_df().set_index("uid").loc["4OgPqTscrYc4veqE0000", "zarr_url"] - roles = [ - {"role": "image", "select": "full_image"}, - {"role": "table", "select": "square_008um"}, - {"role": "shapes", "select": "square_008um"}, - ] - paths = adapter._resolve_elements(url, roles) - sdata = remote.open_sdata(url, lazy=False, elements=paths, cache_dir=tmp_path) - assert "square_002um" not in sdata.tables - assert any("square_008um" in t for t in sdata.tables) - cached = next(tmp_path.glob("*.zarr")) - assert not list(cached.glob("tables/*002um*")) # 2um subtree never copied - - def test_same_basename_stores_get_distinct_cache_dirs(tmp_path): cache = tmp_path / "cache" got = {} @@ -338,60 +190,3 @@ def boom(*a, **k): monkeypatch.setattr(sd, "read_zarr", boom) with pytest.raises(FileNotFoundError, match="Cannot open SpatialData zarr.*bad read"): remote.open_sdata(str(vhd_store), lazy=False, elements=KEEP, cache_dir=tmp_path / "c") - - -def test_resolve_elements_role_canonicalization(vhd_store): - url = str(vhd_store) - for role, want in [ - ("shape", "shapes/s_square_008um"), - ("shapes", "shapes/s_square_008um"), - ("image", "images/s_full_image"), - ("images", "images/s_full_image"), - ("table", "tables/square_008um"), - ]: - select = "full_image" if "image" in role else "square_008um" - assert adapter._resolve_elements(url, [{"role": role, "select": select}]) == [want] - - -class _NoRequirementsAdapter: - name, version, config = "noreq", "0", {} - - def validate_meta(self, meta): - return [] - - def validate(self, sdata): - return [] - - def build(self, sdata): - return sorted(sdata.tables) - - -def test_open_adapter_without_requirements(vhd_store, tmp_path, monkeypatch): - monkeypatch.setenv("SDDB_CACHE_DIR", str(tmp_path / "cache")) - sdata = adapter._open(str(vhd_store), _NoRequirementsAdapter()) - assert "square_002um" in sdata.tables # whole store - - -def test_materialize_adapter_without_requirements(vhd_store, tmp_path, monkeypatch): - monkeypatch.setenv("SDDB_CACHE_DIR", str(tmp_path / "cache")) - out = tmp_path / "prov" - adapter.build(_members_df(vhd_store), _NoRequirementsAdapter(), materialize=True, out=out, verify=False) - prov = json.loads((out / "provenance.json").read_text()) - assert prov["requirements"] == {} - assert prov["members"][0]["elements"] is None - - -class _FakeTableOnlyAdapter(_FakeAdapter): - def requirements(self): - return {"elements": [{"role": "table", "select": "square_008um"}]} - - def validate(self, sdata): - return [] - - -def test_provenance_reflects_region_closure(vhd_store, tmp_path, monkeypatch): - monkeypatch.setenv("SDDB_CACHE_DIR", str(tmp_path / "cache")) - out = tmp_path / "prov" - adapter.build(_members_df(vhd_store), _FakeTableOnlyAdapter(), materialize=True, out=out, verify=False) - prov = json.loads((out / "provenance.json").read_text()) - assert prov["members"][0]["elements"] == ["shapes/s_square_008um", "tables/square_008um"] diff --git a/tests/test_manifest.py b/tests/test_manifest.py index 9221ba5..bf0ca8e 100644 --- a/tests/test_manifest.py +++ b/tests/test_manifest.py @@ -39,12 +39,24 @@ def test_download_and_resume(tmp_path, cohort, monkeypatch): assert calls[2:] == ["uid0002.zarr.part"] -def test_failed_entry_refetched(tmp_path, cohort): - df = cohort.to_df() - df.loc[1, "zarr_url"] = str(tmp_path / "missing.zarr") - m = SpatialDataCohort(df).download(tmp_path / "out") +def test_failed_entry_refetched(tmp_path, cohort, monkeypatch): + # store exists (passes liveness) but the copy fails for one member -> recorded "failed", refetched later + real = manifest._copy_store + bad = str(cohort.to_df()["uid"].iloc[1]) + ".zarr.part" + + def flaky(u, d): + if d.name == bad: + raise OSError("transient copy failure") + return real(u, d) + + monkeypatch.setattr(manifest, "_copy_store", flaky) + dest = tmp_path / "out" + m = cohort.download(dest) assert [e.status for e in m.entries] == ["complete", "failed"] - assert not list((tmp_path / "out").glob("*.part")) + assert not list(dest.glob("*.part")) + monkeypatch.setattr(manifest, "_copy_store", real) + m2 = cohort.download(dest) + assert [e.status for e in m2.entries] == ["complete", "complete"] def test_version_mismatch_and_allow(tmp_path, cohort): @@ -95,11 +107,14 @@ def flaky(url, dest): def test_plan_sizes_local_store(tmp_path): import pandas as pd + from tests._fixtures import make_tiny_sdata_zarr + from sddb.cohort import SpatialDataCohort from sddb.manifest import plan_sizes - from tests._fixtures import make_tiny_sdata_zarr store = make_tiny_sdata_zarr(tmp_path) res = SpatialDataCohort(pd.DataFrame({"uid": ["uidX"], "zarr_url": [store.as_uri()]})) sizes = plan_sizes(res) - assert len(sizes) == 1 and sizes[0][0] == "uidX" and sizes[0][1] > 0 + assert len(sizes) == 1 + assert sizes[0][0] == "uidX" + assert sizes[0][1] > 0 diff --git a/tests/test_mcp_tools.py b/tests/test_mcp_tools.py index ae6c320..b5b49eb 100644 --- a/tests/test_mcp_tools.py +++ b/tests/test_mcp_tools.py @@ -126,6 +126,7 @@ def test_query_tool_expressing_min_fraction(cat_url): def test__query_helper_filters(tmp_path): from tests._fixtures import write_fixture_catalog + from sddb.mcp.tools import _query cat = write_fixture_catalog(tmp_path / "catalog.parquet") @@ -134,33 +135,38 @@ def test__query_helper_filters(tmp_path): def test_facets_tool(tmp_path): + import pytest from tests._fixtures import write_fixture_catalog + from sddb.mcp.tools import facets_tool - import pytest cat = write_fixture_catalog(tmp_path / "catalog.parquet") cols = facets_tool(catalog_url=cat.as_uri()) - assert "organism" in cols and "technology" in cols + assert "organism" in cols + assert "technology" in cols assert facets_tool("organism", catalog_url=cat.as_uri()) == ["human", "mouse"] with pytest.raises(ValueError, match="unknown facet"): facets_tool("not_a_facet", catalog_url=cat.as_uri()) def test_cite_tool(tmp_path): - from tests._fixtures import write_fixture_catalog, make_fixture_citations_bib + from tests._fixtures import make_fixture_citations_bib, write_fixture_catalog + from sddb.mcp.tools import cite_tool cat = write_fixture_catalog(tmp_path / "catalog.parquet") (tmp_path / "citations.bib").write_text(make_fixture_citations_bib(), encoding="utf-8") out = cite_tool(catalog_url=cat.as_uri(), organism="human") assert out["n"] >= 1 - assert "@" in out["bibtex"] and "Smith2023" in out["bibtex"] + assert "@" in out["bibtex"] + assert "Smith2023" in out["bibtex"] empty = cite_tool(catalog_url=cat.as_uri(), organism="nonexistent_xyz") assert empty == {"bibtex": "", "n": 0} def test_download_tool_plan_no_copy(tmp_path, monkeypatch): from tests._fixtures import write_fixture_catalog + import sddb.mcp.tools as t cat = write_fixture_catalog(tmp_path / "catalog.parquet") @@ -172,14 +178,16 @@ def test_download_tool_plan_no_copy(tmp_path, monkeypatch): ) out = t.download_tool(catalog_url=cat.as_uri(), organism="human") assert called["download"] is False # plan mode copies nothing - assert out["datasets"] == 3 and out["total_bytes"] == 3000 + assert out["datasets"] == 3 + assert out["total_bytes"] == 3000 assert {p["uid"] for p in out["per_store"]} == {"uid0001", "uid0002", "uid0005"} def test_download_tool_fetch_sandbox(tmp_path, monkeypatch): from tests._fixtures import write_fixture_catalog - from sddb.manifest import Manifest, ManifestEntry + import sddb.mcp.tools as t + from sddb.manifest import Manifest, ManifestEntry cat = write_fixture_catalog(tmp_path / "catalog.parquet") sandbox = tmp_path / "dl" @@ -212,7 +220,9 @@ def tool(self, name, description=""): def deco(fn): registered[name] = fn return fn + return deco + def run(self): ... monkeypatch.setattr(server, "_server_class", lambda: FakeServer) @@ -240,6 +250,7 @@ def tool(self, name, description=""): def deco(fn): registered[name] = fn return fn + return deco monkeypatch.setattr(server, "_server_class", lambda: FakeServer) diff --git a/tests/test_removals.py b/tests/test_removals.py new file mode 100644 index 0000000..888bc06 --- /dev/null +++ b/tests/test_removals.py @@ -0,0 +1,24 @@ +import pytest + + +def test_removed_names_not_importable_from_sddb(): + import sddb + + for name in ("FrozenCohort", "PreflightReport", "TaskView"): + assert not hasattr(sddb, name), f"{name} should be removed" + + +def test_removed_adapter_functions_gone(): + import sddb.adapter as a + + for name in ("build", "materialize", "check", "TaskView", "PreflightReport", "_resolve_elements"): + assert not hasattr(a, name), f"adapter.{name} should be removed" + + +def test_freeze_module_gone(): + with pytest.raises(ModuleNotFoundError): + import sddb.freeze # noqa: F401 + + +def test_new_names_importable(): + from sddb import CohortDriftError, CohortManifest, SpatialDataCohort, TaskAdapter # noqa: F401 diff --git a/tests/test_split.py b/tests/test_split.py index df1c1ac..8fbee7a 100644 --- a/tests/test_split.py +++ b/tests/test_split.py @@ -1,12 +1,9 @@ from __future__ import annotations -from types import SimpleNamespace - import pandas as pd import pytest -from sddb.freeze import FrozenCohort -from sddb.split import SplitManifest, _largest_remainder, make_split +from sddb.split import _largest_remainder, make_split def _members(studies): @@ -60,14 +57,6 @@ def test_largest_remainder_sums_and_nonempty(n): assert all(v >= 1 for v in c.values()) -def test_write_load_roundtrip_and_cohort_split(tmp_path): - m = _members(["A", "B", "C", "D"]) - fc = SimpleNamespace(members=m, hash="sha256:abc") - s = FrozenCohort.split(fc, train=0.5, val=0.25, test=0.25, seed=3) # type: ignore[arg-type] - assert s.parent_hash == "sha256:abc" - assert SplitManifest.load(s.write(tmp_path / "split.json")) == s - - def test_missing_by_column_raises(): with pytest.raises(KeyError, match="does_not_exist"): make_split(_members(["A", "B"]), "h", by="does_not_exist", train=0.5, val=0.0, test=0.5, seed=1)