Vector Search with pgvector

# script_id: day_026_pgvector_search/hand_example
# Cosine distance by hand for three toy rows. Install: uv add numpy
import numpy as np

rows = {"a": [1.0, 0.0], "b": [0.6, 0.8], "c": [0.0, 1.0]}
q = np.array([1.0, 0.0])
for name, vec in rows.items():
    v = np.array(vec)
    cos = float(v @ q / (np.linalg.norm(v) * np.linalg.norm(q)))
    print(f"{name}  similarity={cos:.2f}  distance={1 - cos:.2f}")

a  similarity=1.00  distance=0.00
b  similarity=0.60  distance=0.40
c  similarity=0.00  distance=1.00

# compose.yaml: Postgres 18 + pgvector 0.8.7 for Atlas. Start: docker compose up -d
services:
  db:
    image: pgvector/pgvector:0.8.7-pg18-trixie
    environment:
      POSTGRES_USER: atlas
      POSTGRES_PASSWORD: atlas   # local dev only; never reuse this password anywhere real
      POSTGRES_DB: atlas
    ports:
      - "5432:5432"
    volumes:
      - pgdata:/var/lib/postgresql   # PG 18 images keep data here, not in .../data
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U atlas -d atlas"]
      interval: 2s
      timeout: 5s
      retries: 15
volumes:
  pgdata:

# script_id: day_026_pgvector_search/hello_pgvector
# Minimal pgvector round trip: extension, table, three rows, one nearest-neighbour query.
# Install: uv add "psycopg[binary]" pgvector numpy
# Env: DATABASE_URL (default matches compose.yaml). Services: Postgres from compose.yaml
import os

import numpy as np
import psycopg
from pgvector.psycopg import register_vector

url = os.environ.get("DATABASE_URL", "postgresql://atlas:atlas@localhost:5432/atlas")
with psycopg.connect(url, autocommit=True) as conn:
    conn.execute("CREATE EXTENSION IF NOT EXISTS vector")  # once per database
    register_vector(conn)  # once per connection, and only after the extension exists
    conn.execute("DROP TABLE IF EXISTS toy")
    conn.execute("CREATE TABLE toy (id text PRIMARY KEY, embedding vector(2))")
    for name, vec in [("a", [1, 0]), ("b", [0.6, 0.8]), ("c", [0, 1])]:
        conn.execute("INSERT INTO toy VALUES (%s, %s)", (name, np.array(vec, dtype=np.float32)))
    q = np.array([1, 0], dtype=np.float32)
    rows = conn.execute("SELECT id, embedding <=> %s AS dist FROM toy ORDER BY dist LIMIT 2", (q,)).fetchall()
    for id_, dist in rows:
        print(f"{id_}  distance={dist:.2f}  similarity={1 - dist:.2f}")

a  distance=0.00  similarity=1.00
b  distance=0.40  similarity=0.60

# script_id: day_026_pgvector_search/break_it
# Three common pgvector failures, caught and printed. Run hello_pgvector first (it creates `toy`).
# Install: uv add "psycopg[binary]" pgvector numpy · Services: Postgres from compose.yaml
import os

import numpy as np
import psycopg
from pgvector.psycopg import register_vector

url = os.environ.get("DATABASE_URL", "postgresql://atlas:atlas@localhost:5432/atlas")
conn = psycopg.connect(url, autocommit=True)  # autocommit: one failed statement doesn't poison the rest
register_vector(conn)


def attempt(label: str, sql: str, params: tuple) -> None:
    try:
        conn.execute(sql, params).fetchall()
        print(f"{label}: no error")
    except psycopg.Error as exc:
        print(f"{label}: {type(exc).__name__}: {str(exc).splitlines()[0]}")


# 1. A plain Python list is sent as a float array, not a vector.
attempt("plain list", "SELECT id FROM toy ORDER BY embedding <=> %s LIMIT 1", ([1.0, 0.0],))
# 2. Wrong size on insert: a 3-d vector into the vector(2) column (same error for 256-d into 1536-d).
attempt("insert 3-d", "INSERT INTO toy VALUES ('d', %s) RETURNING id", (np.zeros(3, dtype=np.float32),))
# 3. Query embedded with a different model or `dimensions=` than the stored rows (3-d vs 2-d here).
attempt("query 3-d", "SELECT id FROM toy ORDER BY embedding <=> %s LIMIT 1", (np.ones(3, dtype=np.float32),))
conn.close()

