"""Parquet lake: layout, atomic publication, versions and snapshots. Layout (relative to FINDATA_LAKE): raw////.zip immutable vendor archives (+ sha256 in raw_archives) curated//delta/session_date=YYYY-MM-DD/part-.parquet one file per session (all symbols, sorted) curated//base/symbol=/year=YYYY/part-.parquet compacted / backfilled history manifests//v.json reproducible snapshot manifests _tmp/ staging (same filesystem → atomic rename) _retired//... files removed from the live set (GC after 7 days) exports/. Invariant: at any dataset version, each (instrument_id, timestamp_utc) is in exactly one live file. Only the ingest worker writes under curated/ (file lock `.writer.lock`); the API only reads. """ from __future__ import annotations import fcntl import hashlib import json import os import shutil import uuid from dataclasses import dataclass from datetime import UTC, date, datetime from pathlib import Path from typing import Any import pyarrow as pa import pyarrow.parquet as pq from findata_core.config import get_settings from findata_core.db import transaction ROW_GROUP_INTRADAY = 65_536 ROW_GROUP_DAILY = 1_000_000 @dataclass class WrittenFile: rel_path: str layout: str symbol_min: str symbol_max: str date_min: date date_max: date row_count: int bytes: int sha256: str def mkdir_retry(path: Path, attempts: int = 20) -> None: """mkdir -p that survives EINTR (seen on a Thunderbolt APFS volume while Spotlight was indexing it).""" import time for i in range(attempts): try: path.mkdir(parents=True, exist_ok=True) return except InterruptedError: time.sleep(0.2 * (i + 1)) path.mkdir(parents=True, exist_ok=True) class WriterLock: def __init__(self, lake: Path): self.path = lake / ".writer.lock" self._fh = None def __enter__(self): self.path.parent.mkdir(parents=True, exist_ok=True) self._fh = open(self.path, "w") try: fcntl.flock(self._fh, fcntl.LOCK_EX | fcntl.LOCK_NB) except BlockingIOError as e: raise RuntimeError("another lake writer holds the lock") from e self._fh.write(f"{os.getpid()} {datetime.now(UTC).isoformat()}\n") self._fh.flush() return self def __exit__(self, *exc): if self._fh: fcntl.flock(self._fh, fcntl.LOCK_UN) self._fh.close() class Lake: """Two physical roots: `root` (internal SSD, hot datasets, manifests, exports) and optional `cold` (external SSD: raw archives + every dataset not listed in FINDATA_HOT_DATASETS). Relative paths stored in `lake_files` are root-agnostic (`curated//…`); `root_for(dataset)` picks the physical root.""" def __init__(self, root: Path | None = None, cold: Path | None = None): s = get_settings() self.root = (root or s.lake).expanduser() self.cold = (cold if cold is not None else s.lake_cold) self.cold = self.cold.expanduser() if self.cold else None self.hot = s.hot_dataset_ids for sub in ("raw", "curated", "manifests", "_tmp", "_retired", "exports"): (self.root / sub).mkdir(parents=True, exist_ok=True) if self.cold: if not (self.cold / ".findata-cold-root").exists(): raise RuntimeError(f"cold lake root {self.cold} is not mounted (marker .findata-cold-root missing)") for sub in ("raw", "curated", "_tmp", "_retired"): (self.cold / sub).mkdir(parents=True, exist_ok=True) # ------------------------------------------------------------------ roots def root_for(self, dataset_id: str) -> Path: if self.cold and dataset_id not in self.hot: return self.cold return self.root def roots(self) -> list[Path]: return [self.root] + ([self.cold] if self.cold else []) def _root_of_rel(self, rel: str) -> Path: parts = rel.split("/") if len(parts) >= 2 and parts[0] in ("curated", "_retired"): return self.root_for(parts[1]) if parts and parts[0] == "raw": return self.cold or self.root return self.root def raw_dir(self, source_id: str, dataset_id: str) -> Path: p = (self.cold or self.root) / "raw" / source_id / dataset_id / str(datetime.now(UTC).year) mkdir_retry(p) return p def abs(self, rel: str) -> Path: return self._root_of_rel(rel) / rel def rel(self, path: Path) -> str: for r in self.roots(): try: return str(path.relative_to(r)) except ValueError: continue raise ValueError(f"{path} is not inside a lake root") def resolve(self, rel: str) -> Path: """Live path, or its retired location (older snapshots stay readable until GC after 7 days).""" p = self.abs(rel) if p.exists(): return p if rel.startswith("curated/"): r = self._root_of_rel(rel) / "_retired" / rel[len("curated/"):] if r.exists(): return r return p def free_gb(self, root: Path | None = None) -> float: st = os.statvfs(root or self.root) return st.f_bavail * st.f_frsize / 1e9 def used_bytes(self, sub: str = "curated") -> int: total = 0 for r in self.roots(): for p in (r / sub).rglob("*"): if p.is_file(): total += p.stat().st_size return total def check_budget(self, dataset_id: str | None = None) -> None: s = get_settings() root = self.root_for(dataset_id) if dataset_id else self.root limit = s.lake_cold_hard_min_free_gb if (self.cold and root == self.cold) else s.lake_hard_min_free_gb if self.free_gb(root) < limit: raise RuntimeError(f"lake hard stop on {root}: only {self.free_gb(root):.1f} GB free (< {limit} GB)") # ------------------------------------------------------------------ writing def write_table(self, dataset_id: str, table: pa.Table, layout: str, partition: dict[str, str], sort: bool = True) -> WrittenFile: """Write a Parquet file into _tmp then move it atomically under curated/. Does NOT register it.""" self.check_budget(dataset_id) if table.num_rows == 0: raise ValueError("refusing to write an empty file") root = self.root_for(dataset_id) is_bars = "symbol_as_reported" in table.column_names and "timestamp_utc" in table.column_names if sort and is_bars: table = table.sort_by([("symbol_as_reported", "ascending"), ("timestamp_utc", "ascending")]) part_dir = "/".join(f"{k}={v}" for k, v in partition.items()) rel = f"curated/{dataset_id}/{layout}/{part_dir}/part-{uuid.uuid4().hex[:12]}.parquet" final = root / rel tmp = root / "_tmp" / f"{uuid.uuid4().hex}.parquet" interval = table.column("interval")[0].as_py() if "interval" in table.column_names else "1d" rg = ROW_GROUP_DAILY if interval == "1d" else ROW_GROUP_INTRADAY dict_cols = [c for c in ("symbol_as_reported", "instrument_id", "asset_class", "interval", "currency", "adjustment_type", "source_id", "underlying", "call_put") if c in table.column_names] pq.write_table(table, tmp, compression="zstd", compression_level=3, row_group_size=rg, write_statistics=True, use_dictionary=dict_cols) with open(tmp, "rb") as f: os.fsync(f.fileno()) digest = _sha256(tmp) mkdir_retry(final.parent) os.replace(tmp, final) import pyarrow.compute as pc syms = table.column("symbol_as_reported") if is_bars else table.column("underlying") dates = table.column("session_date") if is_bars else table.column("trade_date") return WrittenFile(rel_path=rel, layout=layout, symbol_min=pc.min(syms).as_py(), symbol_max=pc.max(syms).as_py(), date_min=pc.min(dates).as_py(), date_max=pc.max(dates).as_py(), row_count=table.num_rows, bytes=final.stat().st_size, sha256=digest) # ------------------------------------------------------------------ publication (atomic, versioned) def publish(self, dataset_id: str, added: list[WrittenFile], retire_paths: list[str], ingest_run_id: str | None, note: str = "") -> dict[str, Any]: """One transaction: register added files, retire replaced files, bump the dataset version, write manifest. Retired files are moved to _retired/ *after* commit (readers that already opened them keep working).""" with transaction() as conn: row = conn.execute("SELECT current_version FROM datasets WHERE id=%s FOR UPDATE", (dataset_id,)).fetchone() if row is None: raise RuntimeError(f"unknown dataset {dataset_id}") version = int(row["current_version"]) + 1 for p in retire_paths: conn.execute("UPDATE lake_files SET version_removed=%s WHERE path=%s AND version_removed IS NULL", (version, p)) for f in added: conn.execute( "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) " "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)", (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)) live = conn.execute("SELECT path, sha256, row_count, bytes FROM lake_files WHERE dataset_id=%s AND version_removed IS NULL ORDER BY path", (dataset_id,)).fetchall() snapshot_id = hashlib.sha256("\n".join(f"{r['path']} {r['sha256']}" for r in live).encode()).hexdigest() file_count = len(live) row_count = sum(int(r["row_count"]) for r in live) nbytes = sum(int(r["bytes"]) for r in live) 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)", (dataset_id, version, snapshot_id, file_count, row_count, nbytes, ingest_run_id, note)) conn.execute("UPDATE datasets SET current_version=%s, updated_at=now() WHERE id=%s", (version, dataset_id)) manifest = {"dataset_id": dataset_id, "version": version, "snapshot_id": snapshot_id, "created_at": datetime.now(UTC).isoformat(), "ingest_run_id": ingest_run_id, "note": note, "files": [{"path": r["path"], "sha256": r["sha256"], "rows": int(r["row_count"]), "bytes": int(r["bytes"])} for r in live]} mpath = self.root / "manifests" / dataset_id / f"v{version}.json" mpath.parent.mkdir(parents=True, exist_ok=True) tmp = mpath.with_suffix(".json.tmp") tmp.write_text(json.dumps(manifest, indent=1, default=str)) os.replace(tmp, mpath) for p in retire_paths: src = self.abs(p) if src.exists(): r = self._root_of_rel(p) dst = r / "_retired" / (p[len("curated/"):] if p.startswith("curated/") else p) mkdir_retry(dst.parent) shutil.move(str(src), str(dst)) return {"version": version, "snapshot_id": snapshot_id, "file_count": file_count, "row_count": row_count, "bytes": nbytes} def gc_orphans(self, older_than_hours: int = 2) -> int: """Delete curated files that are not registered in lake_files (left by a killed job before publication).""" from findata_core.db import connection as _conn with _conn() as conn: known = {r["path"] for r in conn.execute("SELECT path FROM lake_files").fetchall()} cutoff = datetime.now(UTC).timestamp() - older_than_hours * 3600 n = 0 for r in self.roots(): for p in (r / "curated").rglob("*.parquet"): rel = str(p.relative_to(r)) if rel not in known and p.stat().st_mtime < cutoff: p.unlink() n += 1 return n def gc_retired(self, older_than_days: int = 7) -> int: cutoff = datetime.now(UTC).timestamp() - older_than_days * 86400 n = 0 for r in self.roots(): for p in (r / "_retired").rglob("*.parquet"): if p.stat().st_mtime < cutoff: p.unlink() n += 1 for p in (r / "_tmp").glob("*"): if p.stat().st_mtime < datetime.now(UTC).timestamp() - 86400: p.unlink() return n def _sha256(path: Path) -> str: h = hashlib.sha256() with open(path, "rb") as f: for chunk in iter(lambda: f.read(4 * 1024 * 1024), b""): h.update(chunk) return h.hexdigest() def live_files(conn, dataset_id: str, symbols: list[str] | None = None, start: date | None = None, end: date | None = None, version: int | None = None) -> list[dict[str, Any]]: """Files to scan for a query (planner). `version=None` → current live set.""" sql = "SELECT path, layout, symbol_min, symbol_max, date_min, date_max, row_count FROM lake_files WHERE dataset_id=%s" params: list[Any] = [dataset_id] if version is None: sql += " AND version_removed IS NULL" else: sql += " AND version_added <= %s AND (version_removed IS NULL OR version_removed > %s)" params += [version, version] if start is not None: sql += " AND date_max >= %s" params.append(start) if end is not None: sql += " AND date_min <= %s" params.append(end) rows = conn.execute(sql, params).fetchall() if symbols: rows = [r for r in rows if any(r["symbol_min"] <= s <= r["symbol_max"] for s in symbols)] return rows