Pythonpgvector_search.py
Copy
1
import psycopg2
2
import random
3
import time
4
 
5
DIM = 384  # typical sentence-transformer dimension
6
 
7
conn = psycopg2.connect("host=localhost dbname=learndb user=postgres")
8
conn.autocommit = True
9
cur = conn.cursor()
10
 
11
# Setup
12
cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
13
cur.execute("DROP TABLE IF EXISTS documents;")
14
cur.execute(f"""
15
CREATE TABLE documents (
16
    id        SERIAL PRIMARY KEY,
17
    title     TEXT NOT NULL,
18
    category  TEXT NOT NULL,
19
    embedding vector({DIM}) NOT NULL
20
);
21
""")
22
 
23
# Insert sample documents with random embeddings (replace with real embeddings)
24
def random_vector(dim: int) -> list[float]:
25
    v = [random.gauss(0, 1) for _ in range(dim)]
26
    norm = sum(x * x for x in v) ** 0.5
27
    return [x / norm for x in v]  # L2 normalized
28
 
29
categories = ["tech", "health", "finance"]
30
rows = [(f"Document {i}", random.choice(categories), random_vector(DIM)) for i in range(1000)]
31
 
32
cur.executemany(
33
    "INSERT INTO documents (title, category, embedding) VALUES (%s, %s, %s::vector)",
34
    [(title, cat, str(vec)) for title, cat, vec in rows],
35
)
36
print("Inserted 1,000 documents")
37
 
38
# Create HNSW index (approximate, fast)
39
cur.execute(f"""
40
CREATE INDEX IF NOT EXISTS docs_hnsw_idx
41
ON documents
42
USING hnsw (embedding vector_cosine_ops)
43
WITH (m = 16, ef_construction = 64);
44
""")
45
print("HNSW index created")
46
 
47
# KNN search (5 nearest to random query vector)
48
query_vec = random_vector(DIM)
49
start = time.perf_counter()
50
cur.execute(f"""
51
SELECT id, title, category,
52
       1 - (embedding <=> %s::vector) AS cosine_similarity
53
FROM documents
54
WHERE category = 'tech'
55
ORDER BY embedding <=> %s::vector
56
LIMIT 5;
57
""", (str(query_vec), str(query_vec)))
58
elapsed = (time.perf_counter() - start) * 1000
59
results = cur.fetchall()
60
print(f"Top 5 nearest 'tech' docs ({elapsed:.1f}ms):")
61
for row in results:
62
    print(f"  id={row[0]}, title={row[1]}, similarity={row[3]:.4f}")
63
 
64
cur.close()
65
conn.close()
66