mirror of
https://github.com/asg017/sqlite-vec.git
synced 2026-04-25 00:36:56 +02:00
Add ANN search support for vec0 virtual table (#273)
Add approximate nearest neighbor infrastructure to vec0: shared distance dispatch (vec0_distance_full), flat index type with parser, NEON-optimized cosine/Hamming for float32/int8, amalgamation script, and benchmark suite (benchmarks-ann/) with ground-truth generation and profiling tools. Remove unused vec_npy_each/vec_static_blobs code, fix missing stdint.h include.
This commit is contained in:
parent
e9f598abfa
commit
0de765f457
27 changed files with 2177 additions and 2116 deletions
|
|
@ -248,59 +248,6 @@ def bench_libsql(base, query, page_size, k) -> BenchResult:
|
|||
return BenchResult(f"libsql ({page_size})", build_time, times)
|
||||
|
||||
|
||||
def register_np(db, array, name):
|
||||
ptr = array.__array_interface__["data"][0]
|
||||
nvectors, dimensions = array.__array_interface__["shape"]
|
||||
element_type = array.__array_interface__["typestr"]
|
||||
|
||||
assert element_type == "<f4"
|
||||
|
||||
name_escaped = db.execute("select printf('%w', ?)", [name]).fetchone()[0]
|
||||
|
||||
db.execute(
|
||||
"insert into temp.vec_static_blobs(name, data) select ?, vec_static_blob_from_raw(?, ?, ?, ?)",
|
||||
[name, ptr, element_type, dimensions, nvectors],
|
||||
)
|
||||
|
||||
db.execute(
|
||||
f'create virtual table "{name_escaped}" using vec_static_blob_entries({name_escaped})'
|
||||
)
|
||||
|
||||
def bench_sqlite_vec_static(base, query, k) -> BenchResult:
|
||||
print(f"sqlite-vec static...")
|
||||
|
||||
db = sqlite3.connect(":memory:")
|
||||
db.enable_load_extension(True)
|
||||
db.load_extension("../../dist/vec0")
|
||||
|
||||
|
||||
|
||||
t = time.time()
|
||||
register_np(db, base, "base")
|
||||
build_time = time.time() - t
|
||||
|
||||
times = []
|
||||
results = []
|
||||
for (
|
||||
idx,
|
||||
q,
|
||||
) in enumerate(query):
|
||||
t0 = time.time()
|
||||
result = db.execute(
|
||||
"""
|
||||
select
|
||||
rowid
|
||||
from base
|
||||
where vector match ?
|
||||
and k = ?
|
||||
order by distance
|
||||
""",
|
||||
[q.tobytes(), k],
|
||||
).fetchall()
|
||||
assert len(result) == k
|
||||
times.append(time.time() - t0)
|
||||
return BenchResult(f"sqlite-vec static", build_time, times)
|
||||
|
||||
def bench_faiss(base, query, k) -> BenchResult:
|
||||
import faiss
|
||||
dimensions = base.shape[1]
|
||||
|
|
@ -438,8 +385,6 @@ def suite(name, base, query, k, benchmarks):
|
|||
for b in benchmarks:
|
||||
if b == "faiss":
|
||||
results.append(bench_faiss(base, query, k=k))
|
||||
elif b == "vec-static":
|
||||
results.append(bench_sqlite_vec_static(base, query, k=k))
|
||||
elif b.startswith("vec-scalar"):
|
||||
_, page_size = b.split('.')
|
||||
results.append(bench_sqlite_vec_scalar(base, query, page_size, k=k))
|
||||
|
|
@ -541,7 +486,7 @@ def parse_args():
|
|||
help="Number of queries to use. Defaults all",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-x", help="type of runs to make", default="faiss,vec-scalar.4096,vec-static,vec-vec0.4096.16,usearch,duckdb,hnswlib,numpy"
|
||||
"-x", help="type of runs to make", default="faiss,vec-scalar.4096,vec-vec0.4096.16,usearch,duckdb,hnswlib,numpy"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
|
|
|||
|
|
@ -8,10 +8,3 @@ create virtual table vec_items using vec0(
|
|||
embedding float[1536]
|
||||
);
|
||||
|
||||
-- 65s (limit 1e5), ~615MB on disk
|
||||
insert into vec_items
|
||||
select
|
||||
rowid,
|
||||
vector
|
||||
from vec_npy_each(vec_npy_file('examples/dbpedia-openai/data/vectors.npy'))
|
||||
limit 1e5;
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@ def connect(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
|
||||
|
||||
|
|
@ -18,8 +17,6 @@ page_sizes = [ # 4096, 8192,
|
|||
chunk_sizes = [128, 256, 1024, 2048]
|
||||
types = ["f32", "int8", "bit"]
|
||||
|
||||
SRC = "../examples/dbpedia-openai/data/vectors.npy"
|
||||
|
||||
for page_size in page_sizes:
|
||||
for chunk_size in chunk_sizes:
|
||||
for t in types:
|
||||
|
|
@ -42,15 +39,8 @@ for page_size in page_sizes:
|
|||
func = "vec_quantize_i8(vector, 'unit')"
|
||||
if t == "bit":
|
||||
func = "vec_quantize_binary(vector)"
|
||||
db.execute(
|
||||
f"""
|
||||
insert into vec_items
|
||||
select rowid, {func}
|
||||
from vec_npy_each(vec_npy_file(?))
|
||||
limit 100000
|
||||
""",
|
||||
[SRC],
|
||||
)
|
||||
# TODO: replace with non-npy data loading
|
||||
pass
|
||||
elapsed = time.time() - t0
|
||||
print(elapsed)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue