S119 · Vector Databases — pgvector, HNSW, IVF

"""pgvector_bench.py — build HNSW and IVF-Flat over synthetic vectors.
 
Requires: pip install psycopg[binary] numpy
Postgres 15+ with pgvector 0.5+ installed:  CREATE EXTENSION vector;
Set DATABASE_URL, e.g. postgresql://user:pass@localhost:5432/rag
"""
from __future__ import annotations
import os
import time
import numpy as np
import psycopg
 
DB_URL = os.environ["DATABASE_URL"]
N, D = 100_000, 384          # 100K vectors of dim 384 (all-MiniLM-L6-v2 size)
K = 10                       # top-K
 
 
def setup(conn):
    with conn.cursor() as cur:
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
        cur.execute("DROP TABLE IF EXISTS items;")
        cur.execute(f"""
            CREATE TABLE items (
                id      SERIAL PRIMARY KEY,
                cat     TEXT,
                embed   vector({D})
            );
        """)
    conn.commit()
 
 
def ingest(conn):
    print(f"Generating {N} random vectors ...")
    rng = np.random.default_rng(0)
    vecs = rng.standard_normal((N, D)).astype(np.float32)
    vecs /= np.linalg.norm(vecs, axis=1, keepdims=True)   # unit-normalise
    cats = rng.choice(["a", "b", "c"], size=N)
 
    print("Bulk insert ...")
    with conn.cursor() as cur, conn.pipeline():
        for i, (v, c) in enumerate(zip(vecs, cats)):
            cur.execute(
                "INSERT INTO items (cat, embed) VALUES (%s, %s);",
                (c, v.tolist()),
            )
    conn.commit()
    return vecs
 
 
def build_hnsw(conn, m=16, ef_construction=200):
    print(f"Building HNSW (m={m}, ef_construction={ef_construction}) ...")
    with conn.cursor() as cur:
        cur.execute("DROP INDEX IF EXISTS items_hnsw;")
        t0 = time.time()
        cur.execute(f"""
            CREATE INDEX items_hnsw ON items
            USING hnsw (embed vector_cosine_ops)
            WITH (m = {m}, ef_construction = {ef_construction});
        """)
    conn.commit()
    print(f"  built in {time.time()-t0:.1f}s")
 
 
def build_ivfflat(conn, lists=316):     # ~ sqrt(N)
    print(f"Building IVF-Flat (lists={lists}) ...")
    with conn.cursor() as cur:
        cur.execute("DROP INDEX IF EXISTS items_ivf;")
        t0 = time.time()
        cur.execute(f"""
            CREATE INDEX items_ivf ON items
            USING ivfflat (embed vector_cosine_ops)
            WITH (lists = {lists});
        """)
    conn.commit()
    print(f"  built in {time.time()-t0:.1f}s")
 
 
def query(conn, q_vec, k=K, ef_search=None, nprobe=None):
    with conn.cursor() as cur:
        if ef_search is not None:
            cur.execute(f"SET hnsw.ef_search = {ef_search};")
        if nprobe is not None:
            cur.execute(f"SET ivfflat.probes = {nprobe};")
        t0 = time.perf_counter()
        cur.execute(
            "SELECT id FROM items ORDER BY embed <=> %s LIMIT %s;",
            (q_vec.tolist(), k),
        )
        rows = cur.fetchall()
        return [r[0] for r in rows], (time.perf_counter() - t0) * 1000
 
 
def brute_force_gold(vecs, q, k=K):
    """Ground truth: compare q to every vector."""
    sims = vecs @ q                       # both unit-normalised → cosine == dot
    return np.argpartition(-sims, k)[:k]
 
 
def bench(conn, vecs, index_kind: str):
    rng = np.random.default_rng(42)
    queries = rng.standard_normal((50, D)).astype(np.float32)
    queries /= np.linalg.norm(queries, axis=1, keepdims=True)
 
    for knob in (10, 50, 200):
        recalls, lats = [], []
        for q in queries:
            gold = set(int(i) + 1 for i in brute_force_gold(vecs, q))  # +1: SERIAL
            got, ms = query(conn, q,
                            ef_search=(knob if index_kind == "hnsw" else None),
                            nprobe=(knob if index_kind == "ivf" else None))
            recalls.append(len(gold & set(got)) / K)
            lats.append(ms)
        print(f"  {index_kind}  knob={knob:>4}  "
              f"recall@{K}={np.mean(recalls):.3f}  "
              f"p50={np.percentile(lats,50):.1f}ms  "
              f"p95={np.percentile(lats,95):.1f}ms")
 
 
with psycopg.connect(DB_URL) as conn:
    setup(conn)
    vecs = ingest(conn)
 
    print("\n--- HNSW ---")
    build_hnsw(conn)
    bench(conn, vecs, "hnsw")
 
    print("\n--- IVF-Flat ---")
    build_ivfflat(conn)
    bench(conn, vecs, "ivf")