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 def __init__(self, root: Path | None = None):74 self.root = (root or get_settings().lake).expanduser()75 for sub in ("raw", "curated", "manifests", "_tmp", "_retired", "exports"):76 (self.root / sub).mkdir(parents=True, exist_ok=True)7778 # ------------------------------------------------------------------ paths79 def raw_dir(self, source_id: str, dataset_id: str) -> Path:80 p = self.root / "raw" / source_id / dataset_id / str(datetime.now(UTC).year)81 p.mkdir(parents=True, exist_ok=True)82 return p8384 def abs(self, rel: str) -> Path:85 return self.root / rel8687 def resolve(self, rel: str) -> Path:88 """Live path, or its retired location (older snapshots stay readable until GC after 7 days)."""89 p = self.root / rel90 if p.exists():91 return p92 if rel.startswith("curated/"):93 r = self.root / "_retired" / rel[len("curated/"):]94 if r.exists():95 return r96 return p9798 def free_gb(self) -> float:99 st = os.statvfs(self.root)100 return st.f_bavail * st.f_frsize / 1e9101102 def used_bytes(self, sub: str = "curated") -> int:103 total = 0104 for p in (self.root / sub).rglob("*"):105 if p.is_file():106 total += p.stat().st_size107 return total108109 def check_budget(self) -> None:110 s = get_settings()111 if self.free_gb() < s.lake_hard_min_free_gb:112 raise RuntimeError(f"lake hard stop: only {self.free_gb():.1f} GB free (< {s.lake_hard_min_free_gb} GB)")113114 # ------------------------------------------------------------------ writing115 def write_table(self, dataset_id: str, table: pa.Table, layout: str, partition: dict[str, str],116 sort: bool = True) -> WrittenFile:117 """Write a Parquet file into _tmp then move it atomically under curated/. Does NOT register it."""118 self.check_budget()119 if table.num_rows == 0:120 raise ValueError("refusing to write an empty file")121 if sort:122 table = table.sort_by([("symbol_as_reported", "ascending"), ("timestamp_utc", "ascending")])123 part_dir = "/".join(f"{k}={v}" for k, v in partition.items())124 rel = f"curated/{dataset_id}/{layout}/{part_dir}/part-{uuid.uuid4().hex[:12]}.parquet"125 final = self.root / rel126 tmp = self.root / "_tmp" / f"{uuid.uuid4().hex}.parquet"127 interval = table.column("interval")[0].as_py() if "interval" in table.column_names else "1d"128 rg = ROW_GROUP_DAILY if interval == "1d" else ROW_GROUP_INTRADAY129 pq.write_table(table, tmp, compression="zstd", compression_level=3, row_group_size=rg,130 write_statistics=True, use_dictionary=["symbol_as_reported", "instrument_id", "asset_class",131 "interval", "currency", "adjustment_type", "source_id"])132 with open(tmp, "rb") as f:133 os.fsync(f.fileno())134 digest = _sha256(tmp)135 final.parent.mkdir(parents=True, exist_ok=True)136 os.replace(tmp, final)137 syms = table.column("symbol_as_reported")138 dates = table.column("session_date")139 import pyarrow.compute as pc140 return WrittenFile(rel_path=rel, layout=layout, symbol_min=pc.min(syms).as_py(), symbol_max=pc.max(syms).as_py(),141 date_min=pc.min(dates).as_py(), date_max=pc.max(dates).as_py(), row_count=table.num_rows,142 bytes=final.stat().st_size, sha256=digest)143144 # ------------------------------------------------------------------ publication (atomic, versioned)145 def publish(self, dataset_id: str, added: list[WrittenFile], retire_paths: list[str], ingest_run_id: str | None,146 note: str = "") -> dict[str, Any]:147 """One transaction: register added files, retire replaced files, bump the dataset version, write manifest.148 Retired files are moved to _retired/ *after* commit (readers that already opened them keep working)."""149 with transaction() as conn:150 row = conn.execute("SELECT current_version FROM datasets WHERE id=%s FOR UPDATE", (dataset_id,)).fetchone()151 if row is None:152 raise RuntimeError(f"unknown dataset {dataset_id}")153 version = int(row["current_version"]) + 1154 for p in retire_paths:155 conn.execute("UPDATE lake_files SET version_removed=%s WHERE path=%s AND version_removed IS NULL", (version, p))156 for f in added:157 conn.execute(158 "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) "159 "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)",160 (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))161 live = conn.execute("SELECT path, sha256, row_count, bytes FROM lake_files WHERE dataset_id=%s AND version_removed IS NULL ORDER BY path",162 (dataset_id,)).fetchall()163 snapshot_id = hashlib.sha256("\n".join(f"{r['path']} {r['sha256']}" for r in live).encode()).hexdigest()164 file_count = len(live)165 row_count = sum(int(r["row_count"]) for r in live)166 nbytes = sum(int(r["bytes"]) for r in live)167 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)",168 (dataset_id, version, snapshot_id, file_count, row_count, nbytes, ingest_run_id, note))169 conn.execute("UPDATE datasets SET current_version=%s, updated_at=now() WHERE id=%s", (version, dataset_id))170 manifest = {"dataset_id": dataset_id, "version": version, "snapshot_id": snapshot_id,171 "created_at": datetime.now(UTC).isoformat(), "ingest_run_id": ingest_run_id, "note": note,172 "files": [{"path": r["path"], "sha256": r["sha256"], "rows": int(r["row_count"]), "bytes": int(r["bytes"])} for r in live]}173 mpath = self.root / "manifests" / dataset_id / f"v{version}.json"174 mpath.parent.mkdir(parents=True, exist_ok=True)175 tmp = mpath.with_suffix(".json.tmp")176 tmp.write_text(json.dumps(manifest, indent=1, default=str))177 os.replace(tmp, mpath)178 for p in retire_paths:179 src = self.root / p180 if src.exists():181 dst = self.root / "_retired" / p[len("curated/"):] if p.startswith("curated/") else self.root / "_retired" / p182 dst.parent.mkdir(parents=True, exist_ok=True)183 shutil.move(str(src), str(dst))184 return {"version": version, "snapshot_id": snapshot_id, "file_count": file_count, "row_count": row_count, "bytes": nbytes}185186 def gc_retired(self, older_than_days: int = 7) -> int:187 cutoff = datetime.now(UTC).timestamp() - older_than_days * 86400188 n = 0189 for p in (self.root / "_retired").rglob("*.parquet"):190 if p.stat().st_mtime < cutoff:191 p.unlink()192 n += 1193 for p in (self.root / "_tmp").glob("*"):194 if p.stat().st_mtime < datetime.now(UTC).timestamp() - 86400:195 p.unlink()196 return n197198199def _sha256(path: Path) -> str:200 h = hashlib.sha256()201 with open(path, "rb") as f:202 for chunk in iter(lambda: f.read(4 * 1024 * 1024), b""):203 h.update(chunk)204 return h.hexdigest()205206207def live_files(conn, dataset_id: str, symbols: list[str] | None = None, start: date | None = None,208 end: date | None = None, version: int | None = None) -> list[dict[str, Any]]:209 """Files to scan for a query (planner). `version=None` → current live set."""210 sql = "SELECT path, layout, symbol_min, symbol_max, date_min, date_max, row_count FROM lake_files WHERE dataset_id=%s"211 params: list[Any] = [dataset_id]212 if version is None:213 sql += " AND version_removed IS NULL"214 else:215 sql += " AND version_added <= %s AND (version_removed IS NULL OR version_removed > %s)"216 params += [version, version]217 if start is not None:218 sql += " AND date_max >= %s"219 params.append(start)220 if end is not None:221 sql += " AND date_min <= %s"222 params.append(end)223 rows = conn.execute(sql, params).fetchall()224 if symbols:225 rows = [r for r in rows if any(r["symbol_min"] <= s <= r["symbol_max"] for s in symbols)]226 return rows227