plain list: UndefinedFunction: operator does not exist: vector <=> double precision[]
insert 3-d: DataException: expected 2 dimensions, not 3
query 3-d: DataException: different vector dimensions 2 and 3

# script_id: day_026_pgvector_search/support
# support.py: compact copies of earlier helpers so today's files stay self-contained.
# built on Day 23 (embed, hit_rate), Day 24 (load_corpus), Day 25 (DocChunk fields, chunk), Day 12 (cost)
# Install: uv add "openai>=3.46,<4" pydantic numpy
# Env: ORBIT_DATA (path to datasets/orbit); OPENAI_API_KEY only when embed() is called
from __future__ import annotations

import hashlib
import json
import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Sequence

import numpy as np
from pydantic import BaseModel, Field

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
EMBED_MODEL = os.environ.get("ATLAS_EMBED_MODEL", "text-embedding-3-small")
EMBED_DIM = 1536  # text-embedding-3-small's default size; the vector(n) column must match it
DATABASE_URL = os.environ.get("DATABASE_URL", "postgresql://atlas:atlas@localhost:5432/atlas")

# script_id: day_026_pgvector_search/support
class DocChunk(BaseModel):
    id: str
    doc_id: str
    section: str = ""
    applies_to: str = "all"
    text: str
    position: int = 0
    n_tokens: int = 0
    header: str = ""
    metadata: dict[str, Any] = Field(default_factory=dict)


@dataclass
class Document:
    id: str
    title: str
    text: str
    metadata: dict[str, str]

# script_id: day_026_pgvector_search/support
def load_corpus(root: Path = DATA_DIR / "docs") -> list[Document]:
    docs = []
    for path in sorted(root.glob("*.md")):
        _, front, body = path.read_text(encoding="utf-8").split("---", 2)
        meta = {}
        for line in front.strip().splitlines():
            key, _, value = line.partition(":")
            meta[key.strip()] = value.split(" #")[0].strip().strip('"')  # drop inline comments, quotes
        docs.append(Document(meta.get("id", path.stem), meta.get("title", path.stem), body.strip(), meta))
    return docs


def chunk(doc: Document, *, max_chars: int = 1600) -> list[DocChunk]:
    """Split at H2 headings (outside code fences), then pack paragraphs up to max_chars."""
    sections: list[tuple[str, list[str]]] = [("", [])]
    in_fence = False
    for line in doc.text.splitlines():
        in_fence ^= line.startswith("```")
        if line.startswith("## ") and not in_fence:
            sections.append((line[3:].strip(), []))
        else:
            sections[-1][1].append(line)
    out: list[DocChunk] = []
    for heading, lines in sections:
        pieces, buf = [], ""
        for para in "\n".join(lines).split("\n\n"):
            if buf and len(buf) + len(para) > max_chars:
                pieces.append(buf)
                buf = ""
            buf = f"{buf}\n\n{para}" if buf else para
        for text in (p.strip() for p in pieces + [buf]):
            if text:
                out.append(DocChunk(
                    id=f"{doc.id}#{len(out)}", doc_id=doc.id, position=len(out), text=text,
                    section=doc.metadata.get("section", ""), applies_to=doc.metadata.get("applies_to", "all"),
                    header=f"{doc.title} > {heading}" if heading else doc.title))
    return out

# script_id: day_026_pgvector_search/support
def embed_text(c: DocChunk) -> str:
    return f"{c.header}\n\n{c.text}" if c.header else c.text

# script_id: day_026_pgvector_search/support
EMBED_PRICE_PER_1M = {"text-embedding-3-small": 0.02, "text-embedding-3-large": 0.13}  # USD, as of 2026-11, verify


@dataclass
class CostTracker:
    tokens: int = 0
    total_usd: float = 0.0

    def record(self, model: str, tokens: int) -> None:
        self.tokens += tokens
        self.total_usd += tokens * EMBED_PRICE_PER_1M.get(model, 0.0) / 1_000_000


TRACKER = CostTracker()
_client = None


