sqlite_vec_virtual_table_vector_storage_knn_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 byte format for sqlite-vec."""
8    return struct.pack(f"{len(vector)}f", *vector)
9
10# Connect to an in-memory database
11db = sqlite3.connect(":memory:")
12
13# Load the sqlite-vec extension
14db.enable_load_extension(True)
15sqlite_vec.load(db)
16db.enable_load_extension(False)
17
18# Verify the extension loaded correctly
19vec_version, = db.execute("select vec_version()").fetchone()
20print(f"sqlite-vec version: {vec_version}")
21
22# Create a virtual table with a 3-dimensional vector column
23db.execute("""
24  CREATE VIRTUAL TABLE vec_items USING vec0(
25    embedding float[3]
26  );
27""")
28
29# Insert sample data
30items = [
31    (1, [0.1, 0.1, 0.1]),
32    (2, [0.2, 0.2, 0.2]),
33    (3, [0.3, 0.3, 0.3]),
34]
35
36with db:
37    for rowid, vec in items:
38        db.execute(
39            "INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)",
40            [rowid, serialize_f32(vec)],
41        )
42
43# Perform a k-Nearest Neighbor search
44query_vector = [0.12, 0.12, 0.12]
45rows = db.execute("""
46  SELECT
47    rowid,
48    distance
49  FROM vec_items
50  WHERE embedding MATCH ?
51  ORDER BY distance
52  LIMIT 3
53""", [serialize_f32(query_vector)]).fetchall()
54
55for rowid, distance in rows:
56    print(f"Row ID: {rowid}, Distance: {distance}")