LabDataStorageEvaluation/make_queries.py

"""Task 08 - Benchmark query definitions (Q1-Q7 x 5 formats).

Writes the runnable query files under ./out/bench/queries/ (SQL for
SQLite/PostgreSQL, SPARQL text + python drivers for RDF, python drivers
for CSV/JSON-LD), then produces the canonical expected results: executed
on PostgreSQL, cross-validated against SQLite row-by-row within 1e-9, and
stored as ./out/bench/expected/q<N>.rows.csv + q<N>.sha256. CSV, JSON and
RDF implementations are spot-checked on Q1. Writes ./out/.done/08.ok.
Spec: docs/specs/08_benchmark_queries.md.
"""

from __future__ import annotations

import logging
import sqlite3
import sys
from pathlib import Path

import psycopg

from common.pg import get_dsn
from common.pipeline import (
	check_dependencies,
	parse_task_args,
	process_metrics,
	remove_stale_marker,
	write_marker,
)
from common.queries import Q1_TRACK_CODE, csv_queries, json_queries, rdf_queries
from common.results import compare, serialize, sha256_of, sort_rows

logger = logging.getLogger(__name__)

TASK_ID = "08"
DEPENDS_ON = ["03", "04", "05", "06", "07"]

EXACT_ROW_COUNTS = {1: 1000, 3: 480, 5: 4, 6: 480, 7: 2}  # fixed by corpus volumes

# --- SQL implementations -------------------------------------------------
# Q2/Q3/Q4/Q7 use the precomputed track_summary (the relational strength);
# Q5 must scan raw friction_cycles in every format (spec 08).

Q1_SQL = f"""-- Q1: COF vs cycle curve for track {Q1_TRACK_CODE} (point read)
SELECT fc.cycle, fc.cof
FROM friction_cycles fc
JOIN tracks t ON t.track_id = fc.track_id
WHERE t.track_code = '{Q1_TRACK_CODE}'
ORDER BY fc.cycle;
"""

Q2_SQL = """-- Q2: steady-state COF per coupon with mean Au = 10 +/- 0.5 wt%
SELECT c.coupon_code, t.environment, AVG(s.cof_ss_mean) AS cof_ss
FROM coupons c
JOIN tracks t ON t.coupon_id = c.coupon_id
JOIN track_summary s ON s.track_id = t.track_id
WHERE c.au_wtpct_mean BETWEEN 9.5 AND 10.5
GROUP BY c.coupon_code, t.environment
ORDER BY c.coupon_code;
"""

Q3_SQL = """-- Q3: hardness vs steady-state COF per friction coupon, across all batches
SELECT c.coupon_code,
       b.batch_code,
       (SELECT AVG(n.hardness_gpa) FROM nanoindentation n WHERE n.coupon_id = c.coupon_id) AS hardness_gpa,
       AVG(s.cof_ss_mean) AS cof_ss
FROM coupons c
JOIN batches b ON b.batch_id = c.batch_id
JOIN tracks t ON t.coupon_id = c.coupon_id
JOIN track_summary s ON s.track_id = t.track_id
GROUP BY c.coupon_id, c.coupon_code, b.batch_code
ORDER BY c.coupon_code;
"""

Q4_SQL = """-- Q4: anomaly filter - Dry N2 tracks at 100 mN with cof_ss > 0.20
SELECT t.track_code, s.cof_ss_mean
FROM tracks t
JOIN track_summary s ON s.track_id = t.track_id
WHERE t.environment = 'dry_n2' AND t.load_mn = 100 AND s.cof_ss_mean > 0.20
ORDER BY t.track_code;
"""

# Q5: forced full scan over raw cycles; the run-in algorithm is inlined
# (docs/rules/db-sql-schema.md section 5). SQLite has no stddev built-in,
# so the 2-sigma band compares squares: (cof-m)^2 < 4*var  <=>  |cof-m| < 2s.
Q5_SQLITE = """-- Q5: mean run-in cycles per batch over ALL raw cycle rows (track_summary forbidden)
WITH tail AS (
	SELECT track_id, AVG(cof) AS m, AVG(cof * cof) - AVG(cof) * AVG(cof) AS var
	FROM friction_cycles
	WHERE cycle >= 501
	GROUP BY track_id
),
run_in AS (
	SELECT fc.track_id,
	       COALESCE(MIN(CASE WHEN (fc.cof - tl.m) * (fc.cof - tl.m) < 4 * tl.var THEN fc.cycle END), 500) AS run_in
	FROM friction_cycles fc
	JOIN tail tl ON tl.track_id = fc.track_id
	GROUP BY fc.track_id
)
SELECT b.batch_code, AVG(r.run_in) AS avg_run_in_cycles
FROM run_in r
JOIN tracks t ON t.track_id = r.track_id
JOIN coupons c ON c.coupon_id = t.coupon_id
JOIN batches b ON b.batch_id = c.batch_id
GROUP BY b.batch_code
ORDER BY b.batch_code;
"""

