"""HNSW build time vs row count with a known number of parallel workers, merged into build.json.
e2_scale.py did not record how many workers PostgreSQL launched at 250K and 500K rows, so its build
times mix table size with worker count. Here, at each size, we build twice: the default way (counting
the workers PostgreSQL chose) and with ALTER TABLE ... SET (parallel_workers = 3)."""
import json, threading, time
from common import *

path = f"{OUT}/build.json"
out = json.load(open(path))


def build(c, table, label, reloption):
    c.execute("SET maintenance_work_mem = '4GB'")
    c.execute("SET max_parallel_maintenance_workers = 3")
    c.execute(f"ALTER TABLE {table} SET (parallel_workers = {reloption})" if reloption is not None else f"ALTER TABLE {table} RESET (parallel_workers)")
    c.execute(f"DROP INDEX IF EXISTS {table}_h")
    seen, stop = [0], [False]

    def poll():
        with connect() as m:
            while not stop[0]:
                n = m.execute("SELECT count(*) FROM pg_stat_activity WHERE backend_type = 'parallel worker' AND datname = 'bench'").fetchone()[0]
                seen[0] = max(seen[0], n)
                time.sleep(0.5)
    th = threading.Thread(target=poll); th.start()
    t = time.perf_counter()
    c.execute(f"CREATE INDEX {table}_h ON {table} USING hnsw (embedding vector_cosine_ops)")
    el = round(time.perf_counter() - t, 1)
    stop[0] = True; th.join()
    r = {"label": label, "parallel_workers_reloption": reloption, "workers_launched": seen[0], "build_s": el}
    print(r, flush=True)
    return r


sizes = []
with connect() as c:
    for n in (250_000, 500_000):
        load_table(c, "items_s9", corpus(n))
        heap = c.execute("SELECT pg_relation_size('items_s9')").fetchone()[0]
        runs = [build(c, "items_s9", "max_parallel_maintenance_workers = 3 (no reloption)", None),
                build(c, "items_s9", "+ ALTER TABLE SET (parallel_workers = 3)", 3)]
        sizes.append({"n": n, "heap_bytes": heap, "runs": runs})
        print(n, heap, flush=True)
        c.execute("DROP TABLE items_s9")
out["scale_check"] = {"note": "max_parallel_maintenance_workers = 3, maintenance_work_mem = 4GB, m 16 / ef_construction 64", "sizes": sizes}
save("build", out)