forked from OSKYCX/mxokdzntjytqyjzywjz
76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
"""
|
||
sqlite-vec 百万级检索 benchmark（Linux 端侧验证，自包含）
|
||
=======================================================
|
||
|
||
验证题目要求 4「端侧 sqlite-vec 百万级 ≤500ms」——macOS 因系统 sqlite3 不支持 load_extension
|
||
无法验证，故在 Linux（银河麒麟同内核）实测。纯 sqlite-vec + numpy，不依赖嵌入模型，单文件可跑。
|
||
|
||
用法::
|
||
|
||
pip install sqlite-vec numpy
|
||
python3 sqlite_vec_bench.py [N] # N 默认 1000000
|
||
"""
|
||
|
||
import struct
|
||
import sys
|
||
import time
|
||
|
||
import numpy as np
|
||
|
||
try:
|
||
import sqlite3
|
||
import sqlite_vec
|
||
except ImportError:
|
||
print("需要: pip install sqlite-vec numpy")
|
||
sys.exit(1)
|
||
|
||
N = int(sys.argv[1]) if len(sys.argv) > 1 else 1_000_000
|
||
DIM = 512
|
||
|
||
|
||
def ser(v) -> bytes:
|
||
return struct.pack(f"{DIM}f", *v)
|
||
|
||
|
||
def main() -> int:
|
||
print(f"sqlite {sqlite3.sqlite_version} | 规模 N={N} | 维度 DIM={DIM}")
|
||
db = sqlite3.connect(":memory:")
|
||
db.enable_load_extension(True)
|
||
sqlite_vec.load(db)
|
||
db.enable_load_extension(False)
|
||
print("sqlite-vec 版本:", db.execute("select vec_version()").fetchone()[0])
|
||
|
||
db.execute(f"CREATE VIRTUAL TABLE vec USING vec0(embedding float[{DIM}])")
|
||
|
||
rng = np.random.default_rng(0)
|
||
print(f"插入 {N} 条随机向量...")
|
||
t0 = time.perf_counter()
|
||
batch = 10000
|
||
for i in range(0, N, batch):
|
||
m = min(batch, N - i)
|
||
mat = rng.standard_normal((m, DIM)).astype("float32")
|
||
rows = [(i + j, ser(mat[j].tolist())) for j in range(m)]
|
||
db.executemany("INSERT INTO vec(rowid, embedding) VALUES (?, ?)", rows)
|
||
db.commit()
|
||
print(f"插入耗时 {time.perf_counter() - t0:.1f}s")
|
||
|
||
# 检索延迟（100 次随机查询，top-10）
|
||
lat = []
|
||
for _ in range(100):
|
||
q = ser(rng.standard_normal(DIM).astype("float32").tolist())
|
||
t0 = time.perf_counter()
|
||
db.execute(
|
||
"SELECT rowid, distance FROM vec WHERE embedding MATCH ? AND k = 10 ORDER BY distance",
|
||
(q,),
|
||
).fetchall()
|
||
lat.append((time.perf_counter() - t0) * 1000)
|
||
lat.sort()
|
||
p50, p95, p99 = lat[50], lat[95], lat[99]
|
||
print(f"检索 P50={p50:.0f}ms P95={p95:.0f}ms P99={p99:.0f}ms (N={N})")
|
||
print(f"≤500ms 红线: {'✅ 达标' if p95 <= 500 else '⚠️ 超标（sqlite-vec 0.1.x 为暴力扫描，百万级需 HNSW/分片）'}")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|