def embed(texts: Sequence[str], *, model: str = EMBED_MODEL, batch_size: int = 256) -> np.ndarray:
    """(n, d) float32 rows in input order. OpenAI embeddings come back unit-length."""
    global _client
    if _client is None:
        from openai import OpenAI  # lazy, so offline demos never need a key
        _client = OpenAI()
    rows: list[list[float]] = []
    for i in range(0, len(texts), batch_size):
        resp = _client.embeddings.create(model=model, input=list(texts[i:i + batch_size]))
        TRACKER.record(model, resp.usage.prompt_tokens)
        rows += [d.embedding for d in sorted(resp.data, key=lambda d: d.index)]
    return np.asarray(rows, dtype=np.float32)


def fake_embed(texts: Sequence[str], dim: int = EMBED_DIM) -> np.ndarray:
    out = np.zeros((len(texts), dim), dtype=np.float32)
    for row, text in enumerate(texts):
        for word in re.findall(r"[a-z0-9_]+", text.lower()):
            out[row, int(hashlib.md5(word.encode()).hexdigest(), 16) % dim] += 1.0
    norms = np.linalg.norm(out, axis=1, keepdims=True)
    return out / np.where(norms == 0, 1, norms)  # unit length, like real embeddings

# script_id: day_026_pgvector_search/support
@dataclass
class RetrievalCase:
    id: str
    query: str
    relevant_doc_ids: list[str]
    kind: str


def load_retrieval_cases(split: str = "dev") -> list[RetrievalCase]:
    """Day 23's deterministic split: alternate cases by id within each kind."""
    lines = (DATA_DIR / "eval" / "retrieval_queries.jsonl").read_text(encoding="utf-8").splitlines()
    by_kind: dict[str, list[RetrievalCase]] = {}
    for case in sorted((RetrievalCase(**json.loads(l)) for l in lines if l.strip()), key=lambda c: c.id):
        by_kind.setdefault(case.kind, []).append(case)
    keep = 0 if split == "dev" else 1
    return [c for g in by_kind.values() for i, c in enumerate(g) if split == "all" or i % 2 == keep]


def ranked_doc_ids(chunks: Sequence[DocChunk]) -> list[str]:
    return list(dict.fromkeys(c.doc_id for c in chunks))  # dedupe, keep rank order


def hit_rate(cases: Sequence[RetrievalCase], search_fn: Callable[[str, int], list[str]],
             k: int = 5) -> dict[str, float]:
    results: dict[str, list[bool]] = {}
    for case in cases:
        hit = bool(set(search_fn(case.query, k)[:k]) & set(case.relevant_doc_ids))
        for key in ("overall", case.kind):
            results.setdefault(key, []).append(hit)
    return {key: round(sum(v) / len(v), 2) for key, v in results.items()}


if __name__ == "__main__":
    docs = load_corpus()
    chunks = [c for d in docs for c in chunk(d)]
    print(f"{len(docs)} pages · {len(chunks)} chunks · {len(load_retrieval_cases())} dev cases")

45 pages · 367 chunks · 20 dev cases

# script_id: day_026_pgvector_search/index
# index.py: Atlas's durable vector index (atlas/index.py in the capstone).
# Install: uv add "psycopg[binary]" pgvector numpy "openai>=3.46,<4" pydantic
# Env: DATABASE_URL, OPENAI_API_KEY (for embed), ORBIT_DATA. Services: Postgres from compose.yaml
from __future__ import annotations

import hashlib
from dataclasses import dataclass
from typing import Callable, Sequence

import numpy as np
import psycopg
from pgvector.psycopg import register_vector
from psycopg.types.json import Jsonb

from support import DATABASE_URL, EMBED_DIM, TRACKER, DocChunk, embed, embed_text

EmbedFn = Callable[[Sequence[str]], np.ndarray]


@dataclass(frozen=True)
class Hit:
    chunk: DocChunk
    score: float  # 1 - cosine distance: higher is better


@dataclass
class UpsertStats:
    inserted: int = 0
    unchanged: int = 0
    deleted: int = 0
    cost_usd: float = 0.0

    def __str__(self) -> str:
        return f"{self.inserted} inserted · {self.unchanged} unchanged · {self.deleted} deleted · ${self.cost_usd:.5f}"


def chunk_id(doc_id: str, text: str) -> str:
    """Same doc + same embedded text -> same id, like an ETag."""
    return hashlib.sha256(f"{doc_id}\n{text}".encode()).hexdigest()[:16]

