"""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())