corems.molecular_id.search.spectral_library_cache

  1"""Parquet sidecar cache for parsed spectral library DataFrames."""
  2
  3from __future__ import annotations
  4
  5import json
  6from pathlib import Path
  7
  8import numpy as np
  9import pandas as pd
 10
 11# Bump when standardized cache columns or precursor policy change.
 12COREMS_LIB_CACHE_V = 1
 13CACHE_SUFFIX = ".corems-lib.parquet"
 14_META_KEY = b"corems_lib_cache"
 15
 16
 17def default_cache_path(source_path) -> Path:
 18    """Return default sidecar path next to the source library file."""
 19    return Path(source_path).with_name(Path(source_path).name + CACHE_SUFFIX)
 20
 21
 22def _source_fingerprint(source_path) -> dict:
 23    path = Path(source_path)
 24    stat = path.stat()
 25    return {
 26        "schema_version": COREMS_LIB_CACHE_V,
 27        "source_name": path.name,
 28        "source_size": stat.st_size,
 29        "source_mtime_ns": stat.st_mtime_ns,
 30    }
 31
 32
 33def is_valid(cache_path, source_path, extra_meta: dict | None = None) -> bool:
 34    """Return True if cache exists and matches source fingerprint + schema."""
 35    cache_path = Path(cache_path)
 36    source_path = Path(source_path)
 37    if not cache_path.is_file() or not source_path.is_file():
 38        return False
 39    try:
 40        import pyarrow.parquet as pq
 41
 42        meta = pq.read_schema(cache_path).metadata or {}
 43        raw = meta.get(_META_KEY)
 44        if raw is None:
 45            return False
 46        stored = json.loads(raw.decode("utf-8"))
 47    except Exception:
 48        return False
 49
 50    expected = _source_fingerprint(source_path)
 51    if extra_meta:
 52        expected.update(extra_meta)
 53    for key, value in expected.items():
 54        if stored.get(key) != value:
 55            return False
 56    return True
 57
 58
 59def pack_peaks_for_parquet(df: pd.DataFrame) -> pd.DataFrame:
 60    """Convert a peaks column of (N, 2) arrays into mz/intensity list columns."""
 61    out = df.copy()
 62    if "peaks" not in out.columns:
 63        return out
 64    mz_col = []
 65    intensity_col = []
 66    for peaks in out["peaks"]:
 67        if peaks is None or (isinstance(peaks, float) and np.isnan(peaks)):
 68            mz_col.append([])
 69            intensity_col.append([])
 70            continue
 71        arr = np.asarray(peaks, dtype=np.float64)
 72        if arr.size == 0:
 73            mz_col.append([])
 74            intensity_col.append([])
 75        elif arr.ndim == 2 and arr.shape[1] == 2:
 76            mz_col.append(arr[:, 0].tolist())
 77            intensity_col.append(arr[:, 1].tolist())
 78        else:
 79            raise ValueError("peaks entries must be array-like of shape (N, 2)")
 80    out["mz"] = mz_col
 81    out["intensity"] = intensity_col
 82    out = out.drop(columns=["peaks"])
 83    return out
 84
 85
 86def unpack_peaks_from_parquet(df: pd.DataFrame) -> pd.DataFrame:
 87    """Rebuild peaks (N, 2) arrays from mz/intensity list columns."""
 88    out = df.copy()
 89    if "mz" not in out.columns or "intensity" not in out.columns:
 90        return out
 91    peaks = []
 92    for mz, intensity in zip(out["mz"], out["intensity"]):
 93        mz_arr = np.asarray(mz if mz is not None else [], dtype=np.float64)
 94        i_arr = np.asarray(intensity if intensity is not None else [], dtype=np.float64)
 95        if mz_arr.size == 0:
 96            peaks.append(np.zeros((0, 2), dtype=np.float64))
 97        else:
 98            peaks.append(np.column_stack([mz_arr, i_arr]))
 99    out["peaks"] = peaks
100    out = out.drop(columns=["mz", "intensity"])
101    return out
102
103
104def read(cache_path) -> pd.DataFrame:
105    """Load a cached library DataFrame and restore peaks arrays."""
106    df = pd.read_parquet(cache_path)
107    return unpack_peaks_from_parquet(df)
108
109
110def write(
111    cache_path, df: pd.DataFrame, source_path, extra_meta: dict | None = None
112) -> None:
113    """Write standardized library DataFrame as a parquet sidecar."""
114    import pyarrow as pa
115    import pyarrow.parquet as pq
116
117    cache_path = Path(cache_path)
118    packed = pack_peaks_for_parquet(df)
119    table = pa.Table.from_pandas(packed, preserve_index=False)
120    meta = _source_fingerprint(source_path)
121    if extra_meta:
122        meta.update(extra_meta)
123    existing = table.schema.metadata or {}
124    existing = dict(existing)
125    existing[_META_KEY] = json.dumps(meta).encode("utf-8")
126    table = table.replace_schema_metadata(existing)
127    cache_path.parent.mkdir(parents=True, exist_ok=True)
128    pq.write_table(table, cache_path)