# script_id: day_026_pgvector_search/index
def connect(url: str = DATABASE_URL) -> psycopg.Connection:
    conn = psycopg.connect(url, autocommit=True)
    conn.execute("CREATE EXTENSION IF NOT EXISTS vector")  # must exist before register_vector
    register_vector(conn)
    conn.execute("SET hnsw.iterative_scan = relaxed_order")  # pgvector 0.8+: see "ANN Indexes and Tuning"
    return conn

# script_id: day_026_pgvector_search/index
SCHEMA = [
    """CREATE TABLE IF NOT EXISTS {t} (
         id text PRIMARY KEY, doc_id text NOT NULL, position int NOT NULL,
         section text NOT NULL, applies_to text NOT NULL, header text NOT NULL, body text NOT NULL,
         metadata jsonb NOT NULL DEFAULT '{{}}', embedding vector({dim}) NOT NULL)""",
    "CREATE INDEX IF NOT EXISTS {t}_doc_id ON {t} (doc_id)",
]
HNSW = "CREATE INDEX IF NOT EXISTS {t}_hnsw ON {t} USING hnsw (embedding vector_cosine_ops)"
UPSERT = """INSERT INTO {t} (id, doc_id, position, section, applies_to, header, body, metadata, embedding)
    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
    ON CONFLICT (id) DO UPDATE SET position = EXCLUDED.position, metadata = EXCLUDED.metadata,
                                   embedding = EXCLUDED.embedding"""
REFRESH = "UPDATE {t} SET position = %s, section = %s, applies_to = %s WHERE id = %s"


class Index:
    """Durable chunk store + vector search on Postgres/pgvector."""

    def __init__(self, conn: psycopg.Connection, *, table: str = "chunks", dim: int = EMBED_DIM,
                 embed_fn: EmbedFn = embed) -> None:
        if not table.isidentifier():
            raise ValueError(f"bad table name: {table!r}")
        self.conn, self.table, self.dim, self.embed_fn = conn, table, dim, embed_fn

    def create(self, hnsw: bool = True) -> None:
        for stmt in SCHEMA + ([HNSW] if hnsw else []):
            self.conn.execute(stmt.format(t=self.table, dim=self.dim))

    def count(self) -> int:
        return self.conn.execute(f"SELECT count(*) FROM {self.table}").fetchone()[0]

# script_id: day_026_pgvector_search/index
    def upsert(self, chunks: Sequence[DocChunk]) -> UpsertStats:
        keyed = {chunk_id(c.doc_id, embed_text(c)): c for c in chunks}
        doc_ids = sorted({c.doc_id for c in chunks})
        existing = {row[0] for row in self.conn.execute(
            f"SELECT id FROM {self.table} WHERE doc_id = ANY(%s)", (doc_ids,))}
        new = [cid for cid in keyed if cid not in existing]
        kept = [cid for cid in keyed if cid in existing]
        stale = sorted(existing - keyed.keys())  # rows of these docs that the new version no longer has
        before = TRACKER.total_usd
        vectors = self.embed_fn([embed_text(keyed[cid]) for cid in new]) if new else []
        with self.conn.transaction(), self.conn.cursor() as cur:
            if new:
                cur.executemany(UPSERT.format(t=self.table), [
                    (cid, c.doc_id, c.position, c.section, c.applies_to, c.header, c.text, Jsonb(c.metadata), vec)
                    for cid, vec in zip(new, vectors) for c in [keyed[cid]]])
            if kept:
                cur.executemany(REFRESH.format(t=self.table), [
                    (keyed[cid].position, keyed[cid].section, keyed[cid].applies_to, cid) for cid in kept])
            if stale:
                cur.execute(f"DELETE FROM {self.table} WHERE id = ANY(%s)", (stale,))
        return UpsertStats(len(new), len(kept), len(stale), TRACKER.total_usd - before)

    def delete(self, doc_id: str) -> int:
        return self.conn.execute(f"DELETE FROM {self.table} WHERE doc_id = %s", (doc_id,)).rowcount

# script_id: day_026_pgvector_search/index
    def search(self, query: str, k: int = 5, *,
               filters: dict[str, str | list[str]] | None = None) -> list[Hit]:
        where, params = _where(filters)
        rows = self.conn.execute(
            f"""SELECT id, doc_id, position, section, applies_to, header, body, metadata,
                       embedding <=> %(q)s AS distance
                FROM {self.table} {where} ORDER BY embedding <=> %(q)s LIMIT %(k)s""",
            {"q": self.embed_fn([query])[0], "k": k, **params}).fetchall()
        hits = [Hit(DocChunk(id=r[0], doc_id=r[1], position=r[2], section=r[3], applies_to=r[4],
                             header=r[5], text=r[6], metadata=r[7]), 1.0 - r[8]) for r in rows]
        return sorted(hits, key=lambda h: h.score, reverse=True)  # iterative scans may return near-sorted rows


