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 prev = conn.execute("SELECT snapshot_id, file_count, row_count, bytes FROM dataset_versions WHERE dataset_id=%s AND version=%s",210 (dataset_id, version - 1)).fetchone() or {"snapshot_id": "", "file_count": 0, "row_count": 0, "bytes": 0}211 retired_rows = []212 for p in retire_paths:213 r = conn.execute("UPDATE lake_files SET version_removed=%s WHERE path=%s AND version_removed IS NULL RETURNING path, row_count, bytes",214 (version, p)).fetchone()215 if r:216 retired_rows.append(r)217 for f in added:218 conn.execute(219 "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) "220 "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)",221 (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))222 # Chained snapshot id (O(delta), reproducible): H(prev_snapshot ∥ sorted added (path, sha256) ∥ sorted retired paths).223 # The full file list of any version is reconstructed from lake_files (snapshot_files()).224 h = hashlib.sha256(prev["snapshot_id"].encode())225 for f in sorted(added, key=lambda x: x.rel_path):226 h.update(f"+{f.rel_path} {f.sha256}\n".encode())227 for r in sorted(retired_rows, key=lambda x: x["path"]):228 h.update(f"-{r['path']}\n".encode())229 snapshot_id = h.hexdigest()230 file_count = int(prev["file_count"]) + len(added) - len(retired_rows)231 row_count = int(prev["row_count"]) + sum(f.row_count for f in added) - sum(int(r["row_count"]) for r in retired_rows)232 nbytes = int(prev["bytes"]) + sum(f.bytes for f in added) - sum(int(r["bytes"]) for r in retired_rows)233 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)",234 (dataset_id, version, snapshot_id, file_count, row_count, nbytes, ingest_run_id, note))235 conn.execute("UPDATE datasets SET current_version=%s, updated_at=now() WHERE id=%s", (version, dataset_id))236 # Manifest = compact delta (added/retired) + totals; the full file list of any version is reconstructed237 # from lake_files (version_added/version_removed) with `Lake.snapshot_files()`. Writing the full list on238 # every publish produced 83 GB of JSON during the bulk load.239 manifest = {"dataset_id": dataset_id, "version": version, "snapshot_id": snapshot_id,240 "created_at": datetime.now(UTC).isoformat(), "ingest_run_id": ingest_run_id, "note": note,241 "file_count": file_count, "row_count": row_count, "bytes": nbytes,242 "added": [{"path": f.rel_path, "sha256": f.sha256, "rows": f.row_count, "bytes": f.bytes} for f in added],243 "retired": list(retire_paths)}244 mpath = self.root / "manifests" / dataset_id / f"v{version}.json"245 mkdir_retry(mpath.parent)246 tmp = mpath.with_suffix(".json.tmp")247 tmp.write_text(json.dumps(manifest, separators=(",", ":"), default=str))248 os.replace(tmp, mpath)249 for p in retire_paths:250 src = self.abs(p)251 if src.exists():252 r = self._root_of_rel(p)253 dst = r / "_retired" / (p[len("curated/"):] if p.startswith("curated/") else p)254 mkdir_retry(dst.parent)255 shutil.move(str(src), str(dst))256 return {"version": version, "snapshot_id": snapshot_id, "file_count": file_count, "row_count": row_count, "bytes": nbytes}257258 def snapshot_files(self, dataset_id: str, version: int | None = None) -> list[dict[str, Any]]:259 """Full file list (path, sha256, rows, bytes) of a dataset at `version` (default: current), from lake_files."""260 from findata_core.db import connection as _conn261 with _conn() as conn:262 if version is None:263 version = conn.execute("SELECT current_version FROM datasets WHERE id=%s", (dataset_id,)).fetchone()["current_version"]264 rows = conn.execute("SELECT path, sha256, row_count, bytes FROM lake_files WHERE dataset_id=%s AND version_added <= %s "265 "AND (version_removed IS NULL OR version_removed > %s) ORDER BY path", (dataset_id, version, version)).fetchall()266 return [{"path": r["path"], "sha256": r["sha256"], "rows": int(r["row_count"]), "bytes": int(r["bytes"])} for r in rows]267268 def gc_orphans(self, older_than_hours: int = 2) -> int:269 """Delete curated files that are not registered in lake_files (left by a killed job before publication)."""270 from findata_core.db import connection as _conn271 with _conn() as conn:272 known = {r["path"] for r in conn.execute("SELECT path FROM lake_files").fetchall()}273 cutoff = datetime.now(UTC).timestamp() - older_than_hours * 3600274 n = 0275 for r in self.roots():276 for p in (r / "curated").rglob("*.parquet"):277 rel = str(p.relative_to(r))278 if rel not in known and p.stat().st_mtime < cutoff:279 p.unlink()280 n += 1281 return n282283 def gc_retired(self, older_than_days: int = 7) -> int:284 cutoff = datetime.now(UTC).timestamp() - older_than_days * 86400285 n = 0286 for r in self.roots():287 for p in (r / "_retired").rglob("*.parquet"):288 if p.stat().st_mtime < cutoff:289 p.unlink()290 n += 1291 for p in (r / "_tmp").glob("*"):292 if p.stat().st_mtime < datetime.now(UTC).timestamp() - 86400:293 p.unlink()294 for p in (r / "raw").rglob("*.part"): # interrupted downloads295 if p.stat().st_mtime < datetime.now(UTC).timestamp() - 6 * 3600:296 p.unlink()297 return n298299300def _sha256(path: Path) -> str:301 h = hashlib.sha256()302 with open(path, "rb") as f:303 for chunk in iter(lambda: f.read(4 * 1024 * 1024), b""):304 h.update(chunk)305 return h.hexdigest()306307308def live_files(conn, dataset_id: str, symbols: list[str] | None = None, start: date | None = None,309 end: date | None = None, version: int | None = None) -> list[dict[str, Any]]:310 """Files to scan for a query (planner). `version=None` → current live set."""311 sql = "SELECT path, layout, symbol_min, symbol_max, date_min, date_max, row_count FROM lake_files WHERE dataset_id=%s"312 params: list[Any] = [dataset_id]313 if version is None:314 sql += " AND version_removed IS NULL"315 else:316 sql += " AND version_added <= %s AND (version_removed IS NULL OR version_removed > %s)"317 params += [version, version]318 if start is not None:319 sql += " AND date_max >= %s"320 params.append(start)321 if end is not None:322 sql += " AND date_min <= %s"323 params.append(end)324 if symbols:325 # a file is relevant when its symbol range covers any requested symbol; per-symbol base files hit the326 # (dataset_id, symbol_min) index, multi-symbol delta files are matched by range327 sql += " AND EXISTS (SELECT 1 FROM unnest(%s::text[]) s WHERE s BETWEEN symbol_min AND symbol_max)"328 params.append(list(symbols))329 return conn.execute(sql, params).fetchall()330