"""Task 06 - Convert the CSV corpus to PostgreSQL. Loads the same logical schema as task 05 into the PostgreSQL 16 database named by the TRIBO_PG_DSN environment variable: COPY FROM STDIN, RANGE partitioning of the two bulk tables by track_id (12 partitions of 120 tracks), constraints and indexes created after the load (variant B: enforced foreign keys everywhere - docs/research/Tribology_FK_Architecture.html), track_summary as a materialized view (algorithm fixed in docs/rules/db-sql-schema.md section 5, cross-checked against common/track_summary.py). Host CPU/RAM during the load is sampled from the database host's Prometheus (TRIBO_PROM_URL). Appends live (and, when pg_dump is available, dump) sizes to storage_sizes.csv and writes ./out/.done/06.ok. Spec: docs/specs/06_convert_postgresql.md. """ from __future__ import annotations import logging import shutil import subprocess import sys import time from pathlib import Path import psycopg from common.corpus import CorpusReader, ValidationError from common.monitoring import get_prom_url, host_window_metrics from common.pg import get_dsn from common.pipeline import ( check_dependencies, load_lab_config, parse_task_args, process_metrics, remove_stale_marker, write_marker, ) from common.relational import COLUMNS, TABLES, expected_counts, stream_rows from common.storage_sizes import update_storage_sizes from common.track_summary import summarize_track logger = logging.getLogger(__name__) TASK_ID = "06" DEPENDS_ON = ["01", "03"] BATCH_ROWS = 50000 # rows per COPY chunk (mirrors the 50k batching of spec 05) PARTITIONS = 12 # RANGE partitions of 120 tracks each, spec 06 TRACKS_PER_PARTITION = 120 SUMMARY_PARITY_TRACKS = (1, 720, 1440) # matview cross-checked vs common/track_summary.py PARITY_TOLERANCE = 1e-9 # cross-engine float tolerance, test-pipeline-validation # Float columns are DOUBLE PRECISION, not REAL: the corpus is float64 and # REAL (float4) would break the 1e-9 cross-engine result parity that the # whole benchmark depends on (db-sql-schema section 6; spec updated). CREATE_TABLES = [ "CREATE TABLE instruments (instrument_id INTEGER NOT NULL, instrument_code TEXT NOT NULL, name TEXT NOT NULL, role TEXT NOT NULL)", "CREATE TABLE batches (batch_id INTEGER NOT NULL, batch_code TEXT NOT NULL, pt_gun_tilt_deg INTEGER NOT NULL, au_gun_tilt_deg INTEGER NOT NULL, pt_power_w INTEGER NOT NULL, au_power_w INTEGER NOT NULL, pt_discharge_v INTEGER NOT NULL, au_discharge_v INTEGER NOT NULL, deposition_date DATE NOT NULL)", "CREATE TABLE runs (run_id INTEGER NOT NULL, run_code TEXT NOT NULL, environment TEXT NOT NULL, date DATE NOT NULL, plates INTEGER NOT NULL, operator TEXT NOT NULL)", "CREATE TABLE wafers (wafer_id INTEGER NOT NULL, wafer_code TEXT NOT NULL, batch_id INTEGER NOT NULL, wafer_index INTEGER NOT NULL, deposition_date DATE NOT NULL, coupons INTEGER NOT NULL, friction_coupons INTEGER NOT NULL, reserve_coupons INTEGER NOT NULL)", "CREATE TABLE coupons (coupon_id INTEGER NOT NULL, coupon_code TEXT NOT NULL, wafer_id INTEGER NOT NULL, batch_id INTEGER NOT NULL, grid_row INTEGER NOT NULL, grid_col INTEGER NOT NULL, thickness_um DOUBLE PRECISION NOT NULL, ra_nm DOUBLE PRECISION NOT NULL, au_wtpct_mean DOUBLE PRECISION NOT NULL, run_id INTEGER, plate INTEGER, probe INTEGER, square INTEGER)", "CREATE TABLE tracks (track_id INTEGER NOT NULL, track_code TEXT NOT NULL, coupon_id INTEGER NOT NULL, run_id INTEGER NOT NULL, environment TEXT NOT NULL, load_mn DOUBLE PRECISION NOT NULL, stroke_mm DOUBLE PRECISION NOT NULL, speed_mm_s DOUBLE PRECISION NOT NULL, counterface_id TEXT NOT NULL, started_at TIMESTAMPTZ NOT NULL)", "CREATE TABLE simtra_profiles (batch_id INTEGER NOT NULL, row_no INTEGER NOT NULL, angle_deg DOUBLE PRECISION NOT NULL, energy_ev DOUBLE PRECISION NOT NULL, pt_flux DOUBLE PRECISION NOT NULL, au_flux DOUBLE PRECISION NOT NULL)", "CREATE TABLE xrf_points (coupon_id INTEGER NOT NULL, grid_x INTEGER NOT NULL, grid_y INTEGER NOT NULL, pt_wtpct DOUBLE PRECISION NOT NULL, au_wtpct DOUBLE PRECISION NOT NULL)", "CREATE TABLE profilometry_points (coupon_id INTEGER NOT NULL, grid_x INTEGER NOT NULL, grid_y INTEGER NOT NULL, thickness_um DOUBLE PRECISION NOT NULL)", "CREATE TABLE nanoindentation (coupon_id INTEGER NOT NULL, indent_id INTEGER NOT NULL, x_um DOUBLE PRECISION NOT NULL, y_um DOUBLE PRECISION NOT NULL, hardness_gpa DOUBLE PRECISION NOT NULL, reduced_modulus_gpa DOUBLE PRECISION NOT NULL, max_load_mn DOUBLE PRECISION NOT NULL)", "CREATE TABLE afm (coupon_id INTEGER NOT NULL, ra_nm DOUBLE PRECISION NOT NULL, rq_nm DOUBLE PRECISION NOT NULL, image_file TEXT NOT NULL)", "CREATE TABLE friction_cycles (track_id INTEGER NOT NULL, cycle INTEGER NOT NULL, cof DOUBLE PRECISION NOT NULL) PARTITION BY RANGE (track_id)", "CREATE TABLE friction_loop_points (track_id INTEGER NOT NULL, cycle INTEGER NOT NULL, pt INTEGER NOT NULL, position_um DOUBLE PRECISION NOT NULL, friction_force_mn DOUBLE PRECISION NOT NULL) PARTITION BY RANGE (track_id)", "CREATE TABLE wear (track_id INTEGER NOT NULL, wear_volume_um3 DOUBLE PRECISION NOT NULL, k_archard DOUBLE PRECISION NOT NULL, sliding_distance_m DOUBLE PRECISION NOT NULL)", ] CONSTRAINTS = [ "ALTER TABLE instruments ADD PRIMARY KEY (instrument_id)", "ALTER TABLE instruments ADD UNIQUE (instrument_code)", "ALTER TABLE batches ADD PRIMARY KEY (batch_id)", "ALTER TABLE batches ADD UNIQUE (batch_code)", "ALTER TABLE runs ADD PRIMARY KEY (run_id)", "ALTER TABLE runs ADD UNIQUE (run_code)", "ALTER TABLE wafers ADD PRIMARY KEY (wafer_id)", "ALTER TABLE wafers ADD UNIQUE (wafer_code)", "ALTER TABLE coupons ADD PRIMARY KEY (coupon_id)", "ALTER TABLE coupons ADD UNIQUE (coupon_code)", "ALTER TABLE tracks ADD PRIMARY KEY (track_id)", "ALTER TABLE tracks ADD UNIQUE (track_code)", "ALTER TABLE simtra_profiles ADD PRIMARY KEY (batch_id, row_no)", "ALTER TABLE xrf_points ADD PRIMARY KEY (coupon_id, grid_x, grid_y)", "ALTER TABLE profilometry_points ADD PRIMARY KEY (coupon_id, grid_x, grid_y)", "ALTER TABLE nanoindentation ADD PRIMARY KEY (coupon_id, indent_id)", "ALTER TABLE afm ADD PRIMARY KEY (coupon_id)", "ALTER TABLE friction_cycles ADD PRIMARY KEY (track_id, cycle)", "ALTER TABLE friction_loop_points ADD PRIMARY KEY (track_id, cycle, pt)", "ALTER TABLE wear ADD PRIMARY KEY (track_id)", "ALTER TABLE wafers ADD FOREIGN KEY (batch_id) REFERENCES batches", "ALTER TABLE coupons ADD FOREIGN KEY (wafer_id) REFERENCES wafers", "ALTER TABLE coupons ADD FOREIGN KEY (batch_id) REFERENCES batches", "ALTER TABLE coupons ADD FOREIGN KEY (run_id) REFERENCES runs", "ALTER TABLE tracks ADD FOREIGN KEY (coupon_id) REFERENCES coupons", "ALTER TABLE tracks ADD FOREIGN KEY (run_id) REFERENCES runs", "ALTER TABLE simtra_profiles ADD FOREIGN KEY (batch_id) REFERENCES batches", "ALTER TABLE xrf_points ADD FOREIGN KEY (coupon_id) REFERENCES coupons", "ALTER TABLE profilometry_points ADD FOREIGN KEY (coupon_id) REFERENCES coupons", "ALTER TABLE nanoindentation ADD FOREIGN KEY (coupon_id) REFERENCES coupons", "ALTER TABLE afm ADD FOREIGN KEY (coupon_id) REFERENCES coupons", "ALTER TABLE friction_cycles ADD FOREIGN KEY (track_id) REFERENCES tracks", "ALTER TABLE friction_loop_points ADD FOREIGN KEY (track_id) REFERENCES tracks", "ALTER TABLE wear ADD FOREIGN KEY (track_id) REFERENCES tracks", ] INDEX_DDL = [ "CREATE INDEX ix_coupons_au_wtpct_mean ON coupons(au_wtpct_mean)", # Q2 selective filter "CREATE INDEX ix_tracks_run_id ON tracks(run_id)", # Q5/Q7 joins to runs "CREATE INDEX ix_tracks_coupon_id ON tracks(coupon_id)", # Q3/Q6 joins to coupons "CREATE INDEX ix_runs_environment ON runs(environment)", # Q7 grouping "CREATE INDEX ix_tracks_environment_load ON tracks(environment, load_mn)", # Q4 filter # Q5 pruning within partitions; cascades to every partition (spec 06). # nanoindentation(coupon_id) from the spec is omitted: the PK prefix serves it. "CREATE INDEX ix_brin_friction_cycles_cycle ON friction_cycles USING BRIN (cycle)", ] # Materialized view semantics fixed in docs/rules/db-sql-schema.md section 5. MATVIEW_SQL = """ CREATE MATERIALIZED VIEW track_summary AS WITH tail AS ( SELECT track_id, avg(cof) AS m, stddev_pop(cof) AS s FROM friction_cycles WHERE cycle >= 501 GROUP BY track_id ), run_in AS ( SELECT fc.track_id, COALESCE(MIN(fc.cycle) FILTER (WHERE abs(fc.cof - t.m) < 2 * t.s), 500) AS run_in_cycles FROM friction_cycles fc JOIN tail t USING (track_id) GROUP BY fc.track_id ) SELECT r.track_id, avg(fc.cof) AS cof_ss_mean, stddev_pop(fc.cof) AS cof_ss_std, r.run_in_cycles FROM run_in r JOIN friction_cycles fc ON fc.track_id = r.track_id AND fc.cycle > r.run_in_cycles GROUP BY r.track_id, r.run_in_cycles """ def partition_ddl() -> list[str]: ddl = [] for parent in ("friction_cycles", "friction_loop_points"): for i in range(PARTITIONS): lo = i * TRACKS_PER_PARTITION + 1 hi = lo + TRACKS_PER_PARTITION ddl.append(f"CREATE TABLE {parent}_p{i + 1:02d} PARTITION OF {parent} FOR VALUES FROM ({lo}) TO ({hi})") return ddl class PgLoader: def __init__(self, conn: psycopg.Connection) -> None: self.conn = conn self.buffers: dict[str, list[tuple]] = {t: [] for t in TABLES} self.summary_parity: dict[int, tuple] = {} def create_schema(self) -> None: with self.conn.cursor() as cur: cur.execute("DROP MATERIALIZED VIEW IF EXISTS track_summary") for table in reversed(TABLES): if table != "track_summary": cur.execute(f"DROP TABLE IF EXISTS {table} CASCADE") for ddl in CREATE_TABLES + partition_ddl(): cur.execute(ddl) self.conn.commit() logger.info("schema created: %d tables, %d partitions", len(CREATE_TABLES), 2 * PARTITIONS) def load(self, cfg: dict, reader: CorpusReader) -> None: pending = 0 for table, row in stream_rows(cfg, reader): if table == "track_summary": # computed server-side as a materialized view; keep a sample # of the Python values for the cross-engine parity check if row[0] in SUMMARY_PARITY_TRACKS: self.summary_parity[row[0]] = row continue self.buffers[table].append(row) pending += 1 if pending >= BATCH_ROWS: self.flush_all() pending = 0 self.flush_all() def flush_all(self) -> None: for table in TABLES: buf = self.buffers.get(table) if not buf: continue cols = ", ".join(COLUMNS[table]) with self.conn.cursor() as cur: with cur.copy(f"COPY {table} ({cols}) FROM STDIN") as copy: for row in buf: copy.write_row(row) buf.clear() self.conn.commit() def finalize(self) -> None: with self.conn.cursor() as cur: for ddl in CONSTRAINTS: cur.execute(ddl) for ddl in INDEX_DDL: cur.execute(ddl) cur.execute(MATVIEW_SQL) cur.execute("CREATE UNIQUE INDEX ix_track_summary_track_id ON track_summary(track_id)") # Q2/Q3 joins self.conn.commit() logger.info("constraints, indexes and track_summary materialized view created") def validate(self, cfg: dict) -> dict[str, int]: counts: dict[str, int] = {} with self.conn.cursor() as cur: for table, exp in expected_counts(cfg).items(): got = cur.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] counts[table] = got if got != exp: raise ValidationError(f"row count mismatch: {table}: loaded {got} != expected {exp}") reserve = cur.execute("SELECT COUNT(*) FROM coupons WHERE run_id IS NULL").fetchone()[0] if reserve != cfg["friction_assignment"]["reserve_coupons_total"]: raise ValidationError(f"reserve coupons: {reserve} != expected") for track_id, py_row in self.summary_parity.items(): db_row = cur.execute( "SELECT cof_ss_mean, cof_ss_std, run_in_cycles FROM track_summary WHERE track_id = %s", (track_id,), ).fetchone() if db_row is None: raise ValidationError(f"track_summary missing track {track_id}") if (abs(db_row[0] - py_row[1]) > PARITY_TOLERANCE or abs(db_row[1] - py_row[2]) > PARITY_TOLERANCE or db_row[2] != py_row[3]): raise ValidationError( f"track_summary parity failed for track {track_id}: db={db_row} py={py_row[1:]}" ) logger.info( "all %d table counts match, %d reserve coupons, matview parity ok for tracks %s", len(counts), reserve, list(self.summary_parity), ) return counts def table_sizes(self) -> dict[str, int]: sizes: dict[str, int] = {} with self.conn.cursor() as cur: for table in TABLES: total = cur.execute( """SELECT COALESCE(pg_total_relation_size(%s::regclass), 0) + COALESCE((SELECT sum(pg_total_relation_size(inhrelid)) FROM pg_inherits WHERE inhparent = %s::regclass), 0)""", (table, table), ).fetchone()[0] sizes[table] = int(total) return sizes def make_dump(dsn: str, dump_path: Path) -> int: pg_dump = shutil.which("pg_dump") if not pg_dump: logger.warning("pg_dump skipped (client not installed on this host; produce the dump on the database host)") return 0 dump_path.parent.mkdir(parents=True, exist_ok=True) subprocess.run([pg_dump, "--format=custom", f"--file={dump_path}", dsn], check=True) return dump_path.stat().st_size def main() -> int: args = parse_task_args("Task 06: convert the CSV corpus to PostgreSQL") out_root: Path = args.out_root dsn = get_dsn() try: check_dependencies(out_root, DEPENDS_ON) remove_stale_marker(out_root, TASK_ID) cfg = load_lab_config(out_root) t_start = time.time() with psycopg.connect(dsn) as conn: database = conn.info.dbname logger.info("connected to database %r on %s", database, conn.info.host) loader = PgLoader(conn) loader.create_schema() loader.load(cfg, CorpusReader(out_root / "csv")) load_done = time.time() loader.finalize() counts = loader.validate(cfg) sizes = loader.table_sizes() with psycopg.connect(dsn, autocommit=True) as conn: conn.execute("VACUUM ANALYZE") t_end = time.time() live_bytes = sum(sizes.values()) variants = [("live", live_bytes)] dump_bytes = make_dump(dsn, out_root / "pg" / "tribo.dump") if dump_bytes: variants.append(("dump", dump_bytes)) update_storage_sizes(out_root / "bench", "pg", variants) entries = { "database": database, "tables": len(counts), "total_rows": sum(counts.values()), "live_bytes": live_bytes, "friction_cycles_bytes": sizes["friction_cycles"], "friction_loop_points_bytes": sizes["friction_loop_points"], "copy_seconds": round(load_done - t_start, 1), "total_seconds": round(t_end - t_start, 1), "dump": dump_bytes if dump_bytes else "skipped (pg_dump unavailable on this host)", } entries.update(process_metrics()) # client-side converter cost prom_url = get_prom_url() if prom_url: metrics = host_window_metrics(prom_url, t_start, t_end) if metrics: entries.update(metrics) logger.info("db host during load: %s", metrics) marker_path = write_marker(out_root, TASK_ID, entries) except Exception: logger.critical("task %s failed", TASK_ID, exc_info=True) return 1 print(f"task 06 ok: database {database} ({live_bytes / (1024 * 1024):.1f} MiB live, " f"{len(counts)} tables, {sum(counts.values())} rows, {entries['total_seconds']} s)") print(f"bulk rows: cycles={counts['friction_cycles']} loop_points={counts['friction_loop_points']} " f"track_summary={counts['track_summary']}") if prom_url and "host_cpu_avg_pct" in entries: print(f"db host: cpu avg {entries['host_cpu_avg_pct']}% max {entries['host_cpu_max_pct']}%, " f"ram peak {entries['host_ram_used_peak_mb']} MB (delta {entries['host_ram_used_delta_mb']} MB)") print(f"marker: {marker_path}") return 0 if __name__ == "__main__": sys.exit(main())