FILTERABLE = ("doc_id", "section", "applies_to")


def _where(filters: dict[str, str | list[str]] | None) -> tuple[str, dict]:
    clauses, params = [], {}
    for i, (col, value) in enumerate((filters or {}).items()):
        if col not in FILTERABLE:  # column names go into SQL text, so allowlist them
            raise ValueError(f"cannot filter on {col!r}")
        params[f"f{i}"] = [value] if isinstance(value, str) else list(value)
        clauses.append(f"{col} = ANY(%(f{i})s)")
    return ("WHERE " + " AND ".join(clauses)) if clauses else "", params


if __name__ == "__main__":
    import support  # module import keeps the function support.chunk apart from the field Hit.chunk

    index = Index(connect())
    index.create()
    print(index.upsert([c for doc in support.load_corpus() for c in support.chunk(doc)]), f"· {index.count()} rows")
    for hit in index.search("why does my job run twice?", k=3, filters={"applies_to": ["3.x", "all"]}):
        print(f"{hit.score:.3f}  {hit.chunk.doc_id:28} {hit.chunk.header[:45]}")
    search_docs = lambda q, k: support.ranked_doc_ids([h.chunk for h in index.search(q, k=4 * k)])
    print("hit@5", support.hit_rate(support.load_retrieval_cases("dev"), search_docs))

367 inserted · 0 unchanged · 0 deleted · $0.00131 · 367 rows
0.684  troubleshooting-duplicates   Troubleshooting: jobs running twice > Cause 1
0.641  leases-and-heartbeats        Leases and heartbeats > Why duplicate executi
0.612  idempotency                  Idempotency > What happens on delivery
hit@5 {'overall': 0.9, 'exact_term': 0.75, 'semantic': 1.0, 'multi_hop': 1.0}

0 inserted · 367 unchanged · 0 deleted · $0.00000 · 367 rows

# script_id: day_026_pgvector_search/memory_index
# memory_index.py: same interface as Index, in RAM. No Postgres; runs offline with fake_embed.
# Install: uv add numpy pydantic "psycopg[binary]" pgvector (index.py imports them). Env: ORBIT_DATA
from __future__ import annotations

from typing import Sequence

import numpy as np

from index import EmbedFn, Hit, UpsertStats, chunk_id
from support import EMBED_DIM, TRACKER, DocChunk, embed, embed_text


class MemoryIndex:
    def __init__(self, *, dim: int = EMBED_DIM, embed_fn: EmbedFn = embed) -> None:
        self.dim, self.embed_fn = dim, embed_fn
        self.rows: dict[str, tuple[DocChunk, np.ndarray]] = {}

    def create(self, hnsw: bool = True) -> None:
        """Nothing to create; kept so callers can treat both indexes the same way."""

    def count(self) -> int:
        return len(self.rows)

    def upsert(self, chunks: Sequence[DocChunk]) -> UpsertStats:
        keyed = {chunk_id(c.doc_id, embed_text(c)): c for c in chunks}
        doc_ids = {c.doc_id for c in chunks}
        existing = {cid for cid, (c, _) in self.rows.items() if c.doc_id in doc_ids}
        new = [cid for cid in keyed if cid not in existing]
        before = TRACKER.total_usd
        vectors = self.embed_fn([embed_text(keyed[cid]) for cid in new]) if new else []
        for cid, vec in zip(new, vectors):
            self.rows[cid] = (keyed[cid].model_copy(update={"id": cid}), np.asarray(vec))
        for cid in existing & keyed.keys():  # refresh position/metadata, keep the vector
            self.rows[cid] = (keyed[cid].model_copy(update={"id": cid}), self.rows[cid][1])
        stale = existing - keyed.keys()
        for cid in stale:
            del self.rows[cid]
        return UpsertStats(len(new), len(keyed) - len(new), len(stale), TRACKER.total_usd - before)

