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 pq3031from findata_core.config import get_settings32from findata_core.db import transaction3334ROW_GROUP_INTRADAY = 65_53635ROW_GROUP_DAILY = 1_000_000363738@dataclass39class WrittenFile:40 rel_path: str41 layout: str42 symbol_min: str43 symbol_max: str44 date_min: date45 date_max: date46 row_count: int47 bytes: int48 sha256: str495051class WriterLock:52 def __init__(self, lake: Path):53 self.path = lake / ".writer.lock"54 self._fh = None5556 def __enter__(self):57 self.path.parent.mkdir(parents=True, exist_ok=True)58 self._fh = open(self.path, "w")59 try:60 fcntl.flock(self._fh, fcntl.LOCK_EX | fcntl.LOCK_NB)61 except BlockingIOError as e:62 raise RuntimeError("another lake writer holds the lock") from e63 self._fh.write(f"{os.getpid()} {datetime.now(UTC).isoformat()}\n")64 self._fh.flush()65 return self6667 def __exit__(self, *exc):68 if self._fh:69 fcntl.flock(self._fh, fcntl.LOCK_UN)70 self._fh.close()717273class Lake:74 def __init__(self, root: Path | None = None):75 self.root = (root or get_settings().lake).expanduser()76 for sub in ("raw", "curated", "manifests", "_tmp", "_retired", "exports"):77 (self.root / sub).mkdir(parents=True, exist_ok=True)7879 # ------------------------------------------------------------------ paths80 def raw_dir(self, source_id: str, dataset_id: str) -> Path:81 p = self.root / "raw" / source_id / dataset_id / str(datetime.now(UTC).year)82 p.mkdir(parents=True, exist_ok=True)83 return p8485 def abs(self, rel: str) -> Path:86 return self.root / rel8788 def resolve(self, rel: str) -> Path:89 """Live path, or its retired location (older snapshots stay readable until GC after 7 days)."""90 p = self.root / rel91 if p.exists():92 return p93 if rel.startswith("curated/"):94 r = self.root / "_retired" / rel[len("curated/"):]95 if r.exists():96 return r97 return p9899 def free_gb(self) -> float:100 st = os.statvfs(self.root)101 return st.f_bavail * st.f_frsize / 1e9102103 def used_bytes(self, sub: str = "curated") -> int:104 total = 0105 for p in (self.root / sub).rglob("*"):106 if p.is_file():107 total += p.stat().st_size108 return total109110 def check_budget(self) -> None:111 s = get_settings()112 if self.free_gb() < s.lake_hard_min_free_gb:113 raise RuntimeError(f"lake hard stop: only {self.free_gb():.1f} GB free (< {s.lake_hard_min_free_gb} GB)")114115 # ------------------------------------------------------------------ writing116 def write_table(self, dataset_id: str, table: pa.Table, layout: str, partition: dict[str, str],117 sort: bool = True) -> WrittenFile:118 """Write a Parquet file into _tmp then move it atomically under curated/. Does NOT register it."""119 self.check_budget()120 if table.num_rows == 0:121 raise ValueError("refusing to write an empty file")122 if sort:123 table = table.sort_by([("symbol_as_reported", "ascending"), ("timestamp_utc", "ascending")])124 part_dir = "/".join(f"{k}={v}" for k, v in partition.items())125 rel = f"curated/{dataset_id}/{layout}/{part_dir}/part-{uuid.uuid4().hex[:12]}.parquet"126 final = self.root / rel127 tmp = self.root / "_tmp" / f"{uuid.uuid4().hex}.parquet"128 interval = table.column("interval")[0].as_py() if "interval" in table.column_names else "1d"129 rg = ROW_GROUP_DAILY if interval == "1d" else ROW_GROUP_INTRADAY130 pq.write_table(table, tmp, compression="zstd", compression_level=3, row_group_size=rg,131 write_statistics=True, use_dictionary=["symbol_as_reported", "instrument_id", "asset_class",132 "interval", "currency", "adjustment_type", "source_id"])133 with open(tmp, "rb") as f:134 os.fsync(f.fileno())135 digest = _sha256(tmp)136 final.parent.mkdir(parents=True, exist_ok=True)137 os.replace(tmp, final)138 syms = table.column("symbol_as_reported")139 dates = table.column("session_date")140 import pyarrow.compute as pc141 return WrittenFile(rel_path=rel, layout=layout, symbol_min=pc.min(syms).as_py(), symbol_max=pc.max(syms).as_py(),142 date_min=pc.min(dates).as_py(), date_max=pc.max(dates).as_py(), row_count=table.num_rows,143 bytes=final.stat().st_size, sha256=digest)144145 # ------------------------------------------------------------------ publication (atomic, versioned)146 def publish(self, dataset_id: str, added: list[WrittenFile], retire_paths: list[str], ingest_run_id: str | None,147 note: str = "") -> dict[str, Any]:148 """One transaction: register added files, retire replaced files, bump the dataset version, write manifest.149 Retired files are moved to _retired/ *after* commit (readers that already opened them keep working)."""150 with transaction() as conn:151 row = conn.execute("SELECT current_version FROM datasets WHERE id=%s FOR UPDATE", (dataset_id,)).fetchone()152 if row is None:153 raise RuntimeError(f"unknown dataset {dataset_id}")154 version = int(row["current_version"]) + 1155 for p in retire_paths:156 conn.execute("UPDATE lake_files SET version_removed=%s WHERE path=%s AND version_removed IS NULL", (version, p))157 for f in added:158 conn.execute(159 "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) "160 "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)",161 (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))162 live = conn.execute("SELECT path, sha256, row_count, bytes FROM lake_files WHERE dataset_id=%s AND version_removed IS NULL ORDER BY path",163 (dataset_id,)).fetchall()164 snapshot_id = hashlib.sha256("\n".join(f"{r['path']} {r['sha256']}" for r in live).encode()).hexdigest()165 file_count = len(live)166 row_count = sum(int(r["row_count"]) for r in live)167 nbytes = sum(int(r["bytes"]) for r in live)168 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)",169 (dataset_id, version, snapshot_id, file_count, row_count, nbytes, ingest_run_id, note))170 conn.execute("UPDATE datasets SET current_version=%s, updated_at=now() WHERE id=%s", (version, dataset_id))171 manifest = {"dataset_id": dataset_id, "version": version, "snapshot_id": snapshot_id,172 "created_at": datetime.now(UTC).isoformat(), "ingest_run_id": ingest_run_id, "note": note,173 "files": [{"path": r["path"], "sha256": r["sha256"], "rows": int(r["row_count"]), "bytes": int(r["bytes"])} for r in live]}174 mpath = self.root / "manifests" / dataset_id / f"v{version}.json"175 mpath.parent.mkdir(parents=True, exist_ok=True)176 tmp = mpath.with_suffix(".json.tmp")177 tmp.write_text(json.dumps(manifest, indent=1, default=str))178 os.replace(tmp, mpath)179 for p in retire_paths:180 src = self.root / p181 if src.exists():182 dst = self.root / "_retired" / p[len("curated/"):] if p.startswith("curated/") else self.root / "_retired" / p183 dst.parent.mkdir(parents=True, exist_ok=True)184 shutil.move(str(src), str(dst))185 return {"version": version, "snapshot_id": snapshot_id, "file_count": file_count, "row_count": row_count, "bytes": nbytes}186187 def gc_retired(self, older_than_days: int = 7) -> int:188 cutoff = datetime.now(UTC).timestamp() - older_than_days * 86400189 n = 0190 for p in (self.root / "_retired").rglob("*.parquet"):191 if p.stat().st_mtime < cutoff:192 p.unlink()193 n += 1194 for p in (self.root / "_tmp").glob("*"):195 if p.stat().st_mtime < datetime.now(UTC).timestamp() - 86400:196 p.unlink()197 return n198199200def _sha256(path: Path) -> str:201 h = hashlib.sha256()202 with open(path, "rb") as f:203 for chunk in iter(lambda: f.read(4 * 1024 * 1024), b""):204 h.update(chunk)205 return h.hexdigest()206207208def live_files(conn, dataset_id: str, symbols: list[str] | None = None, start: date | None = None,209 end: date | None = None, version: int | None = None) -> list[dict[str, Any]]:210 """Files to scan for a query (planner). `version=None` → current live set."""211 sql = "SELECT path, layout, symbol_min, symbol_max, date_min, date_max, row_count FROM lake_files WHERE dataset_id=%s"212 params: list[Any] = [dataset_id]213 if version is None:214 sql += " AND version_removed IS NULL"215 else:216 sql += " AND version_added <= %s AND (version_removed IS NULL OR version_removed > %s)"217 params += [version, version]218 if start is not None:219 sql += " AND date_max >= %s"220 params.append(start)221 if end is not None:222 sql += " AND date_min <= %s"223 params.append(end)224 rows = conn.execute(sql, params).fetchall()225 if symbols:226 rows = [r for r in rows if any(r["symbol_min"] <= s <= r["symbol_max"] for s in symbols)]227 return rows228