Q5_PG = """-- Q5: mean run-in cycles per batch over ALL raw cycle rows (track_summary forbidden)
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 - tl.m) < 2 * tl.s), 500) AS run_in
	FROM friction_cycles fc
	JOIN tail tl ON tl.track_id = fc.track_id
	GROUP BY fc.track_id
)
SELECT b.batch_code, AVG(r.run_in)::double precision AS avg_run_in_cycles
FROM run_in r
JOIN tracks t ON t.track_id = r.track_id
JOIN coupons c ON c.coupon_id = t.coupon_id
JOIN batches b ON b.batch_id = c.batch_id
GROUP BY b.batch_code
ORDER BY b.batch_code;
"""

Q6_SQL = """-- Q6: mean wear volume vs load per friction coupon
SELECT c.coupon_code, t.load_mn, AVG(w.wear_volume_um3) AS wear_volume_um3
FROM wear w
JOIN tracks t ON t.track_id = w.track_id
JOIN coupons c ON c.coupon_id = t.coupon_id
GROUP BY c.coupon_id, c.coupon_code, t.load_mn
ORDER BY c.coupon_code;
"""

Q7_SQL = """-- Q7: Stribeck-style mean COF by (environment, speed*load bucket) past run-in
SELECT t.environment,
       CAST(ROUND(t.speed_mm_s * t.load_mn) AS INTEGER) AS speed_load_bucket,
       AVG(fc.cof) AS mean_cof
FROM friction_cycles fc
JOIN tracks t ON t.track_id = fc.track_id
JOIN track_summary s ON s.track_id = fc.track_id
WHERE fc.cycle > s.run_in_cycles
GROUP BY t.environment, speed_load_bucket
ORDER BY t.environment, speed_load_bucket;
"""

SQLITE_SQL = {1: Q1_SQL, 2: Q2_SQL, 3: Q3_SQL, 4: Q4_SQL, 5: Q5_SQLITE, 6: Q6_SQL, 7: Q7_SQL}
PG_SQL = {1: Q1_SQL, 2: Q2_SQL, 3: Q3_SQL, 4: Q4_SQL, 5: Q5_PG, 6: Q6_SQL, 7: Q7_SQL}

RDF_SPARQL_FILES = {
	1: ("q1",), 2: ("q2",), 3: ("q3_hardness", "q3_cycles"),
	4: ("q4",), 5: ("q5",), 6: ("q6",), 7: ("q7",),
}

WRAPPER = '''"""Benchmark query q{n} for the {fmt} format (generated by task 08).

Prints the canonical result rows to stdout at full float precision; the
benchmark runner (task 09) captures and validates them against
out/bench/expected/q{n}.rows.csv.
"""
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[4]
sys.path.insert(0, str(REPO_ROOT))

from common.queries.{module} import q{n}
from common.results import print_rows

if __name__ == "__main__":
	print_rows(q{n}(REPO_ROOT / "out" / {subpath}))
'''

PY_FORMATS = [
	("csv", "csv_queries", '"csv"'),
	("json", "json_queries", '"json"'),
	("rdf", "rdf_queries", '"rdf" / "oxigraph_store"'),
]

SQLITE_WRAPPER = '''"""Benchmark query q{n} for the sqlite format (generated by task 08).

Executes the adjacent q{n}.sql against out/sqlite/tribo.db (read-only) and
prints the canonical result rows to stdout at full float precision; the
benchmark runner (task 09) captures and validates them against
out/bench/expected/q{n}.rows.csv.
"""
import sqlite3
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[4]
sys.path.insert(0, str(REPO_ROOT))

from common.results import print_rows

if __name__ == "__main__":
	sql = Path(__file__).with_suffix(".sql").read_text(encoding="ascii")
	db = REPO_ROOT / "out" / "sqlite" / "tribo.db"
	conn = sqlite3.connect(f"file:{{db.as_posix()}}?mode=ro", uri=True)
	print_rows([tuple(r) for r in conn.execute(sql).fetchall()])
'''

PG_WRAPPER = '''"""Benchmark query q{n} for the pg format (generated by task 08).

Executes the adjacent q{n}.sql on the PostgreSQL instance addressed by
TRIBO_PG_DSN and prints the canonical result rows to stdout at full float
precision; the benchmark runner (task 09) captures and validates them
against out/bench/expected/q{n}.rows.csv.
"""
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[4]
sys.path.insert(0, str(REPO_ROOT))

import psycopg

from common.pg import get_dsn
from common.results import print_rows

if __name__ == "__main__":
	sql = Path(__file__).with_suffix(".sql").read_text(encoding="ascii")
	with psycopg.connect(get_dsn()) as conn:
		print_rows([tuple(r) for r in conn.execute(sql).fetchall()])
'''

SQL_WRAPPERS = {"sqlite": SQLITE_WRAPPER, "pg": PG_WRAPPER}


def write_query_files(bench_dir: Path) -> int:
	queries_dir = bench_dir / "queries"
	written = 0
	for fmt, module, subpath in PY_FORMATS:
		fmt_dir = queries_dir / fmt
		fmt_dir.mkdir(parents=True, exist_ok=True)
		for n in range(1, 8):
			text = WRAPPER.format(n=n, fmt=fmt, module=module, subpath=subpath)
			(fmt_dir / f"q{n}.py").write_text(text, encoding="ascii", newline="\n")
			written += 1
	for fmt, sql in (("sqlite", SQLITE_SQL), ("pg", PG_SQL)):
		fmt_dir = queries_dir / fmt
		fmt_dir.mkdir(parents=True, exist_ok=True)
		for n, text in sql.items():
			(fmt_dir / f"q{n}.sql").write_text(text, encoding="ascii", newline="\n")
			(fmt_dir / f"q{n}.py").write_text(SQL_WRAPPERS[fmt].format(n=n), encoding="ascii", newline="\n")
			written += 2
	rdf_dir = queries_dir / "rdf"
	for n, keys in RDF_SPARQL_FILES.items():
		text = "\n\n".join(f"# retrieval query: {key}\n{rdf_queries.SPARQL[key]}" for key in keys)
		(rdf_dir / f"q{n}.sparql").write_text(text + "\n", encoding="ascii", newline="\n")
		written += 1
	logger.info("wrote %d query files under %s", written, queries_dir)
	return written


def run_sqlite(db_path: Path) -> dict[int, list[tuple]]:
	conn = sqlite3.connect(db_path)
	try:
		return {n: [tuple(r) for r in conn.execute(sql).fetchall()] for n, sql in SQLITE_SQL.items()}
	finally:
		conn.close()


def run_pg(dsn: str) -> dict[int, list[tuple]]:
	results: dict[int, list[tuple]] = {}
	with psycopg.connect(dsn) as conn:
		for n, sql in PG_SQL.items():
			results[n] = [tuple(r) for r in conn.execute(sql).fetchall()]
			logger.info("pg q%d: %d rows", n, len(results[n]))
	return results


def main() -> int:
	args = parse_task_args("Task 08: define benchmark queries and canonical expected results")
	out_root: Path = args.out_root
	bench_dir = out_root / "bench"
	try:
		check_dependencies(out_root, DEPENDS_ON)
		remove_stale_marker(out_root, TASK_ID)
		files_written = write_query_files(bench_dir)

		logger.info("computing canonical results on PostgreSQL")
		pg_rows = run_pg(get_dsn())
		logger.info("cross-validating against SQLite")
		sqlite_rows = run_sqlite(out_root / "sqlite" / "tribo.db")
		for n in range(1, 8):
			diff = compare(pg_rows[n], sqlite_rows[n])
			if diff:
				raise ValueError(f"q{n}: SQLite disagrees with PostgreSQL: {diff}")
			exact = EXACT_ROW_COUNTS.get(n)
			if exact is not None and len(pg_rows[n]) != exact:
				raise ValueError(f"q{n}: {len(pg_rows[n])} rows, expected {exact}")
			if not pg_rows[n]:
				raise ValueError(f"q{n}: empty result set")

		logger.info("spot-checking csv/json/rdf on Q1")
		for label, rows in (
			("csv", csv_queries.q1(out_root / "csv")),
			("json", json_queries.q1(out_root / "json")),
			("rdf", rdf_queries.q1(out_root / "rdf" / "oxigraph_store")),
		):
			diff = compare(pg_rows[1], rows)
			if diff:
				raise ValueError(f"q1 via {label} disagrees with canonical: {diff}")

		expected_dir = bench_dir / "expected"
		expected_dir.mkdir(parents=True, exist_ok=True)
		checksums = {}
		for n in range(1, 8):
			rows = sort_rows(pg_rows[n])
			(expected_dir / f"q{n}.rows.csv").write_text(serialize(rows), encoding="ascii", newline="\n")
			checksums[n] = sha256_of(rows)
			(expected_dir / f"q{n}.sha256").write_text(checksums[n] + "\n", encoding="ascii", newline="\n")

		entries = {"query_files": files_written}
		for n in range(1, 8):
			entries[f"q{n}_rows"] = len(pg_rows[n])
			entries[f"q{n}_sha256"] = checksums[n][:12]
		entries.update(process_metrics())
		marker_path = write_marker(out_root, TASK_ID, entries)
	except Exception:
		logger.critical("task %s failed", TASK_ID, exc_info=True)
		return 1

	counts = ", ".join(f"q{n}={len(pg_rows[n])}" for n in range(1, 8))
	print(f"task 08 ok: {files_written} query files, canonical results validated (pg == sqlite, 1e-9)")
	print(f"rows: {counts}")
	print(f"expected: {bench_dir / 'expected'} (q<N>.rows.csv + q<N>.sha256)")
	print(f"marker: {marker_path}")
	return 0


if __name__ == "__main__":
	sys.exit(main())