# script_id: day_026_pgvector_search/memory_index
    def delete(self, doc_id: str) -> int:
        gone = [cid for cid, (c, _) in self.rows.items() if c.doc_id == doc_id]
        for cid in gone:
            del self.rows[cid]
        return len(gone)

    def search(self, query: str, k: int = 5, *,
               filters: dict[str, str | list[str]] | None = None) -> list[Hit]:
        allowed = {col: {v} if isinstance(v, str) else set(v) for col, v in (filters or {}).items()}
        items = [(c, v) for c, v in self.rows.values()
                 if all(getattr(c, col) in vals for col, vals in allowed.items())]
        if not items:
            return []
        sims = np.stack([v for _, v in items]) @ self.embed_fn([query])[0]  # unit vectors: dot = cosine
        return [Hit(items[i][0], float(sims[i])) for i in np.argsort(-sims)[:k]]


if __name__ == "__main__":
    from support import chunk, fake_embed, hit_rate, load_corpus, load_retrieval_cases, ranked_doc_ids

    index = MemoryIndex(embed_fn=fake_embed)  # swap in embed for real numbers
    chunks = [c for doc in load_corpus() for c in chunk(doc)]
    print(index.upsert(chunks))
    print(index.upsert(chunks))  # idempotent: nothing new to embed
    for hit in index.search("ORB-1301 RetryExhausted", k=3, filters={"applies_to": ["3.x", "all"]}):
        print(f"{hit.score:.3f}  {hit.chunk.doc_id}")
    search_docs = lambda q, k: ranked_doc_ids([h.chunk for h in index.search(q, k=4 * k)])
    print("hit@5 (fake_embed)", hit_rate(load_retrieval_cases("dev"), search_docs))

367 inserted · 0 unchanged · 0 deleted · $0.00000
0 inserted · 367 unchanged · 0 deleted · $0.00000
0.299  troubleshooting-broker
0.239  dead-letter-queue
0.233  migrating-2-to-3
hit@5 (fake_embed) {'overall': 1.0, 'exact_term': 1.0, 'semantic': 1.0, 'multi_hop': 1.0}

# script_id: day_026_pgvector_search/reingest
# Edit one page, re-ingest, and watch only the changed chunk get re-embedded. Then restore it.
# Env: ATLAS_INDEX=memory (no Postgres) or DATABASE_URL; OPENAI_API_KEY; ORBIT_DATA
import os
import shutil
import tempfile
from pathlib import Path

from index import Index, connect
from memory_index import MemoryIndex
from support import DATA_DIR, chunk, load_corpus


def build_index() -> Index | MemoryIndex:
    if os.environ.get("ATLAS_INDEX") == "memory":
        return MemoryIndex()
    index = Index(connect())
    index.create()
    return index


def ingest(index: Index | MemoryIndex, docs_dir: Path, label: str) -> None:
    print(f"{label:9}", index.upsert([c for doc in load_corpus(docs_dir) for c in chunk(doc)]))


with tempfile.TemporaryDirectory() as tmp:
    docs = Path(tmp) / "docs"
    shutil.copytree(DATA_DIR / "docs", docs)
    index = build_index()
    ingest(index, docs, "baseline")
    page = docs / "retries-and-backoff.md"
    original = page.read_text(encoding="utf-8")
    edited = original.replace("These are also the built-in defaults", "These match the built-in defaults", 1)
    assert edited != original, "sentence not found; pick another one to edit"
    page.write_text(edited, encoding="utf-8")
    ingest(index, docs, "edited")
    page.write_text(original, encoding="utf-8")
    ingest(index, docs, "restored")

baseline  0 inserted · 367 unchanged · 0 deleted · $0.00000
edited    1 inserted · 366 unchanged · 1 deleted · $0.00000
restored  1 inserted · 366 unchanged · 1 deleted · $0.00000

# script_id: day_026_pgvector_search/ann_bench
# Synthetic ANN benchmark: recall@10 vs exact numpy search, ms/query, and the filtered-query gotcha.
# ~1 minute on a laptop. Synthetic data, so your numbers will differ.
# Install: uv add "psycopg[binary]" pgvector numpy · Env: DATABASE_URL · Services: Postgres (compose.yaml)
import os
import time

import numpy as np
import psycopg
from pgvector.psycopg import register_vector

