diff --git a/infra/.gitignore b/infra/.gitignore index 5795e87..c3ac61f 100644 --- a/infra/.gitignore +++ b/infra/.gitignore @@ -15,6 +15,10 @@ age.key .DS_Store +# lake-iceberg's pytest suite runs in place. +__pycache__/ +*.pyc + # Host-side deploy state (deploy.sh): the live per-service image tags and the # cross-repo deploy lock. Untracked on purpose -- they must survive the # `git reset --hard` at the top of every deploy (deploy.sh's comments already diff --git a/infra/lake-iceberg/Dockerfile b/infra/lake-iceberg/Dockerfile new file mode 100644 index 0000000..da79b53 --- /dev/null +++ b/infra/lake-iceberg/Dockerfile @@ -0,0 +1,20 @@ +# One-shot Iceberg sync for the ERA5 lake: registers new hive tiles from the +# era5-thermograph bucket into the iceberg/era5_daily table (see +# sync_iceberg.py). Everything ships as manylinux wheels — no toolchain. +FROM python:3.12-slim + +# The pyiceberg extras pull sqlalchemy (sqlite catalog), pyarrow (FileIO + +# parquet footers) and s3fs (bucket listing/manifest reads). +COPY requirements.txt /tmp/requirements.txt +RUN pip install --no-cache-dir -r /tmp/requirements.txt + +WORKDIR /app +COPY sync_iceberg.py . + +# The sqlite catalog lives on a volume so runs share it; losing it is fine +# (the next run re-registers the table from its latest metadata JSON). +ENV THERMOGRAPH_ICEBERG_CATALOG_DB=/state/iceberg-catalog.db \ + PYTHONUNBUFFERED=1 +VOLUME /state + +ENTRYPOINT ["python", "sync_iceberg.py"] diff --git a/infra/lake-iceberg/README.md b/infra/lake-iceberg/README.md new file mode 100644 index 0000000..d3dac78 --- /dev/null +++ b/infra/lake-iceberg/README.md @@ -0,0 +1,123 @@ +# lake-iceberg — Iceberg conversion for the ERA5 lake + +Maintains an Apache Iceberg v2 table `era5_daily` at +`s3://era5-thermograph/iceberg/era5_daily` over the lake's existing hive +parquet (`era5/daily/tile=_/year=/month=/part.parquet`). + +Key properties: + +- **No data rewrite.** `sync_iceberg.py` uses pyiceberg `add_files`: the + Iceberg metadata references the existing part files in place. Storage is + not doubled and nothing under `era5/` is ever written, moved, or deleted — + the tool only writes under the `iceberg/` prefix. +- **Incremental + idempotent.** Each run diffs the tiles in + `era5/manifest.parquet` (the source of truth for *complete* tiles — the + extractor uploads all of a tile's objects before its manifest rows appear) + against the `thermograph.synced-tiles` table property, and registers only + the new ones. Files and the property update land in the same commit, so an + interrupted run never half-syncs a tile; re-running is always safe, even + while an extraction job is writing new tiles. +- **Self-contained catalog.** A pyiceberg SqlCatalog on local sqlite + (`THERMOGRAPH_ICEBERG_CATALOG_DB`, default `/state/iceberg-catalog.db`). + The catalog is only a pointer: if the db is lost, the next run finds the + newest `*.metadata.json` under the table location and re-registers it — + the synced-tiles set lives in the table metadata itself, so nothing is + re-added. +- **Partitioning `(tile, year, month)`** via order-preserving transforms of + real columns: `truncate[12](lat_idx)` / `truncate[12](lon_idx)` (partition + values are `ti*12` / `tj*12` — the part files carry no tile column, the 12 + matches the lake's 12x12-point tile grid) plus `year(date)` and + `month(date)`. Predicates on `lat_idx`/`lon_idx`/`date` prune partitions + directly — no synthetic tile/year/month columns needed in queries. + +## Build + +```sh +docker build -t thermograph-lake-iceberg infra/lake-iceberg/ +``` + +## Run (one-shot) + +Config is the backend's lake variables (same names and defaults as +`backend/data/era5lake.py`): `THERMOGRAPH_LAKE_S3_ACCESS_KEY`, +`THERMOGRAPH_LAKE_S3_SECRET_KEY`, and optionally +`THERMOGRAPH_LAKE_S3_ENDPOINT` / `_BUCKET` / `_REGION`. + +```sh +docker run --rm --env-file /path/to/lake.env \ + -v lake-iceberg-state:/state \ + thermograph-lake-iceberg +``` + +Flags: `--dry-run` (list pending tiles, write nothing), `--limit N` (cap +tiles per run), `--batch N` (tiles per Iceberg commit, default 5), +`--catalog-db PATH`. `PYICEBERG_MAX_WORKERS` controls the parquet-footer +read parallelism inside `add_files` (32 is a good value for the ~1000 small +footers per tile). + +Each run prints the row-count reconciliation (records added vs the +manifest's per-tile row counts) and ends with the table's current metadata +JSON path. It also refreshes `metadata/version-hint.text` (Hadoop-catalog +convention) so readers can scan the bare table root without knowing that +path. Transient S3 errors (Contabo intermittently returns SLOW_DOWN / +connection timeouts under load) are retried per batch with backoff +(`--retries`, default 3); a run that still dies resumes at the next batch +boundary on re-run. + +## Scheduling + +Meant to run as a one-shot on prod after extraction batches finish — e.g. a +cron or an ops-cron workflow step invoking the `docker run` above, with the +secrets rendered from the SOPS vault like every other consumer of the +`THERMOGRAPH_LAKE_S3_*` pair. Deliberately **not** wired into the Swarm +stack or CI: the sync is an offline maintenance job, not a service, and each +run costs one footer read per new part file (~1000 per tile), so it belongs +after extraction batches, not on a tight loop. + +## Querying from DuckDB + +No catalog server needed — the version-hint pointer lets DuckDB scan the +table root directly: + +```sql +INSTALL iceberg; LOAD iceberg; +INSTALL httpfs; LOAD httpfs; +CREATE SECRET lake ( + TYPE s3, + KEY_ID '', + SECRET '', + ENDPOINT 'eu2.contabostorage.com', + REGION 'default', + URL_STYLE 'path' +); +SELECT count(*) +FROM iceberg_scan('s3://era5-thermograph/iceberg/era5_daily'); +``` + +To pin an exact table version, scan the metadata JSON path the sync run +printed (`iceberg_scan('s3://…/metadata/.metadata.json')`). + +Partition pruning works on the raw columns, e.g. one tile-month: + +```sql +SELECT lat_idx, lon_idx, avg(tmax) +FROM iceberg_scan('s3://era5-thermograph/iceberg/era5_daily') +WHERE lat_idx BETWEEN 120 AND 131 AND lon_idx BETWEEN 120 AND 131 + AND date >= DATE '2020-07-01' AND date < DATE '2020-08-01' +GROUP BY 1, 2; +``` + +The part files carry no Iceberg field ids; the table sets +`schema.name-mapping.default` so engines resolve columns by name (validated +with pyiceberg 0.11.1 and DuckDB 1.5.5, values byte-identical to the hive +side). + +## Tests + +```sh +pip install -r infra/lake-iceberg/requirements.txt pytest +pytest infra/lake-iceberg/test_sync.py +``` + +Local-filesystem fixtures only (tiny parquet files, `file://` warehouse, +sqlite catalog in tmp) — no network, no credentials. diff --git a/infra/lake-iceberg/requirements.txt b/infra/lake-iceberg/requirements.txt new file mode 100644 index 0000000..22943aa --- /dev/null +++ b/infra/lake-iceberg/requirements.txt @@ -0,0 +1,5 @@ +# Pinned to the versions the sync was validated with (add_files partition +# inference from footer stats, sqlite catalog, path-style S3). +pyiceberg[sql-sqlite,pyarrow,s3fs]==0.11.1 +pyarrow==25.0.0 +s3fs==2026.6.0 diff --git a/infra/lake-iceberg/sync_iceberg.py b/infra/lake-iceberg/sync_iceberg.py new file mode 100644 index 0000000..9bf3921 --- /dev/null +++ b/infra/lake-iceberg/sync_iceberg.py @@ -0,0 +1,451 @@ +"""Maintain an Apache Iceberg table over the ERA5 lake's hive parquet. + +The lake extractor (backend/gen_era5_lake.py) writes plain hive-partitioned +parquet to the era5-thermograph bucket: + + era5/daily/tile=_/year=/month=/part.parquet + +This tool registers those files, in place, into an Iceberg v2 table +``era5_daily`` located at ``s3:///iceberg/era5_daily`` via pyiceberg +``add_files`` — no data is rewritten or copied, the Iceberg metadata simply +points at the existing objects. Everything this tool writes lives under the +``iceberg/`` prefix; the ``era5/`` prefixes are read-only to it by design. + +Partitioning: the part files carry no tile/year/month columns (those live in +the object path), so the table partitions by order-preserving transforms of +real columns instead — ``truncate[12](lat_idx)`` / ``truncate[12](lon_idx)`` +(a 12x12-point tile, so the partition value is ``ti*12``/``tj*12``) and +``year(date)`` / ``month(date)``. Each part file holds exactly one tile and +one calendar month, so every partition value is single-valued per file and +``add_files`` can derive it from the parquet footer statistics alone. + +Incremental and safe under concurrent extraction: ``era5/manifest.parquet`` +is the source of truth for complete tiles (the extractor uploads all of a +tile's objects before adding it to the manifest), so each run diffs the +manifest's tiles against the ``thermograph.synced-tiles`` table property and +adds only the difference. The property is updated in the same commit as the +files, so a run can die at any point and the next one resumes cleanly. + +Catalog: a local pyiceberg SqlCatalog (sqlite, path from +THERMOGRAPH_ICEBERG_CATALOG_DB, default /state/iceberg-catalog.db). The +catalog is a convenience pointer only — if it is lost, the next run finds the +latest ``*.metadata.json`` under the table location and re-registers it. Each +run ends by printing the current metadata JSON path, which DuckDB can read +directly: ``iceberg_scan('s3://.../metadata/NNNNN-....metadata.json')``. + +S3 config comes from the same THERMOGRAPH_LAKE_S3_* variables (and defaults) +as backend/data/era5lake.py; THERMOGRAPH_LAKE_LOCAL_DIR switches to a local +mirror of the lake layout (tests, rehearsal). +""" +import argparse +import os +import re +import sys +import time + +import pyarrow.parquet as pq +from pyiceberg.catalog.sql import SqlCatalog +from pyiceberg.exceptions import NoSuchTableError +from pyiceberg.partitioning import PartitionField, PartitionSpec +from pyiceberg.schema import Schema +from pyiceberg.table.name_mapping import create_mapping_from_schema +from pyiceberg.transforms import MonthTransform, TruncateTransform, YearTransform +from pyiceberg.types import ( + DateType, + DoubleType, + FloatType, + IntegerType, + NestedField, +) + +PREFIX = "era5" +MANIFEST_KEY = f"{PREFIX}/manifest.parquet" +ICEBERG_PREFIX = "iceberg" + +NAMESPACE = "lake" +TABLE_NAME = "era5_daily" +IDENTIFIER = f"{NAMESPACE}.{TABLE_NAME}" +SYNCED_PROP = "thermograph.synced-tiles" + +TILE = 12 # grid points per tile side, as in era5lake.TILE +DEFAULT_BATCH = 5 # tiles per Iceberg commit +DEFAULT_RETRIES = 3 # per-batch retries on transient S3 errors +BACKOFF_SECONDS = 15 # first retry delay, doubled per attempt +DEFAULT_CATALOG_DB = "/state/iceberg-catalog.db" + + +class LakeInvariantError(RuntimeError): + """The lake violated an extractor invariant; never retried.""" + +# The part-file schema as the extractor writes it (verified against the +# bucket): float32 for the heat-index-derived and raw met columns, float64 +# where polars promoted (fmin/sun/feels), int16 doy, int32 grid indices. +# Iceberg has no 16-bit type, so doy maps to int; all fields optional. +SCHEMA = Schema( + NestedField(1, "date", DateType(), required=False), + NestedField(2, "tmax", FloatType(), required=False), + NestedField(3, "tmin", FloatType(), required=False), + NestedField(4, "precip", FloatType(), required=False), + NestedField(5, "wind", FloatType(), required=False), + NestedField(6, "gust", FloatType(), required=False), + NestedField(7, "humid", FloatType(), required=False), + NestedField(8, "fmax", FloatType(), required=False), + NestedField(9, "fmin", DoubleType(), required=False), + NestedField(10, "sun", DoubleType(), required=False), + NestedField(11, "feels", DoubleType(), required=False), + NestedField(12, "doy", IntegerType(), required=False), + NestedField(13, "lat_idx", IntegerType(), required=False), + NestedField(14, "lon_idx", IntegerType(), required=False), +) + +# (tile, year, month): truncate[12] of the grid indices is the tile (values +# are ti*12 / tj*12), year/month come from the date column. All four +# transforms preserve order, which is what lets add_files infer the partition +# value from footer min/max stats (single value per file by construction). +PARTITION_SPEC = PartitionSpec( + PartitionField(source_id=13, field_id=1000, + transform=TruncateTransform(TILE), name="tile_lat"), + PartitionField(source_id=14, field_id=1001, + transform=TruncateTransform(TILE), name="tile_lon"), + PartitionField(source_id=1, field_id=1002, + transform=YearTransform(), name="year"), + PartitionField(source_id=1, field_id=1003, + transform=MonthTransform(), name="month"), +) + + +def s3_config() -> "dict | None": + """Bucket access, or None when unconfigured. Same variable names and + defaults as backend/data/era5lake.py's s3_config (Contabo: path-style + addressing, region is a signing formality).""" + key = os.environ.get("THERMOGRAPH_LAKE_S3_ACCESS_KEY", "") + secret = os.environ.get("THERMOGRAPH_LAKE_S3_SECRET_KEY", "") + if not (key and secret): + return None + return { + "endpoint": os.environ.get("THERMOGRAPH_LAKE_S3_ENDPOINT", + "https://eu2.contabostorage.com"), + "bucket": os.environ.get("THERMOGRAPH_LAKE_S3_BUCKET", "era5-thermograph"), + "region": os.environ.get("THERMOGRAPH_LAKE_S3_REGION", "default"), + "access_key": key, + "secret_key": secret, + } + + +class HiveLake: + """Read-only view of the extracted lake (bucket or local mirror) plus the + iceberg/ prefix the table lives under. ``root`` is an fsspec path prefix + ("" or "/abs/dir"), ``uri`` the equivalent URI prefix pyiceberg + sees ("s3://" or "file:///abs/dir").""" + + def __init__(self, fs, root: str, uri: str, cfg: "dict | None"): + self.fs = fs + self.root = root.rstrip("/") + self.uri = uri.rstrip("/") + self.cfg = cfg + + @property + def warehouse(self) -> str: + return f"{self.uri}/{ICEBERG_PREFIX}" + + def to_uri(self, path: str) -> str: + rel = path.removeprefix(self.root).lstrip("/") + return f"{self.uri}/{rel}" + + def tile_rows(self) -> "dict[str, int]": + """tile -> row count, from the manifest (the source of truth for + complete tiles: the extractor uploads every object of a tile before + its manifest rows appear, so listing incomplete tiles is impossible + here even while extraction is running).""" + with self.fs.open(f"{self.root}/{MANIFEST_KEY}", "rb") as f: + m = pq.read_table(f, columns=["tile", "rows"]) + out: dict[str, int] = {} + for tile, rows in zip(m.column("tile").to_pylist(), + m.column("rows").to_pylist()): + out[tile] = out.get(tile, 0) + rows + return out + + def tile_files(self, tile: str) -> "list[str]": + """URIs of one complete tile's hive part files. A flat find() on the + tile prefix, not a glob: fsspec expands year=*/month=* by walking + every directory level (~1000 LIST calls per tile on S3), while find() + is one paginated listing.""" + prefix = f"{self.root}/{PREFIX}/daily/tile={tile}" + part = re.compile(r"/year=\d+/month=\d+/part\.parquet$") + try: + paths = self.fs.find(prefix) + except FileNotFoundError: + return [] + return sorted(self.to_uri(p) for p in paths if part.search(p)) + + def latest_metadata_uri(self) -> "str | None": + """Newest *.metadata.json under the table location, for re-registering + the table into a fresh catalog db (pyiceberg names them with a + zero-padded, monotonically increasing version prefix).""" + prefix = f"{self.root}/{ICEBERG_PREFIX}/{TABLE_NAME}/metadata" + try: + paths = [p for p in self.fs.find(prefix) + if p.endswith(".metadata.json")] + except FileNotFoundError: + return None + if not paths: + return None + + def version(p: str) -> int: + m = re.match(r"(\d+)", os.path.basename(p)) + return int(m.group(1)) if m else -1 + + return self.to_uri(max(paths, key=version)) + + +def lake_from_env() -> HiveLake: + local = os.environ.get("THERMOGRAPH_LAKE_LOCAL_DIR", "") + if local: + from fsspec.implementations.local import LocalFileSystem + root = os.path.abspath(local) + return HiveLake(LocalFileSystem(), root, f"file://{root}", None) + cfg = s3_config() + if cfg is None: + sys.exit("no lake configured: set THERMOGRAPH_LAKE_S3_ACCESS_KEY/" + "_SECRET_KEY (or THERMOGRAPH_LAKE_LOCAL_DIR for a local mirror)") + import s3fs + fs = s3fs.S3FileSystem( + key=cfg["access_key"], secret=cfg["secret_key"], + endpoint_url=cfg["endpoint"], + client_kwargs={"region_name": cfg["region"]}, + config_kwargs={"s3": {"addressing_style": "path"}}) + return HiveLake(fs, cfg["bucket"], f"s3://{cfg['bucket']}", cfg) + + +def open_catalog(catalog_db: str, lake: HiveLake) -> SqlCatalog: + parent = os.path.dirname(os.path.abspath(catalog_db)) + os.makedirs(parent, exist_ok=True) + props = {"uri": f"sqlite:///{os.path.abspath(catalog_db)}", + "warehouse": lake.warehouse} + if lake.cfg: + props |= { + "py-io-impl": "pyiceberg.io.pyarrow.PyArrowFileIO", + "s3.endpoint": lake.cfg["endpoint"], + "s3.region": lake.cfg["region"], + "s3.access-key-id": lake.cfg["access_key"], + "s3.secret-access-key": lake.cfg["secret_key"], + # Contabo serves buckets on the path, not as a subdomain. + "s3.force-virtual-addressing": "false", + } + return SqlCatalog("thermograph", **props) + + +def ensure_table(catalog: SqlCatalog, lake: HiveLake, create: bool = True): + """Load the table; adopt existing S3 metadata into a fresh catalog db; + otherwise create it. ``create=False`` (dry runs) returns None instead of + creating anything.""" + try: + return catalog.load_table(IDENTIFIER) + except NoSuchTableError: + pass + existing = lake.latest_metadata_uri() + if existing: + catalog.create_namespace_if_not_exists(NAMESPACE) + print(f"registering existing table metadata: {existing}") + return catalog.register_table(IDENTIFIER, existing) + if not create: + return None + catalog.create_namespace_if_not_exists(NAMESPACE) + return catalog.create_table( + IDENTIFIER, + schema=SCHEMA, + location=f"{lake.warehouse}/{TABLE_NAME}", + partition_spec=PARTITION_SPEC, + properties={ + "format-version": "2", + # The part files carry no Iceberg field ids; every reader needs + # this name mapping to resolve columns (DuckDB included). + "schema.name-mapping.default": + create_mapping_from_schema(SCHEMA).model_dump_json(), + }) + + +def synced_tiles(table) -> "set[str]": + if table is None: + return set() + raw = table.properties.get(SYNCED_PROP, "") + return {t for t in raw.split(",") if t} + + +def tile_key(tile: str) -> "tuple[int, int]": + ti, tj = tile.split("_") + return int(ti), int(tj) + + +def pending_tiles(manifest_tiles, synced: "set[str]") -> "list[str]": + return sorted(set(manifest_tiles) - synced, key=tile_key) + + +def sync_batch(table, lake: HiveLake, chunk: "list[str]", + retries: int = DEFAULT_RETRIES) -> "list[str]": + """Commit one batch of tiles (files + synced-tiles property, atomically), + retrying transient S3 failures with backoff. Each attempt re-reads the + synced set from the (refreshed) table, so an attempt that committed but + errored afterwards is not re-added. Returns the tiles actually added.""" + for attempt in range(retries + 1): + try: + table.refresh() + left = [t for t in chunk if t not in synced_tiles(table)] + if not left: + return [] + files = [] + for t in left: + tf = lake.tile_files(t) + if not tf: + raise LakeInvariantError( + f"tile {t} is in the manifest but has no daily part " + f"files — refusing to mark it synced") + files.extend(tf) + new_prop = ",".join( + sorted(synced_tiles(table) | set(left), key=tile_key)) + with table.transaction() as tx: + tx.add_files( + file_paths=files, + snapshot_properties={"thermograph.tiles-added": + ",".join(left)}) + tx.set_properties({SYNCED_PROP: new_prop}) + return left + except LakeInvariantError: + raise + except Exception as e: # noqa: BLE001 - S3 throttling has many faces + if attempt >= retries: + raise + wait = BACKOFF_SECONDS * 2 ** attempt + print(f"batch {chunk[0]}..{chunk[-1]} attempt {attempt + 1} " + f"failed ({e}); retrying in {wait}s", file=sys.stderr) + time.sleep(wait) + return [] # unreachable + + +def write_version_hint(lake: HiveLake, table) -> "str | None": + """Hadoop-catalog-style pointer next to the metadata files so engines can + scan the bare table root — DuckDB: iceberg_scan('s3://…/era5_daily') — + without knowing the exact metadata path. Content is the current metadata + file's basename minus '.metadata.json' (DuckDB's default + version_name_format tries '%s.metadata.json', which resolves it). A + convenience pointer only: refreshed every run, so a failed write just + lags one version.""" + hint = os.path.basename(table.metadata_location) \ + .removesuffix(".metadata.json") + path = f"{lake.root}/{ICEBERG_PREFIX}/{TABLE_NAME}/metadata/version-hint.text" + for attempt in range(3): + try: + with lake.fs.open(path, "wb") as f: + f.write(hint.encode()) + return hint + except Exception as e: # noqa: BLE001 + if attempt == 2: + print(f"WARNING: version-hint.text write failed ({e}); " + f"table-root scans will lag one version", file=sys.stderr) + return None + time.sleep(BACKOFF_SECONDS) + return None + + +def run(lake: HiveLake, catalog_db: str, *, limit: "int | None" = None, + batch: int = DEFAULT_BATCH, retries: int = DEFAULT_RETRIES, + dry_run: bool = False) -> dict: + tile_rows = lake.tile_rows() + catalog = open_catalog(catalog_db, lake) + table = ensure_table(catalog, lake, create=not dry_run) + done = synced_tiles(table) + todo = pending_tiles(tile_rows, done) + if limit: + todo = todo[:limit] + print(f"manifest tiles={len(tile_rows)} synced={len(done)} " + f"pending this run={len(todo)}") + + added_files = added_records = 0 + mismatches = [] + if dry_run: + for t in todo: + files = lake.tile_files(t) + print(f"[dry-run] tile {t}: {len(files)} part files, " + f"{tile_rows[t]} manifest rows") + else: + for k in range(0, len(todo), batch): + chunk = todo[k:k + batch] + # One commit per batch: the files and the synced-tiles marker land + # atomically, so an interrupted run never half-syncs a tile. + added = sync_batch(table, lake, chunk, retries) + done |= set(chunk) + if not added: + print(f"[{min(k + batch, len(todo))}/{len(todo)}] tiles " + f"{chunk[0]}..{chunk[-1]}: already synced") + continue + summary = table.current_snapshot().summary + got = int(summary.get("added-records", "0")) + expected = sum(tile_rows[t] for t in added) + n_files = int(summary.get("added-data-files", "0")) + added_files += n_files + added_records += got + status = "ok" if got == expected else "MISMATCH" + if got != expected: + mismatches.append((added, expected, got)) + print(f"[{min(k + batch, len(todo))}/{len(todo)}] tiles " + f"{added[0]}..{added[-1]}: +{n_files} files, " + f"+{got} records (manifest says {expected}: {status})") + + result = {"pending": todo, "added_files": added_files, + "added_records": added_records, "mismatches": mismatches, + "metadata_location": None, "total_records": 0, + "total_files": 0} + if table is not None: + snap = table.current_snapshot() + if snap is not None: + result["total_records"] = int(snap.summary.get("total-records", "0")) + result["total_files"] = int(snap.summary.get("total-data-files", "0")) + result["metadata_location"] = table.metadata_location + if not dry_run: + result["version_hint"] = write_version_hint(lake, table) + expected_total = sum(tile_rows[t] for t in done if t in tile_rows) + print(f"table: {result['total_files']} data files, " + f"{result['total_records']} records " + f"(manifest rows for synced tiles: {expected_total})") + print(f"metadata: {table.metadata_location}") + if lake.cfg: + print("read it from DuckDB (no catalog needed):\n" + " LOAD iceberg; LOAD httpfs; -- plus an S3 secret for the " + "endpoint\n" + f" SELECT count(*) FROM iceberg_scan(" + f"'{lake.warehouse}/{TABLE_NAME}');\n" + f" -- or pin the exact version: iceberg_scan(" + f"'{table.metadata_location}')") + if mismatches: + print(f"WARNING: {len(mismatches)} batch(es) added a record count " + f"that differs from the manifest — investigate before trusting " + f"the table", file=sys.stderr) + return result + + +def main() -> int: + ap = argparse.ArgumentParser( + description="Register new ERA5 lake tiles into the Iceberg table") + ap.add_argument("--catalog-db", + default=os.environ.get("THERMOGRAPH_ICEBERG_CATALOG_DB", + DEFAULT_CATALOG_DB), + help="sqlite catalog path (default: %(default)s)") + ap.add_argument("--batch", type=int, default=DEFAULT_BATCH, + help="tiles per Iceberg commit (default: %(default)s)") + ap.add_argument("--limit", type=int, + help="sync at most this many new tiles this run") + ap.add_argument("--retries", type=int, default=DEFAULT_RETRIES, + help="per-batch retries on transient S3 errors " + "(default: %(default)s)") + ap.add_argument("--dry-run", action="store_true", + help="list pending tiles; write nothing") + args = ap.parse_args() + + lake = lake_from_env() + result = run(lake, args.catalog_db, limit=args.limit, batch=args.batch, + retries=args.retries, dry_run=args.dry_run) + return 1 if result["mismatches"] else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/infra/lake-iceberg/test_sync.py b/infra/lake-iceberg/test_sync.py new file mode 100644 index 0000000..fc29c84 --- /dev/null +++ b/infra/lake-iceberg/test_sync.py @@ -0,0 +1,247 @@ +"""Unit tests for the Iceberg lake sync: incremental tile diffing and +add_files planning against a local filesystem fixture (a miniature lake in a +tmp dir, a file:// warehouse, a sqlite catalog). No network.""" +import datetime +import os +import sys + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from fsspec.implementations.local import LocalFileSystem + +sys.path.insert(0, os.path.dirname(__file__)) +import sync_iceberg as si # noqa: E402 + +# Mirrors the real part-file schema (float32/float64 split, int16 doy, +# all-null float64 sun) so the compatibility check is exercised for real. +PART_SCHEMA = pa.schema([ + ("date", pa.date32()), + ("tmax", pa.float32()), + ("tmin", pa.float32()), + ("precip", pa.float32()), + ("wind", pa.float32()), + ("gust", pa.float32()), + ("humid", pa.float32()), + ("fmax", pa.float32()), + ("fmin", pa.float64()), + ("sun", pa.float64()), + ("feels", pa.float64()), + ("doy", pa.int16()), + ("lat_idx", pa.int32()), + ("lon_idx", pa.int32()), +]) + +POINTS = 2 # points per fixture tile +DAYS = 3 # rows per point per month + + +def write_tile(root, ti, tj, months): + """Write one tile's hive part files; return its total row count.""" + total = 0 + for year, month in months: + rows = {name: [] for name in PART_SCHEMA.names} + for p in range(POINTS): + i, j = ti * si.TILE + p, tj * si.TILE + p + for d in range(DAYS): + day = datetime.date(year, month, d + 1) + rows["date"].append(day) + for col in ("tmax", "tmin", "precip", "wind", "gust", + "humid", "fmax", "fmin", "feels"): + rows[col].append(float(10 * p + d)) + rows["sun"].append(None) + rows["doy"].append(day.timetuple().tm_yday) + rows["lat_idx"].append(i) + rows["lon_idx"].append(j) + part = pa.table(rows, schema=PART_SCHEMA) + path = os.path.join(root, si.PREFIX, "daily", f"tile={ti}_{tj}", + f"year={year}", f"month={month:02d}", + "part.parquet") + os.makedirs(os.path.dirname(path), exist_ok=True) + pq.write_table(part, path) + total += part.num_rows + return total + + +def write_manifest(root, tile_rows): + """Manifest with one row per (tile, point) like the extractor writes; + the sync only reads tile+rows.""" + tiles, rows = [], [] + for tile, total in tile_rows.items(): + per_point = total // POINTS + for _ in range(POINTS): + tiles.append(tile) + rows.append(per_point) + m = pa.table({ + "lat_idx": pa.array([0] * len(tiles), pa.int64()), + "lon_idx": pa.array([0] * len(tiles), pa.int64()), + "lat": pa.array([0.0] * len(tiles)), + "lon": pa.array([0.0] * len(tiles)), + "tile": pa.array(tiles, pa.large_string()), + "rows": pa.array(rows, pa.int64()), + "date_min": pa.array(["1940-01-01"] * len(tiles), pa.large_string()), + "date_max": pa.array(["2026-01-01"] * len(tiles), pa.large_string()), + }) + path = os.path.join(root, si.MANIFEST_KEY) + os.makedirs(os.path.dirname(path), exist_ok=True) + pq.write_table(m, path) + + +@pytest.fixture +def lake(tmp_path): + root = str(tmp_path / "lake") + counts = { + "10_10": write_tile(root, 10, 10, [(2000, 1), (2000, 2)]), + "10_11": write_tile(root, 10, 11, [(2000, 1)]), + } + write_manifest(root, counts) + hl = si.HiveLake(LocalFileSystem(), root, f"file://{root}", None) + hl.fixture_counts = counts + hl.fixture_root = root + return hl + + +@pytest.fixture +def catalog_db(tmp_path): + return str(tmp_path / "catalog.db") + + +def test_pending_tiles_diff(): + manifest = {"10_10": 1, "10_11": 1, "2_3": 1} + assert si.pending_tiles(manifest, set()) == ["2_3", "10_10", "10_11"] + assert si.pending_tiles(manifest, {"10_10"}) == ["2_3", "10_11"] + assert si.pending_tiles(manifest, set(manifest)) == [] + # numeric, not lexicographic, tile ordering + assert si.pending_tiles({"10_2": 1, "9_11": 1}, set()) == ["9_11", "10_2"] + + +def test_initial_sync_registers_all_files_in_place(lake, catalog_db): + result = si.run(lake, catalog_db, batch=1) + expected = sum(lake.fixture_counts.values()) + assert result["added_records"] == expected + assert result["total_records"] == expected + assert result["added_files"] == 3 # 2 months + 1 month + assert not result["mismatches"] + + table = si.open_catalog(catalog_db, lake).load_table(si.IDENTIFIER) + assert si.synced_tiles(table) == set(lake.fixture_counts) + # add_files, not a rewrite: every data file is an original hive part file + files = table.inspect.files().column("file_path").to_pylist() + assert len(files) == 3 + assert all(f"/{si.PREFIX}/daily/tile=" in f for f in files) + # nothing was written outside iceberg/ (the hive tree is untouched, the + # table tree holds only metadata) + fs = LocalFileSystem() + written = fs.find(os.path.join(lake.fixture_root, si.ICEBERG_PREFIX)) + assert written and all("/metadata/" in p for p in written) + + +def test_rerun_is_noop(lake, catalog_db): + first = si.run(lake, catalog_db) + again = si.run(lake, catalog_db) + assert again["pending"] == [] + assert again["added_records"] == 0 + assert again["total_records"] == first["total_records"] + assert again["metadata_location"] == first["metadata_location"] + + +def test_incremental_new_tile_only(lake, catalog_db): + si.run(lake, catalog_db) + counts = dict(lake.fixture_counts) + counts["11_10"] = write_tile(lake.fixture_root, 11, 10, [(2001, 6)]) + write_manifest(lake.fixture_root, counts) + + result = si.run(lake, catalog_db) + assert result["pending"] == ["11_10"] + assert result["added_records"] == counts["11_10"] + assert result["total_records"] == sum(counts.values()) + + +def test_version_hint_tracks_current_metadata(lake, catalog_db): + first = si.run(lake, catalog_db) + hint_path = os.path.join(lake.fixture_root, si.ICEBERG_PREFIX, + si.TABLE_NAME, "metadata", "version-hint.text") + with open(hint_path) as f: + assert f.read() == os.path.basename( + first["metadata_location"]).removesuffix(".metadata.json") + + counts = dict(lake.fixture_counts) + counts["11_10"] = write_tile(lake.fixture_root, 11, 10, [(2001, 6)]) + write_manifest(lake.fixture_root, counts) + second = si.run(lake, catalog_db) + assert second["metadata_location"] != first["metadata_location"] + with open(hint_path) as f: + assert f.read() == os.path.basename( + second["metadata_location"]).removesuffix(".metadata.json") + + +def test_partition_spec_and_pruned_scan(lake, catalog_db): + si.run(lake, catalog_db) + table = si.open_catalog(catalog_db, lake).load_table(si.IDENTIFIER) + assert [f.name for f in table.spec().fields] == \ + ["tile_lat", "tile_lon", "year", "month"] + + parts = table.inspect.partitions() + keys = {(r["partition"]["tile_lat"], r["partition"]["tile_lon"]) + for r in parts.to_pylist()} + assert keys == {(120, 120), (120, 132)} # ti*12, tj*12 + + # a tile-shaped predicate resolves to that tile's rows only + scan = table.scan(row_filter="lat_idx >= 120 and lat_idx < 132 " + "and lon_idx >= 132 and lon_idx < 144") + assert scan.to_arrow().num_rows == lake.fixture_counts["10_11"] + # and file planning pruned to that tile's single part file + assert len(list(scan.plan_files())) == 1 + + +def test_lost_catalog_db_reregisters_without_readding(lake, catalog_db): + first = si.run(lake, catalog_db) + os.remove(catalog_db) # simulate losing local state + again = si.run(lake, catalog_db) + # synced-tiles lives in table metadata, so nothing is re-added + assert again["pending"] == [] + assert again["added_records"] == 0 + assert again["total_records"] == first["total_records"] + + +def test_dry_run_writes_nothing(lake, catalog_db): + result = si.run(lake, catalog_db, dry_run=True) + assert sorted(result["pending"]) == sorted(lake.fixture_counts) + assert result["metadata_location"] is None + assert not os.path.exists( + os.path.join(lake.fixture_root, si.ICEBERG_PREFIX)) + + +def test_manifest_tile_without_files_refuses(lake, catalog_db): + counts = dict(lake.fixture_counts) + counts["50_50"] = 6 # in the manifest, no files + write_manifest(lake.fixture_root, counts) + # an invariant break aborts immediately: no retries, nothing marked synced + with pytest.raises(si.LakeInvariantError, match="50_50"): + si.run(lake, catalog_db) + + +def test_transient_failure_retries_and_succeeds(lake, catalog_db, monkeypatch): + monkeypatch.setattr(si, "BACKOFF_SECONDS", 0) + fails = {"n": 2} + original = si.HiveLake.tile_files + + def flaky(self, tile): + if fails["n"]: + fails["n"] -= 1 + raise OSError("AWS Error SLOW_DOWN during CompleteMultipartUpload") + return original(self, tile) + + monkeypatch.setattr(si.HiveLake, "tile_files", flaky) + result = si.run(lake, catalog_db, retries=2) + assert result["added_records"] == sum(lake.fixture_counts.values()) + assert not result["mismatches"] + + +def test_transient_failure_exhausts_retries(lake, catalog_db, monkeypatch): + monkeypatch.setattr(si, "BACKOFF_SECONDS", 0) + monkeypatch.setattr( + si.HiveLake, "tile_files", + lambda self, tile: (_ for _ in ()).throw(OSError("SLOW_DOWN"))) + with pytest.raises(OSError): + si.run(lake, catalog_db, retries=1)