"""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 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: def __init__(self, root: Path | None = None): self.root = (root or get_settings().lake).expanduser() for sub in ("raw", "curated", "manifests", "_tmp", "_retired", "exports"): (self.root / sub).mkdir(parents=True, exist_ok=True) # ------------------------------------------------------------------ paths def raw_dir(self, source_id: str, dataset_id: str) -> Path: p = self.root / "raw" / source_id / dataset_id / str(datetime.now(UTC).year) p.mkdir(parents=True, exist_ok=True) return p def abs(self, rel: str) -> Path: return self.root / rel def resolve(self, rel: str) -> Path: """Live path, or its retired location (older snapshots stay readable until GC after 7 days).""" p = self.root / rel if p.exists(): return p if rel.startswith("curated/"): r = self.root / "_retired" / rel[len("curated/"):] if r.exists(): return r return p def free_gb(self) -> float: st = os.statvfs(self.root) return st.f_bavail * st.f_frsize / 1e9 def used_bytes(self, sub: str = "curated") -> int: total = 0 for p in (self.root / sub).rglob("*"): if p.is_file(): total += p.stat().st_size return total def check_budget(self) -> None: s = get_settings() if self.free_gb() < s.lake_hard_min_free_gb: raise RuntimeError(f"lake hard stop: only {self.free_gb():.1f} GB free (< {s.lake_hard_min_free_gb} 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() if table.num_rows == 0: raise ValueError("refusing to write an empty file") if sort: 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 = self.root / rel tmp = self.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 pq.write_table(table, tmp, compression="zstd", compression_level=3, row_group_size=rg, write_statistics=True, use_dictionary=["symbol_as_reported", "instrument_id", "asset_class", "interval", "currency", "adjustment_type", "source_id"]) with open(tmp, "rb") as f: os.fsync(f.fileno()) digest = _sha256(tmp) final.parent.mkdir(parents=True, exist_ok=True) os.replace(tmp, final) syms = table.column("symbol_as_reported") dates = table.column("session_date") import pyarrow.compute as pc 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.root / p if src.exists(): dst = self.root / "_retired" / p[len("curated/"):] if p.startswith("curated/") else self.root / "_retired" / p dst.parent.mkdir(parents=True, exist_ok=True) 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_retired(self, older_than_days: int = 7) -> int: cutoff = datetime.now(UTC).timestamp() - older_than_days * 86400 n = 0 for p in (self.root / "_retired").rglob("*.parquet"): if p.stat().st_mtime < cutoff: p.unlink() n += 1 for p in (self.root / "_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