Initial commit

benchmarks/self-params/knn.py

@@ -0,0 +1,83 @@
+import sqlite3
+import time
+from random import randrange
+from statistics import mean
+
+
+def connect(path):
+    print(path)
+    db = sqlite3.connect(path)
+    db.enable_load_extension(True)
+    db.load_extension("../dist/vec0")
+    db.execute("select load_extension('../dist/vec0', 'sqlite3_vec_fs_read_init')")
+    db.enable_load_extension(False)
+    return db
+
+
+page_sizes = [  # 4096, 8192,
+    16384,
+    32768,
+]
+chunk_sizes = [128, 256, 1024, 2048]
+types = ["f32", "int8", "bit"]
+
+types.reverse()
+
+for t in types:
+    for page_size in page_sizes:
+        for chunk_size in chunk_sizes:
+            print(f"page_size={page_size}, chunk_size={chunk_size}")
+
+            func = "embedding"
+            if t == "int8":
+                func = "vec_quantize_i8(embedding, 'unit')"
+            if t == "bit":
+                func = "vec_quantize_binary(embedding)"
+
+            times = []
+            trials = 20
+            db = connect(f"dbs/test.{page_size}.{chunk_size}.{t}.db")
+
+            for trial in range(trials):
+                t0 = time.time()
+                results = db.execute(
+                    f"""
+                      select rowid
+                      from vec_items
+                      where embedding match (select {func} from vec_items where rowid = ?)
+                        and k = 10
+                      order by distance
+                    """,
+                    [randrange(100000)],
+                ).fetchall()
+
+                times.append(time.time() - t0)
+            print(mean(times))
+
+"""
+
+page_size=4096, chunk_size=256
+0.2635102152824402
+page_size=4096, chunk_size=1024
+0.2609449863433838
+page_size=4096, chunk_size=2048
+0.275589919090271
+page_size=8192, chunk_size=256
+0.18621582984924318
+page_size=8192, chunk_size=1024
+0.20939643383026124
+page_size=8192, chunk_size=2048
+0.22376316785812378
+page_size=16384, chunk_size=256
+0.16012665033340454
+page_size=16384, chunk_size=1024
+0.18346318006515502
+page_size=16384, chunk_size=2048
+0.18224761486053467
+page_size=32768, chunk_size=256
+0.14202518463134767
+page_size=32768, chunk_size=1024
+0.15340715646743774
+page_size=32768, chunk_size=2048
+0.18018823862075806
+"""