Merge pull request #1661 from tursodatabase/vector-search-full-suppor…

…t-all-column-types vector search: full support all column types
tursodatabase · Aug 14, 2024 · fd1f5c4 · fd1f5c4
2 parents d86b5b7 + 8604065
commit fd1f5c4
Show file tree

Hide file tree

Showing 15 changed files with 1,371 additions and 669 deletions.
diff --git a/libsql-ffi/bundled/SQLite3MultipleCiphers/src/sqlite3.c b/libsql-ffi/bundled/SQLite3MultipleCiphers/src/sqlite3.c
diff --git a/libsql-ffi/bundled/src/sqlite3.c b/libsql-ffi/bundled/src/sqlite3.c
diff --git a/libsql-sqlite3/Makefile.in b/libsql-sqlite3/Makefile.in
@@ -195,7 +195,7 @@ LIBOBJS0 = alter.lo analyze.lo attach.lo auth.lo \
          sqlite3session.lo select.lo sqlite3rbu.lo status.lo stmt.lo \
          table.lo threads.lo tokenize.lo treeview.lo trigger.lo \
          update.lo userauth.lo upsert.lo util.lo vacuum.lo \
-         vector.lo vectorfloat32.lo vectorfloat64.lo vector1bit.lo \
+         vector.lo vectorfloat32.lo vectorfloat64.lo vectorfloat1bit.lo \
          vectorIndex.lo vectordiskann.lo vectorvtab.lo \
          vdbe.lo vdbeapi.lo vdbeaux.lo vdbeblob.lo vdbemem.lo vdbesort.lo \
          vdbetrace.lo vdbevtab.lo \
@@ -302,8 +302,8 @@ SRC = \
   $(TOP)/src/util.c \
   $(TOP)/src/vacuum.c \
   $(TOP)/src/vector.c \
-  $(TOP)/src/vector1bit.c \
   $(TOP)/src/vectorInt.h \
+  $(TOP)/src/vectorfloat1bit.c \
   $(TOP)/src/vectorfloat32.c \
   $(TOP)/src/vectorfloat64.c \
   $(TOP)/src/vectorIndexInt.h \
@@ -1139,8 +1139,8 @@ vacuum.lo:	$(TOP)/src/vacuum.c $(HDR)
 vector.lo:	$(TOP)/src/vector.c $(HDR)
 	$(LTCOMPILE) $(TEMP_STORE) -c $(TOP)/src/vector.c
 
-vector1bit.lo:	$(TOP)/src/vector1bit.c $(HDR)
-	$(LTCOMPILE) $(TEMP_STORE) -c $(TOP)/src/vector1bit.c
+vectorfloat1bit.lo:	$(TOP)/src/vectorfloat1bit.c $(HDR)
+	$(LTCOMPILE) $(TEMP_STORE) -c $(TOP)/src/vectorfloat1bit.c
 
 vectorfloat32.lo:	$(TOP)/src/vectorfloat32.c $(HDR)
 	$(LTCOMPILE) $(TEMP_STORE) -c $(TOP)/src/vectorfloat32.c

diff --git a/libsql-sqlite3/benchmark/workload.py b/libsql-sqlite3/benchmark/workload.py
@@ -10,10 +10,10 @@ def recall_uniform(dim, n, q):
     print(f'CREATE TABLE queries ( emb FLOAT32({dim}) );')
     print(f'BEGIN TRANSACTION;')
     for i in range(n):
-        vector = f"[{','.join(map(str, np.random.uniform(size=dim)))}]"
+        vector = f"[{','.join(map(str, np.random.uniform(-1, 1, size=dim)))}]"
         print(f'INSERT INTO data VALUES ({i}, vector(\'{vector}\'));')
     for i in range(q):
-        vector = f"[{','.join(map(str, np.random.uniform(size=dim)))}]"
+        vector = f"[{','.join(map(str, np.random.uniform(-1, 1, size=dim)))}]"
         print(f'INSERT INTO queries VALUES (vector(\'{vector}\'));')
     print(f'COMMIT;')
     print('---insert everything')
@@ -29,7 +29,7 @@ def recall_normal(dim, n, q):
         vector = f"[{','.join(map(str, np.random.uniform(size=64)))}]"
         print(f'INSERT INTO data VALUES ({i}, \'{vector}\');')
     for i in range(q):
-        vector = f"[{','.join(map(str, np.random.uniform(size=64)))}]"
+        vector = f"[{','.join(map(str, np.random.uniform(-1, 1, size=64)))}]"
         print(f'INSERT INTO queries VALUES (\'{vector}\');')
     print(f'COMMIT;')
     print('---insert everything')
@@ -40,7 +40,7 @@ def no_vectors(n, q):
     print('PRAGMA journal_mode=WAL;')
     print(f'CREATE TABLE x ( id INTEGER PRIMARY KEY, value TEXT );')
     for i in range(n):
-        vector = f"[{','.join(map(str, np.random.uniform(size=64)))}]"
+        vector = f"[{','.join(map(str, np.random.uniform(-1, 1, size=64)))}]"
         print(f'INSERT INTO x VALUES ({i}, \'{vector}\');')
     print('---inserts')
     for i in range(q):
@@ -54,11 +54,11 @@ def bruteforce(dim, n, q):
     print('PRAGMA journal_mode=WAL;')
     print(f'CREATE TABLE x ( id INTEGER PRIMARY KEY, embedding FLOAT32({dim}) );')
     for i in range(n):
-        vector = f"[{','.join(map(str, np.random.uniform(size=dim)))}]"
+        vector = f"[{','.join(map(str, np.random.uniform(-1, 1, size=dim)))}]"
         print(f'INSERT INTO x VALUES ({i}, vector(\'{vector}\'));')
     print('---inserts')
     for i in range(q):
-        vector = f"[{','.join(map(str, np.random.uniform(size=dim)))}]"
+        vector = f"[{','.join(map(str, np.random.uniform(-1, 1, size=dim)))}]"
         print(f'SELECT id FROM x ORDER BY vector_distance_cos(embedding, vector(\'{vector}\')) LIMIT 1;')
     print('---search')
 
@@ -68,13 +68,13 @@ def diskann(dim, n, q):
     q = int(q)
     print('PRAGMA journal_mode=WAL;')
     print(f'CREATE TABLE x ( id INTEGER PRIMARY KEY, embedding FLOAT32({dim}) );')
-    print(f'CREATE INDEX x_idx ON x( libsql_vector_idx(embedding) );')
+    print(f"CREATE INDEX x_idx ON x( libsql_vector_idx(embedding) );")
     for i in range(n):
-        vector = f"[{','.join(map(str, np.random.uniform(size=dim)))}]"
+        vector = f"[{','.join(map(str, np.random.uniform(-1, 1, size=dim)))}]"
         print(f'INSERT INTO x VALUES ({i}, vector(\'{vector}\'));')
     print('---inserts')
     for i in range(q):
-        vector = f"[{','.join(map(str, np.random.uniform(size=dim)))}]"
+        vector = f"[{','.join(map(str, np.random.uniform(-1, 1, size=dim)))}]"
         print(f'SELECT id FROM vector_top_k(\'x_idx\', vector(\'{vector}\'), 1);')
     print('---search')