f87e22be51
All checks were successful
PR build (required check) / changes (pull_request) Successful in 9s
secrets-guard / encrypted (pull_request) Successful in 6s
PR build (required check) / build-frontend (pull_request) Has been skipped
PR build (required check) / validate-observability (pull_request) Has been skipped
PR build (required check) / build-backend (pull_request) Successful in 44s
PR build (required check) / gate (pull_request) Successful in 2s
Registers the hive part files under era5/daily into an Apache Iceberg v2 table at iceberg/era5_daily via pyiceberg add_files -- the metadata points at the existing parquet in place, no data rewrite. Incremental: each run diffs era5/manifest.parquet (the completeness signal while extraction is running) against the thermograph.synced-tiles table property, updated in the same commit as the files, so runs are idempotent and resume at batch boundaries. Partitioned by truncate[12](lat_idx), truncate[12](lon_idx), year(date), month(date) -- order-preserving transforms add_files derives from footer stats; queries prune on plain column predicates. Local sqlite catalog that re-registers from the latest metadata JSON if lost; version-hint.text refreshed per run so DuckDB can iceberg_scan the bare table root. Per-batch retries with backoff for Contabo SLOW_DOWN and slow-transfer timeouts. Slim pinned one-shot image; local-filesystem pytest suite (no network).
247 lines
9.4 KiB
Python
247 lines
9.4 KiB
Python
"""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)
|