sqlite_vec_virtual_table_knn_vector_search.py

1import sqlite3
2import sqlite_vec
3import struct
4from typing import List
5
6def serialize_f32(vector: List[float]) -> bytes:
7    """Serializes a list of floats into a compact bitstring for sqlite-vec."""
8    return struct.pack(f"{len(vector)}f", *vector)
9
10db = sqlite3.connect(":memory:")
11db.enable_load_extension(True)
12sqlite_vec.load(db)
13db.enable_load_extension(False)
14
15# Check that the extension loaded correctly
16vec_version, = db.execute("select vec_version()").fetchone()
17print(f"sqlite-vec version: {vec_version}")
18
19# Create a virtual table for 3-dimensional float vectors
20db.execute("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[3])")
21
22# Insert some sample data
23items = [
24    (1, [0.1, 0.1, 0.1]),
25    (2, [0.2, 0.2, 0.2]),
26    (3, [0.3, 0.3, 0.3]),
27]
28
29with db:
30    for rowid, vec in items:
31        db.execute(
32            "INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)",
33            [rowid, serialize_f32(vec)],
34        )
35
36# Perform a k-NN search (k=2) for vectors closest to [0.12, 0.12, 0.12]
37query_vec = serialize_f32([0.12, 0.12, 0.12])
38rows = db.execute(
39    """
40    SELECT
41      rowid,
42      distance
43    FROM vec_items
44    WHERE embedding MATCH ?
45    ORDER BY distance
46    LIMIT 2
47    """,
48    [query_vec],
49).fetchall()
50
51for rowid, distance in rows:
52    print(f"Result: rowid={rowid}, distance={distance}")