"""Parallel HNSW build check (second run, idle machine), merged into build.json.
Hypothesis under test: 1,536-dim vectors are TOASTed, so the main heap is tiny and PostgreSQL plans
no parallel workers for CREATE INDEX unless the table's parallel_workers reloption is set.
We count launched parallel workers by polling pg_stat_activity from a second connection."""
import json, threading, time
from common import *

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


def build(c, label, workers, reloption):
    c.execute("SET maintenance_work_mem = '4GB'")
    c.execute(f"SET max_parallel_maintenance_workers = {workers}")
    c.execute(f"ALTER TABLE items_b3 SET (parallel_workers = {reloption})" if reloption is not None else "ALTER TABLE items_b3 RESET (parallel_workers)")
    c.execute("DROP INDEX IF EXISTS items_b3_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("CREATE INDEX items_b3_h ON items_b3 USING hnsw (embedding vector_cosine_ops)")
    el = round(time.perf_counter() - t, 1)
    stop[0] = True; th.join()
    r = {"label": label, "max_parallel_maintenance_workers": workers, "parallel_workers_reloption": reloption, "workers_launched": seen[0], "build_s": el}
    print(r, flush=True)
    return r


with connect() as c:
    load_table(c, "items_b3", corpus(100_000))
    heap = c.execute("SELECT pg_relation_size('items_b3')").fetchone()[0]
    total = c.execute("SELECT pg_table_size('items_b3')").fetchone()[0]
    runs = [
        build(c, "max_parallel_maintenance_workers = 3 (no reloption)", 3, None),
        build(c, "max_parallel_maintenance_workers = 0", 0, None),
        build(c, "workers 3 + ALTER TABLE SET (parallel_workers = 3)", 3, 3),
        build(c, "workers 1 + ALTER TABLE SET (parallel_workers = 1)", 1, 1),
    ]
    c.execute("DROP TABLE items_b3")
out["parallel_check"] = {"heap_bytes": heap, "table_bytes_incl_toast": total, "runs": runs}
save("build", out)