1"""Parquet lake: layout, atomic publication, versions and snapshots.23Layout (relative to FINDATA_LAKE):4 raw/<source>/<dataset>/<YYYY>/<archive>.zip immutable vendor archives (+ sha256 in raw_archives)5 curated/<dataset>/delta/session_date=YYYY-MM-DD/part-<id>.parquet one file per session (all symbols, sorted)6 curated/<dataset>/base/symbol=<SYM>/year=YYYY/part-<id>.parquet compacted / backfilled history7 manifests/<dataset>/v<version>.json reproducible snapshot manifests8 _tmp/ staging (same filesystem → atomic rename)9 _retired/<dataset>/... files removed from the live set (GC after 7 days)10 exports/<export_id>.<ext>1112Invariant: at any dataset version, each (instrument_id, timestamp_utc) is in exactly one live file.13Only the ingest worker writes under curated/ (file lock `.writer.lock`); the API only reads.14"""15from __future__ import annotations1617import fcntl18import hashlib19import json20import os21import shutil22import uuid23from dataclasses import dataclass24from datetime import UTC, date, datetime25from pathlib import Path26from typing import Any2728import pyarrow as pa29import pyarrow.parquet as pq30from findata_core.config import get_settings31from findata_core.db import transaction3233ROW_GROUP_INTRADAY = 65_53634ROW_GROUP_DAILY = 1_000_000353637@dataclass38class WrittenFile:39 rel_path: str40 layout: str41 symbol_min: str42 symbol_max: str43 date_min: date44 date_max: date45 row_count: int46 bytes: int47 sha256: str484950def mkdir_retry(path: Path, attempts: int = 20) -> None:51 """mkdir -p that survives EINTR (seen on a Thunderbolt APFS volume while Spotlight was indexing it)."""52 import time53 for i in range(attempts):54 try:55 path.mkdir(parents=True, exist_ok=True)56 return57 except InterruptedError:58 time.sleep(0.2 * (i + 1))59 path.mkdir(parents=True, exist_ok=True)606162class WriterLock:63 def __init__(self, lake: Path):64 self.path = lake / ".writer.lock"65 self._fh = None6667 def __enter__(self):68 self.path.parent.mkdir(parents=True, exist_ok=True)69 self._fh = open(self.path, "w")70 try:71 fcntl.flock(self._fh, fcntl.LOCK_EX | fcntl.LOCK_NB)72 except BlockingIOError as e:73 raise RuntimeError("another lake writer holds the lock") from e74 self._fh.write(f"{os.getpid()} {datetime.now(UTC).isoformat()}\n")75 self._fh.flush()76 return self7778 def __exit__(self, *exc):79 if self._fh:80 fcntl.flock(self._fh, fcntl.LOCK_UN)81 self._fh.close()828384class Lake:85 """Two physical roots: `root` (internal SSD, hot datasets, manifests, exports) and optional `cold`86 (external SSD: raw archives + every dataset not listed in FINDATA_HOT_DATASETS). Relative paths stored in87 `lake_files` are root-agnostic (`curated/<dataset>/…`); `root_for(dataset)` picks the physical root."""8889 def __init__(self, root: Path | None = None, cold: Path | None = None):90 s = get_settings()91 self.root = (root or s.lake).expanduser()92 self.cold = (cold if cold is not None else s.lake_cold)93 self.cold = self.cold.expanduser() if self.cold else None94 self.hot = s.hot_dataset_ids95 for sub in ("raw", "curated", "manifests", "_tmp", "_retired", "exports"):96 (self.root / sub).mkdir(parents=True, exist_ok=True)97 if self.cold:98 if not (self.cold / ".findata-cold-root").exists():99 raise RuntimeError(f"cold lake root {self.cold} is not mounted (marker .findata-cold-root missing)")100 for sub in ("raw", "curated", "_tmp", "_retired"):101 (self.cold / sub).mkdir(parents=True, exist_ok=True)102103 # ------------------------------------------------------------------ roots104 def root_for(self, dataset_id: str) -> Path:105 if self.cold and dataset_id not in self.hot:106 return self.cold107 return self.root108109 def roots(self) -> list[Path]:110 return [self.root] + ([self.cold] if self.cold else [])111112 def _root_of_rel(self, rel: str) -> Path:113 parts = rel.split("/")114 if len(parts) >= 2 and parts[0] in ("curated", "_retired"):115 return self.root_for(parts[1])116 if parts and parts[0] == "raw":117 return self.cold or self.root118 return self.root119120 def raw_dir(self, source_id: str, dataset_id: str) -> Path:121 p = (self.cold or self.root) / "raw" / source_id / dataset_id / str(datetime.now(UTC).year)122 mkdir_retry(p)123 return p124125 def abs(self, rel: str) -> Path:126 return self._root_of_rel(rel) / rel127128 def rel(self, path: Path) -> str:129 for r in self.roots():130 try:131 return str(path.relative_to(r))132 except ValueError:133 continue134 raise ValueError(f"{path} is not inside a lake root")135136 def resolve(self, rel: str) -> Path:137 """Live path, or its retired location (older snapshots stay readable until GC after 7 days)."""138 p = self.abs(rel)139 if p.exists():140 return p141 if rel.startswith("curated/"):142 r = self._root_of_rel(rel) / "_retired" / rel[len("curated/"):]143 if r.exists():144 return r145 return p146147 def free_gb(self, root: Path | None = None) -> float:148 st = os.statvfs(root or self.root)149 return st.f_bavail * st.f_frsize / 1e9150151 def used_bytes(self, sub: str = "curated") -> int:152 total = 0153 for r in self.roots():154 for p in (r / sub).rglob("*"):155 if p.is_file():156 total += p.stat().st_size157 return total158159 def check_budget(self, dataset_id: str | None = None) -> None:160 s = get_settings()161 root = self.root_for(dataset_id) if dataset_id else self.root162 limit = s.lake_cold_hard_min_free_gb if (self.cold and root == self.cold) else s.lake_hard_min_free_gb163 if self.free_gb(root) < limit:164 raise RuntimeError(f"lake hard stop on {root}: only {self.free_gb(root):.1f} GB free (< {limit} GB)")165166 # ------------------------------------------------------------------ writing167 def write_table(self, dataset_id: str, table: pa.Table, layout: str, partition: dict[str, str],168 sort: bool = True) -> WrittenFile:169 """Write a Parquet file into _tmp then move it atomically under curated/. Does NOT register it."""170 self.check_budget(dataset_id)171 if table.num_rows == 0:172 raise ValueError("refusing to write an empty file")173 root = self.root_for(dataset_id)174 is_bars = "symbol_as_reported" in table.column_names and "timestamp_utc" in table.column_names175 if sort and is_bars:176 table = table.sort_by([("symbol_as_reported", "ascending"), ("timestamp_utc", "ascending")])177 part_dir = "/".join(f"{k}={v}" for k, v in partition.items())178 rel = f"curated/{dataset_id}/{layout}/{part_dir}/part-{uuid.uuid4().hex[:12]}.parquet"179 final = root / rel180 tmp = root / "_tmp" / f"{uuid.uuid4().hex}.parquet"181 interval = table.column("interval")[0].as_py() if "interval" in table.column_names else "1d"182 rg = ROW_GROUP_DAILY if interval == "1d" else ROW_GROUP_INTRADAY183 dict_cols = [c for c in ("symbol_as_reported", "instrument_id", "asset_class", "interval", "currency", "adjustment_type",184 "source_id", "underlying", "call_put") if c in table.column_names]185 pq.write_table(table, tmp, compression="zstd", compression_level=3, row_group_size=rg,186 write_statistics=True, use_dictionary=dict_cols)187 with open(tmp, "rb") as f:188 os.fsync(f.fileno())189 digest = _sha256(tmp)190 mkdir_retry(final.parent)191 os.replace(tmp, final)192 import pyarrow.compute as pc193 syms = table.column("symbol_as_reported") if is_bars else table.column("underlying")194 dates = table.column("session_date") if is_bars else table.column("trade_date")195 return WrittenFile(rel_path=rel, layout=layout, symbol_min=pc.min(syms).as_py(), symbol_max=pc.max(syms).as_py(),196 date_min=pc.min(dates).as_py(), date_max=pc.max(dates).as_py(), row_count=table.num_rows,197 bytes=final.stat().st_size, sha256=digest)198199 # ------------------------------------------------------------------ publication (atomic, versioned)200 def publish(self, dataset_id: str, added: list[WrittenFile], retire_paths: list[str], ingest_run_id: str | None,201 note: str = "") -> dict[str, Any]:202 """One transaction: register added files, retire replaced files, bump the dataset version, write manifest.203 Retired files are moved to _retired/ *after* commit (readers that already opened them keep working)."""204 with transaction() as conn:205 row = conn.execute("SELECT current_version FROM datasets WHERE id=%s FOR UPDATE", (dataset_id,)).fetchone()206 if row is None:207 raise RuntimeError(f"unknown dataset {dataset_id}")208 version = int(row["current_version"]) + 1209 for p in retire_paths:210 conn.execute("UPDATE lake_files SET version_removed=%s WHERE path=%s AND version_removed IS NULL", (version, p))211 for f in added:212 conn.execute(213 "INSERT INTO lake_files(dataset_id, path, layout, symbol_min, symbol_max, date_min, date_max, row_count, bytes, sha256, version_added, ingest_run_id) "214 "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)",215 (dataset_id, f.rel_path, f.layout, f.symbol_min, f.symbol_max, f.date_min, f.date_max, f.row_count, f.bytes, f.sha256, version, ingest_run_id))216 live = conn.execute("SELECT path, sha256, row_count, bytes FROM lake_files WHERE dataset_id=%s AND version_removed IS NULL ORDER BY path",217 (dataset_id,)).fetchall()218 snapshot_id = hashlib.sha256("\n".join(f"{r['path']} {r['sha256']}" for r in live).encode()).hexdigest()219 file_count = len(live)220 row_count = sum(int(r["row_count"]) for r in live)221 nbytes = sum(int(r["bytes"]) for r in live)222 conn.execute("INSERT INTO dataset_versions(dataset_id, version, snapshot_id, file_count, row_count, bytes, ingest_run_id, note) VALUES (%s,%s,%s,%s,%s,%s,%s,%s)",223 (dataset_id, version, snapshot_id, file_count, row_count, nbytes, ingest_run_id, note))224 conn.execute("UPDATE datasets SET current_version=%s, updated_at=now() WHERE id=%s", (version, dataset_id))225 # Manifest = compact delta (added/retired) + totals; the full file list of any version is reconstructed226 # from lake_files (version_added/version_removed) with `Lake.snapshot_files()`. Writing the full list on227 # every publish produced 83 GB of JSON during the bulk load.228 manifest = {"dataset_id": dataset_id, "version": version, "snapshot_id": snapshot_id,229 "created_at": datetime.now(UTC).isoformat(), "ingest_run_id": ingest_run_id, "note": note,230 "file_count": file_count, "row_count": row_count, "bytes": nbytes,231 "added": [{"path": f.rel_path, "sha256": f.sha256, "rows": f.row_count, "bytes": f.bytes} for f in added],232 "retired": list(retire_paths)}233 mpath = self.root / "manifests" / dataset_id / f"v{version}.json"234 mkdir_retry(mpath.parent)235 tmp = mpath.with_suffix(".json.tmp")236 tmp.write_text(json.dumps(manifest, separators=(",", ":"), default=str))237 os.replace(tmp, mpath)238 for p in retire_paths:239 src = self.abs(p)240 if src.exists():241 r = self._root_of_rel(p)242 dst = r / "_retired" / (p[len("curated/"):] if p.startswith("curated/") else p)243 mkdir_retry(dst.parent)244 shutil.move(str(src), str(dst))245 return {"version": version, "snapshot_id": snapshot_id, "file_count": file_count, "row_count": row_count, "bytes": nbytes}246247 def snapshot_files(self, dataset_id: str, version: int | None = None) -> list[dict[str, Any]]:248 """Full file list (path, sha256, rows, bytes) of a dataset at `version` (default: current), from lake_files."""249 from findata_core.db import connection as _conn250 with _conn() as conn:251 if version is None:252 version = conn.execute("SELECT current_version FROM datasets WHERE id=%s", (dataset_id,)).fetchone()["current_version"]253 rows = conn.execute("SELECT path, sha256, row_count, bytes FROM lake_files WHERE dataset_id=%s AND version_added <= %s "254 "AND (version_removed IS NULL OR version_removed > %s) ORDER BY path", (dataset_id, version, version)).fetchall()255 return [{"path": r["path"], "sha256": r["sha256"], "rows": int(r["row_count"]), "bytes": int(r["bytes"])} for r in rows]256257 def gc_orphans(self, older_than_hours: int = 2) -> int:258 """Delete curated files that are not registered in lake_files (left by a killed job before publication)."""259 from findata_core.db import connection as _conn260 with _conn() as conn:261 known = {r["path"] for r in conn.execute("SELECT path FROM lake_files").fetchall()}262 cutoff = datetime.now(UTC).timestamp() - older_than_hours * 3600263 n = 0264 for r in self.roots():265 for p in (r / "curated").rglob("*.parquet"):266 rel = str(p.relative_to(r))267 if rel not in known and p.stat().st_mtime < cutoff:268 p.unlink()269 n += 1270 return n271272 def gc_retired(self, older_than_days: int = 7) -> int:273 cutoff = datetime.now(UTC).timestamp() - older_than_days * 86400274 n = 0275 for r in self.roots():276 for p in (r / "_retired").rglob("*.parquet"):277 if p.stat().st_mtime < cutoff:278 p.unlink()279 n += 1280 for p in (r / "_tmp").glob("*"):281 if p.stat().st_mtime < datetime.now(UTC).timestamp() - 86400:282 p.unlink()283 for p in (r / "raw").rglob("*.part"): # interrupted downloads284 if p.stat().st_mtime < datetime.now(UTC).timestamp() - 6 * 3600:285 p.unlink()286 return n287288289def _sha256(path: Path) -> str:290 h = hashlib.sha256()291 with open(path, "rb") as f:292 for chunk in iter(lambda: f.read(4 * 1024 * 1024), b""):293 h.update(chunk)294 return h.hexdigest()295296297def live_files(conn, dataset_id: str, symbols: list[str] | None = None, start: date | None = None,298 end: date | None = None, version: int | None = None) -> list[dict[str, Any]]:299 """Files to scan for a query (planner). `version=None` → current live set."""300 sql = "SELECT path, layout, symbol_min, symbol_max, date_min, date_max, row_count FROM lake_files WHERE dataset_id=%s"301 params: list[Any] = [dataset_id]302 if version is None:303 sql += " AND version_removed IS NULL"304 else:305 sql += " AND version_added <= %s AND (version_removed IS NULL OR version_removed > %s)"306 params += [version, version]307 if start is not None:308 sql += " AND date_max >= %s"309 params.append(start)310 if end is not None:311 sql += " AND date_min <= %s"312 params.append(end)313 rows = conn.execute(sql, params).fetchall()314 if symbols:315 rows = [r for r in rows if any(r["symbol_min"] <= s <= r["symbol_max"] for s in symbols)]316 return rows317