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: str484950class WriterLock:51 def __init__(self, lake: Path):52 self.path = lake / ".writer.lock"53 self._fh = None5455 def __enter__(self):56 self.path.parent.mkdir(parents=True, exist_ok=True)57 self._fh = open(self.path, "w")58 try:59 fcntl.flock(self._fh, fcntl.LOCK_EX | fcntl.LOCK_NB)60 except BlockingIOError as e:61 raise RuntimeError("another lake writer holds the lock") from e62 self._fh.write(f"{os.getpid()} {datetime.now(UTC).isoformat()}\n")63 self._fh.flush()64 return self6566 def __exit__(self, *exc):67 if self._fh:68 fcntl.flock(self._fh, fcntl.LOCK_UN)69 self._fh.close()707172class Lake:73 """Two physical roots: `root` (internal SSD, hot datasets, manifests, exports) and optional `cold`74 (external SSD: raw archives + every dataset not listed in FINDATA_HOT_DATASETS). Relative paths stored in75 `lake_files` are root-agnostic (`curated/<dataset>/…`); `root_for(dataset)` picks the physical root."""7677 def __init__(self, root: Path | None = None, cold: Path | None = None):78 s = get_settings()79 self.root = (root or s.lake).expanduser()80 self.cold = (cold if cold is not None else s.lake_cold)81 self.cold = self.cold.expanduser() if self.cold else None82 self.hot = s.hot_dataset_ids83 for sub in ("raw", "curated", "manifests", "_tmp", "_retired", "exports"):84 (self.root / sub).mkdir(parents=True, exist_ok=True)85 if self.cold:86 if not (self.cold / ".findata-cold-root").exists():87 raise RuntimeError(f"cold lake root {self.cold} is not mounted (marker .findata-cold-root missing)")88 for sub in ("raw", "curated", "_tmp", "_retired"):89 (self.cold / sub).mkdir(parents=True, exist_ok=True)9091 # ------------------------------------------------------------------ roots92 def root_for(self, dataset_id: str) -> Path:93 if self.cold and dataset_id not in self.hot:94 return self.cold95 return self.root9697 def roots(self) -> list[Path]:98 return [self.root] + ([self.cold] if self.cold else [])99100 def _root_of_rel(self, rel: str) -> Path:101 parts = rel.split("/")102 if len(parts) >= 2 and parts[0] in ("curated", "_retired"):103 return self.root_for(parts[1])104 if parts and parts[0] == "raw":105 return self.cold or self.root106 return self.root107108 def raw_dir(self, source_id: str, dataset_id: str) -> Path:109 p = (self.cold or self.root) / "raw" / source_id / dataset_id / str(datetime.now(UTC).year)110 p.mkdir(parents=True, exist_ok=True)111 return p112113 def abs(self, rel: str) -> Path:114 return self._root_of_rel(rel) / rel115116 def rel(self, path: Path) -> str:117 for r in self.roots():118 try:119 return str(path.relative_to(r))120 except ValueError:121 continue122 raise ValueError(f"{path} is not inside a lake root")123124 def resolve(self, rel: str) -> Path:125 """Live path, or its retired location (older snapshots stay readable until GC after 7 days)."""126 p = self.abs(rel)127 if p.exists():128 return p129 if rel.startswith("curated/"):130 r = self._root_of_rel(rel) / "_retired" / rel[len("curated/"):]131 if r.exists():132 return r133 return p134135 def free_gb(self, root: Path | None = None) -> float:136 st = os.statvfs(root or self.root)137 return st.f_bavail * st.f_frsize / 1e9138139 def used_bytes(self, sub: str = "curated") -> int:140 total = 0141 for r in self.roots():142 for p in (r / sub).rglob("*"):143 if p.is_file():144 total += p.stat().st_size145 return total146147 def check_budget(self, dataset_id: str | None = None) -> None:148 s = get_settings()149 root = self.root_for(dataset_id) if dataset_id else self.root150 limit = s.lake_cold_hard_min_free_gb if (self.cold and root == self.cold) else s.lake_hard_min_free_gb151 if self.free_gb(root) < limit:152 raise RuntimeError(f"lake hard stop on {root}: only {self.free_gb(root):.1f} GB free (< {limit} GB)")153154 # ------------------------------------------------------------------ writing155 def write_table(self, dataset_id: str, table: pa.Table, layout: str, partition: dict[str, str],156 sort: bool = True) -> WrittenFile:157 """Write a Parquet file into _tmp then move it atomically under curated/. Does NOT register it."""158 self.check_budget(dataset_id)159 if table.num_rows == 0:160 raise ValueError("refusing to write an empty file")161 root = self.root_for(dataset_id)162 is_bars = "symbol_as_reported" in table.column_names and "timestamp_utc" in table.column_names163 if sort and is_bars:164 table = table.sort_by([("symbol_as_reported", "ascending"), ("timestamp_utc", "ascending")])165 part_dir = "/".join(f"{k}={v}" for k, v in partition.items())166 rel = f"curated/{dataset_id}/{layout}/{part_dir}/part-{uuid.uuid4().hex[:12]}.parquet"167 final = root / rel168 tmp = root / "_tmp" / f"{uuid.uuid4().hex}.parquet"169 interval = table.column("interval")[0].as_py() if "interval" in table.column_names else "1d"170 rg = ROW_GROUP_DAILY if interval == "1d" else ROW_GROUP_INTRADAY171 dict_cols = [c for c in ("symbol_as_reported", "instrument_id", "asset_class", "interval", "currency", "adjustment_type",172 "source_id", "underlying", "call_put") if c in table.column_names]173 pq.write_table(table, tmp, compression="zstd", compression_level=3, row_group_size=rg,174 write_statistics=True, use_dictionary=dict_cols)175 with open(tmp, "rb") as f:176 os.fsync(f.fileno())177 digest = _sha256(tmp)178 final.parent.mkdir(parents=True, exist_ok=True)179 os.replace(tmp, final)180 import pyarrow.compute as pc181 syms = table.column("symbol_as_reported") if is_bars else table.column("underlying")182 dates = table.column("session_date") if is_bars else table.column("trade_date")183 return WrittenFile(rel_path=rel, layout=layout, symbol_min=pc.min(syms).as_py(), symbol_max=pc.max(syms).as_py(),184 date_min=pc.min(dates).as_py(), date_max=pc.max(dates).as_py(), row_count=table.num_rows,185 bytes=final.stat().st_size, sha256=digest)186187 # ------------------------------------------------------------------ publication (atomic, versioned)188 def publish(self, dataset_id: str, added: list[WrittenFile], retire_paths: list[str], ingest_run_id: str | None,189 note: str = "") -> dict[str, Any]:190 """One transaction: register added files, retire replaced files, bump the dataset version, write manifest.191 Retired files are moved to _retired/ *after* commit (readers that already opened them keep working)."""192 with transaction() as conn:193 row = conn.execute("SELECT current_version FROM datasets WHERE id=%s FOR UPDATE", (dataset_id,)).fetchone()194 if row is None:195 raise RuntimeError(f"unknown dataset {dataset_id}")196 version = int(row["current_version"]) + 1197 for p in retire_paths:198 conn.execute("UPDATE lake_files SET version_removed=%s WHERE path=%s AND version_removed IS NULL", (version, p))199 for f in added:200 conn.execute(201 "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) "202 "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)",203 (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))204 live = conn.execute("SELECT path, sha256, row_count, bytes FROM lake_files WHERE dataset_id=%s AND version_removed IS NULL ORDER BY path",205 (dataset_id,)).fetchall()206 snapshot_id = hashlib.sha256("\n".join(f"{r['path']} {r['sha256']}" for r in live).encode()).hexdigest()207 file_count = len(live)208 row_count = sum(int(r["row_count"]) for r in live)209 nbytes = sum(int(r["bytes"]) for r in live)210 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)",211 (dataset_id, version, snapshot_id, file_count, row_count, nbytes, ingest_run_id, note))212 conn.execute("UPDATE datasets SET current_version=%s, updated_at=now() WHERE id=%s", (version, dataset_id))213 manifest = {"dataset_id": dataset_id, "version": version, "snapshot_id": snapshot_id,214 "created_at": datetime.now(UTC).isoformat(), "ingest_run_id": ingest_run_id, "note": note,215 "files": [{"path": r["path"], "sha256": r["sha256"], "rows": int(r["row_count"]), "bytes": int(r["bytes"])} for r in live]}216 mpath = self.root / "manifests" / dataset_id / f"v{version}.json"217 mpath.parent.mkdir(parents=True, exist_ok=True)218 tmp = mpath.with_suffix(".json.tmp")219 tmp.write_text(json.dumps(manifest, indent=1, default=str))220 os.replace(tmp, mpath)221 for p in retire_paths:222 src = self.abs(p)223 if src.exists():224 r = self._root_of_rel(p)225 dst = r / "_retired" / (p[len("curated/"):] if p.startswith("curated/") else p)226 dst.parent.mkdir(parents=True, exist_ok=True)227 shutil.move(str(src), str(dst))228 return {"version": version, "snapshot_id": snapshot_id, "file_count": file_count, "row_count": row_count, "bytes": nbytes}229230 def gc_orphans(self, older_than_hours: int = 2) -> int:231 """Delete curated files that are not registered in lake_files (left by a killed job before publication)."""232 from findata_core.db import connection as _conn233 with _conn() as conn:234 known = {r["path"] for r in conn.execute("SELECT path FROM lake_files").fetchall()}235 cutoff = datetime.now(UTC).timestamp() - older_than_hours * 3600236 n = 0237 for r in self.roots():238 for p in (r / "curated").rglob("*.parquet"):239 rel = str(p.relative_to(r))240 if rel not in known and p.stat().st_mtime < cutoff:241 p.unlink()242 n += 1243 return n244245 def gc_retired(self, older_than_days: int = 7) -> int:246 cutoff = datetime.now(UTC).timestamp() - older_than_days * 86400247 n = 0248 for r in self.roots():249 for p in (r / "_retired").rglob("*.parquet"):250 if p.stat().st_mtime < cutoff:251 p.unlink()252 n += 1253 for p in (r / "_tmp").glob("*"):254 if p.stat().st_mtime < datetime.now(UTC).timestamp() - 86400:255 p.unlink()256 return n257258259def _sha256(path: Path) -> str:260 h = hashlib.sha256()261 with open(path, "rb") as f:262 for chunk in iter(lambda: f.read(4 * 1024 * 1024), b""):263 h.update(chunk)264 return h.hexdigest()265266267def live_files(conn, dataset_id: str, symbols: list[str] | None = None, start: date | None = None,268 end: date | None = None, version: int | None = None) -> list[dict[str, Any]]:269 """Files to scan for a query (planner). `version=None` → current live set."""270 sql = "SELECT path, layout, symbol_min, symbol_max, date_min, date_max, row_count FROM lake_files WHERE dataset_id=%s"271 params: list[Any] = [dataset_id]272 if version is None:273 sql += " AND version_removed IS NULL"274 else:275 sql += " AND version_added <= %s AND (version_removed IS NULL OR version_removed > %s)"276 params += [version, version]277 if start is not None:278 sql += " AND date_max >= %s"279 params.append(start)280 if end is not None:281 sql += " AND date_min <= %s"282 params.append(end)283 rows = conn.execute(sql, params).fetchall()284 if symbols:285 rows = [r for r in rows if any(r["symbol_min"] <= s <= r["symbol_max"] for s in symbols)]286 return rows287