scale_bench.py

"""scale benchmark: indexed (sqlite-vec) vs brute-force (numpy).

honest version. it measures SEARCH ONLY, with the query vector precomputed, so
we compare the thing that actually differs. (embedding the query costs the same
for both and dominates at small sizes, so wall-clock recall is identical until
the corpus is large.) run at several sizes to see where the index starts to win.
"""
import os
import tempfile
import time

os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ["IGLEGAIS_DB"] = os.path.join(tempfile.mkdtemp(), "scale.db")

import numpy as np
from iglegais import LocalMemoryGraph

mg = LocalMemoryGraph()
mg.setup()
print("sqlite-vec index active:", mg._vec, "\n")

rng = np.random.default_rng(0)
DIM = 384
# one real-ish needle vector plus a big pile of random normalized vectors.
needle = rng.normal(0, 1, DIM).astype(np.float32); needle /= np.linalg.norm(needle)
qv = needle.copy()  # query the needle exactly


def add_rows(n, start_id):
    rows = rng.normal(0, 1, (n, DIM)).astype(np.float32)
    rows /= np.linalg.norm(rows, axis=1, keepdims=True)
    for i in range(n):
        blob = rows[i].tobytes()
        cur = mg.conn.execute(
            "INSERT INTO memories(content, embedding, ts, namespace) VALUES(?,?,?,?)",
            (f"row {start_id+i}", blob, i, "default"))
        if mg._vec:
            mg.conn.execute("INSERT INTO vec_memories(rowid, namespace, embedding) VALUES(?,?,?)",
                            (cur.lastrowid, "default", blob))
    mg.conn.commit()


def search_numpy(k=3):
    rows = mg.conn.execute("SELECT id, content, embedding FROM memories WHERE namespace='default'").fetchall()
    scored = [(float(qv @ np.frombuffer(b, dtype=np.float32)), i, c) for i, c, b in rows]
    scored.sort(key=lambda t: t[0], reverse=True)
    return scored[:k]


def search_vec(k=3):
    return mg.conn.execute(
        "SELECT v.rowid, v.distance FROM vec_memories v "
        "WHERE v.namespace='default' AND v.embedding MATCH ? ORDER BY v.distance LIMIT ?",
        (qv.tobytes(), k)).fetchall()


def bench(fn, runs=15):
    fn()  # warm
    t0 = time.time()
    for _ in range(runs):
        fn()
    return (time.time() - t0) / runs * 1000


# insert the needle first so it is always in the corpus
mg.conn.execute("INSERT INTO memories(content, embedding, ts, namespace) VALUES(?,?,?,?)",
                ("NEEDLE", needle.tobytes(), 0, "default"))
if mg._vec:
    mg.conn.execute("INSERT INTO vec_memories(rowid, namespace, embedding) VALUES(?,?,?)",
                    (1, "default", needle.tobytes()))
mg.conn.commit()

total = 1
print(f"{'corpus':>10} {'numpy ms':>10} {'sqlite-vec ms':>14} {'speedup':>9}")
for target in [5_000, 25_000, 100_000, 300_000]:
    add_rows(target - total, total)
    total = target
    np_ms = bench(search_numpy)
    if mg._vec:
        vec_ms = bench(search_vec)
        print(f"{total:>10,} {np_ms:>10.1f} {vec_ms:>14.2f} {np_ms/vec_ms:>8.1f}x")
    else:
        print(f"{total:>10,} {np_ms:>10.1f} {'n/a':>14} {'n/a':>9}")

# correctness: both return the needle as the top hit
if mg._vec:
    print("\ncorrectness: numpy top id =", search_numpy()[0][1],
          "| sqlite-vec top id =", search_vec()[0][0], "(1 = the needle)")