N, DIM, K, N_QUERIES = 50_000, 128, 10, 50
rng = np.random.default_rng(0)
centers = rng.normal(size=(1000, DIM))
X = centers[rng.integers(0, 1000, N)] + 0.3 * rng.normal(size=(N, DIM))  # clustered, like real embeddings
X = (X / np.linalg.norm(X, axis=1, keepdims=True)).astype(np.float32)
Q = X[rng.choice(N, N_QUERIES, replace=False)] + 0.1 * rng.normal(size=(N_QUERIES, DIM))
Q = (Q / np.linalg.norm(Q, axis=1, keepdims=True)).astype(np.float32)
tags = rng.integers(0, 10, N)  # each tag value matches ~10% of rows
tag3 = np.flatnonzero(tags == 3)
truth = [set(np.argsort(-(X @ q))[:K].tolist()) for q in Q]  # exact top-10 by cosine
truth_tag3 = [set(tag3[np.argsort(-(X[tag3] @ q))[:K]].tolist()) for q in Q]

conn = psycopg.connect(os.environ.get("DATABASE_URL", "postgresql://atlas:atlas@localhost:5432/atlas"),
                       autocommit=True)
conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
register_vector(conn)
conn.execute("DROP TABLE IF EXISTS bench")
conn.execute(f"CREATE TABLE bench (id int PRIMARY KEY, tag int NOT NULL, embedding vector({DIM}) NOT NULL)")
with conn.transaction(), conn.cursor() as cur:  # one commit for all rows, not 50,000
    cur.executemany("INSERT INTO bench VALUES (%s, %s, %s)", [(i, int(tags[i]), X[i]) for i in range(N)])

# script_id: day_026_pgvector_search/ann_bench
TOP10 = "SELECT id FROM bench ORDER BY embedding <=> %s LIMIT 10"
TOP10_TAG3 = "SELECT id FROM bench WHERE tag = 3 ORDER BY embedding <=> %s LIMIT 10"


def run(label: str, expected: list[set[int]], sql: str = TOP10) -> None:
    start, recalls, counts = time.perf_counter(), [], []
    for q, want in zip(Q, expected):
        got = {row[0] for row in conn.execute(sql, (q,))}
        recalls.append(len(want & got) / K)
        counts.append(len(got))
    ms = (time.perf_counter() - start) * 1000 / len(Q)
    print(f"{label:30} recall@10={np.mean(recalls):.2f}  rows={np.mean(counts):4.1f}  {ms:6.2f} ms/query")


run("exact scan (no index)", truth)
conn.execute("CREATE INDEX bench_hnsw ON bench USING hnsw (embedding vector_cosine_ops)")
for ef in (10, 40, 100):
    conn.execute(f"SET hnsw.ef_search = {ef}")  # SET can't take %s parameters; ef is our own int
    run(f"hnsw ef_search={ef}", truth)
conn.execute("SET hnsw.ef_search = 40")
conn.execute("SET enable_seqscan = off")  # force the index so the filter gotcha shows on every machine
run("hnsw + WHERE tag = 3", truth_tag3, TOP10_TAG3)
conn.execute("SET hnsw.iterative_scan = relaxed_order")  # pgvector 0.8+
run("hnsw + WHERE + iterative scan", truth_tag3, TOP10_TAG3)
conn.execute("RESET ALL")
conn.execute("DROP INDEX bench_hnsw")
conn.execute("CREATE INDEX bench_ivf ON bench USING ivfflat (embedding vector_cosine_ops) WITH (lists = 50)")
for probes in (1, 7, 20):
    conn.execute(f"SET ivfflat.probes = {probes}")
    run(f"ivfflat lists=50 probes={probes}", truth)
conn.close()

exact scan (no index)          recall@10=1.00  rows=10.0    6.85 ms/query
hnsw ef_search=10              recall@10=0.54  rows=10.0    0.23 ms/query
hnsw ef_search=40              recall@10=0.70  rows=10.0    0.85 ms/query
hnsw ef_search=100             recall@10=0.88  rows=10.0    0.82 ms/query
hnsw + WHERE tag = 3           recall@10=0.32  rows= 3.9    0.29 ms/query
ivfflat lists=50 probes=1      recall@10=0.62  rows=10.0    0.22 ms/query
ivfflat lists=50 probes=7      recall@10=1.00  rows=10.0    1.14 ms/query
ivfflat lists=50 probes=20     recall@10=1.00  rows=10.0    3.22 ms/query