From e97e2e4078458ecda81d294bfcd9713b74cb8af3 Mon Sep 17 00:00:00 2001 From: thzxx Date: Sun, 5 Apr 2026 01:47:57 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20IVF=20=E5=80=92=E6=8E=92=E7=B4=A2?= =?UTF-8?q?=E5=BC=95=E4=BC=98=E5=8C=96=E5=90=91=E9=87=8F=E6=90=9C=E7=B4=A2?= =?UTF-8?q?=20-=20K-Means=20=E8=81=9A=E7=B1=BB=20+=20nprobe=20=E8=BF=91?= =?UTF-8?q?=E9=82=BB=E6=A3=80=E7=B4=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - N < 200 自动降级暴力搜索 - N ≥ 200 K-Means++ 聚类构建 IVF 索引 - nprobe=5 探测最近 5 个聚类，搜索量降 75% - 索引持久化到 IndexedDB，数据变化自动重建 - 内存向量缓存 + Map O(1) 查找 --- js/vector-store.js | 339 +++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 312 insertions(+), 27 deletions(-) diff --git a/js/vector-store.js b/js/vector-store.js index 70141f1..b0d0da1 100644 --- a/js/vector-store.js +++ b/js/vector-store.js @@ -1,18 +1,32 @@ /** * VectorStore - 向量存储与相似度检索 * - * 基于 IndexedDB 持久化，支持余弦相似度搜索 + * 优化：IVF 倒排索引 + * - 小数据集 (N < MIN_IVF_SIZE)：自动降级为暴力搜索 + * - 大数据集：K-Means 聚类 + nprobe 近邻搜索 + * - 索引在 addVectors() 后自动构建，持久化到 IndexedDB */ +const MIN_IVF_SIZE = 200; // 低于此数量不建索引，直接暴力搜 +const DEFAULT_K = 20; // 聚类中心数上限 +const DEFAULT_NPROBE = 5; // 搜索时探测的聚类数（覆盖 ~25% 数据） +const KMEANS_ITERS = 10; // K-Means 迭代次数 + export class VectorStore { constructor(dbName = 'metona-ollama-vectors') { this.dbName = dbName; this.db = null; + + // 内存索引缓存 { colId → { centroids, invertedLists, version } } + this._indexCache = new Map(); + + // 向量内存缓存 { colId → Map } + this._vectorCache = new Map(); } async init() { return new Promise((resolve, reject) => { - const req = indexedDB.open(this.dbName, 1); + const req = indexedDB.open(this.dbName, 2); req.onerror = () => reject(req.error); req.onsuccess = () => { this.db = req.result; resolve(); }; req.onupgradeneeded = (e) => { @@ -24,6 +38,9 @@ export class VectorStore { if (!db.objectStoreNames.contains('collections')) { db.createObjectStore('collections', { keyPath: 'id' }); } + if (!db.objectStoreNames.contains('indexes')) { + db.createObjectStore('indexes', { keyPath: 'collectionId' }); + } }; }); } @@ -32,17 +49,16 @@ export class VectorStore { return this.db.transaction(store, mode).objectStore(store); } - // ── 集合管理 ── + // ═══════════════════════════════════════════ + // 集合管理 + // ═══════════════════════════════════════════ async createCollection(name, embeddingModel = '') { const col = { id: `kb_${Date.now()}_${Math.random().toString(36).slice(2, 7)}`, - name, - embeddingModel, - docCount: 0, - chunkCount: 0, - createdAt: Date.now(), - updatedAt: Date.now() + name, embeddingModel, + docCount: 0, chunkCount: 0, + createdAt: Date.now(), updatedAt: Date.now() }; return new Promise((resolve, reject) => { const req = this._tx('collections', 'readwrite').put(col); @@ -77,79 +93,348 @@ export class VectorStore { } async deleteCollection(colId) { - // 删除集合下所有向量 const vectors = await this.getVectorsByCollection(colId); - const store = this._tx('vectors', 'readwrite'); - for (const v of vectors) store.delete(v.id); - return new Promise((resolve, reject) => { - const req = this._tx('collections', 'readwrite').delete(colId); - req.onsuccess = () => resolve(); - req.onerror = () => reject(req.error); + await new Promise((resolve, reject) => { + const tx = this.db.transaction(['vectors', 'collections', 'indexes'], 'readwrite'); + const vStore = tx.objectStore('vectors'); + for (const v of vectors) vStore.delete(v.id); + tx.objectStore('collections').delete(colId); + tx.objectStore('indexes').delete(colId); + tx.oncomplete = () => resolve(); + tx.onerror = () => reject(tx.error); }); + this._indexCache.delete(colId); + this._vectorCache.delete(colId); } - // ── 向量操作 ── + // ═══════════════════════════════════════════ + // 向量操作 + // ═══════════════════════════════════════════ async addVectors(items) { - return new Promise((resolve, reject) => { + if (items.length === 0) return; + + await new Promise((resolve, reject) => { const tx = this.db.transaction('vectors', 'readwrite'); const store = tx.objectStore('vectors'); for (const item of items) store.put(item); tx.oncomplete = () => resolve(); tx.onerror = () => reject(tx.error); }); + + // 更新内存缓存 + const colId = items[0].collectionId; + if (!this._vectorCache.has(colId)) { + this._vectorCache.set(colId, new Map()); + } + const cache = this._vectorCache.get(colId); + for (const item of items) cache.set(item.id, item); + + // 标记索引需重建 + this._indexCache.delete(colId); } async getVectorsByCollection(colId) { - return new Promise((resolve, reject) => { + if (this._vectorCache.has(colId)) { + return Array.from(this._vectorCache.get(colId).values()); + } + + const vectors = await new Promise((resolve, reject) => { const idx = this._tx('vectors').index('collectionId'); const req = idx.getAll(colId); req.onsuccess = () => resolve(req.result || []); req.onerror = () => reject(req.error); }); + + const cache = new Map(); + for (const v of vectors) cache.set(v.id, v); + this._vectorCache.set(colId, cache); + + return vectors; } async deleteVectorsByCollection(colId) { const vectors = await this.getVectorsByCollection(colId); - return new Promise((resolve, reject) => { + await new Promise((resolve, reject) => { const tx = this.db.transaction('vectors', 'readwrite'); const store = tx.objectStore('vectors'); for (const v of vectors) store.delete(v.id); tx.oncomplete = () => resolve(); tx.onerror = () => reject(tx.error); }); + this._indexCache.delete(colId); + this._vectorCache.delete(colId); } async deleteVectorsByDocument(colId, docId) { const vectors = await this.getVectorsByCollection(colId); const docVectors = vectors.filter(v => v.docId === docId); - return new Promise((resolve, reject) => { + await new Promise((resolve, reject) => { const tx = this.db.transaction('vectors', 'readwrite'); const store = tx.objectStore('vectors'); for (const v of docVectors) store.delete(v.id); tx.oncomplete = () => resolve(); tx.onerror = () => reject(tx.error); }); + if (this._vectorCache.has(colId)) { + const cache = this._vectorCache.get(colId); + for (const v of docVectors) cache.delete(v.id); + } + this._indexCache.delete(colId); } - // ── 相似度检索 ── + // ═══════════════════════════════════════════ + // 搜索：自动选择暴力 / IVF + // ═══════════════════════════════════════════ async search(colId, queryEmbedding, topK = 5) { const vectors = await this.getVectorsByCollection(colId); if (vectors.length === 0) return []; + // 小数据集 → 暴力搜索 + if (vectors.length < MIN_IVF_SIZE) { + return this._bruteForceSearch(vectors, queryEmbedding, topK); + } + + // 大数据集 → IVF 搜索 + const index = await this._getIndex(colId, vectors); + if (!index) { + return this._bruteForceSearch(vectors, queryEmbedding, topK); + } + + return this._ivfSearch(vectors, index, queryEmbedding, topK); + } + + _bruteForceSearch(vectors, queryEmbedding, topK) { return vectors - .map(v => ({ - ...v, - score: VectorStore.cosineSimilarity(queryEmbedding, v.embedding) - })) + .map(v => ({ ...v, score: VectorStore.cosineSimilarity(queryEmbedding, v.embedding) })) .sort((a, b) => b.score - a.score) .slice(0, topK); } + _ivfSearch(vectors, index, queryEmbedding, topK) { + const { centroids, invertedLists } = index; + + // 1. 找最近的 nprobe 个聚类中心 + const clusterScores = centroids.map((c, i) => ({ + id: i, + score: VectorStore.cosineSimilarity(queryEmbedding, c) + })); + clusterScores.sort((a, b) => b.score - a.score); + + const nprobe = Math.min(DEFAULT_NPROBE, centroids.length); + const targetClusters = clusterScores.slice(0, nprobe); + + // 2. 收集候选向量 ID + const candidateIds = new Set(); + for (const { id: clusterId } of targetClusters) { + const list = invertedLists.get(clusterId); + if (list) for (const id of list) candidateIds.add(id); + } + + // 3. 候选向量构建快速查找（避免逐个 find） + const vectorMap = new Map(); + for (const v of vectors) vectorMap.set(v.id, v); + + const candidates = []; + for (const id of candidateIds) { + const v = vectorMap.get(id); + if (v) candidates.push(v); + } + + return candidates + .map(v => ({ ...v, score: VectorStore.cosineSimilarity(queryEmbedding, v.embedding) })) + .sort((a, b) => b.score - a.score) + .slice(0, topK); + } + + // ═══════════════════════════════════════════ + // IVF 索引 + // ═══════════════════════════════════════════ + + async _getIndex(colId, vectors) { + // 内存缓存 + if (this._indexCache.has(colId)) { + const cached = this._indexCache.get(colId); + if (cached.version === this._indexVersion(vectors)) { + return cached; + } + } + + // IndexedDB 缓存 + const saved = await new Promise((resolve, reject) => { + const req = this._tx('indexes').get(colId); + req.onsuccess = () => resolve(req.result || null); + req.onerror = () => reject(req.error); + }); + + if (saved && saved.version === this._indexVersion(vectors)) { + // 反序列化 invertedLists（IndexedDB 不存 Map，存的是对象） + if (!(saved.invertedLists instanceof Map)) { + saved.invertedLists = new Map(Object.entries(saved.invertedLists || {})); + } + this._indexCache.set(colId, saved); + return saved; + } + + // 重建 + console.log(`[VectorStore] 构建 IVF 索引: ${colId} (${vectors.length} 向量)`); + const index = this._buildIVFIndex(vectors); + index.collectionId = colId; + index.version = this._indexVersion(vectors); + + // 持久化（Map → 普通对象） + const toSave = { + collectionId: index.collectionId, + version: index.version, + centroids: index.centroids, + invertedLists: Object.fromEntries(index.invertedLists) + }; + await new Promise((resolve, reject) => { + const req = this._tx('indexes', 'readwrite').put(toSave); + req.onsuccess = () => resolve(); + req.onerror = () => reject(req.error); + }); + + this._indexCache.set(colId, index); + return index; + } + + _indexVersion(vectors) { + if (vectors.length === 0) return '0'; + return `${vectors.length}_${vectors[0]?.id || ''}_${vectors[vectors.length - 1]?.id || ''}`; + } + /** - * 余弦相似度 + * K-Means 聚类构建 IVF 索引 */ + _buildIVFIndex(vectors) { + const N = vectors.length; + const dim = vectors[0].embedding.length; + + const K = Math.max(2, Math.min(DEFAULT_K, Math.floor(N / 10))); + + // 1. K-Means++ 初始化 + const centroids = this._kmeansPPInit(vectors, K, dim); + + // 2. K-Means 迭代 + const assignments = new Array(N); + for (let iter = 0; iter < KMEANS_ITERS; iter++) { + // E-step + for (let i = 0; i < N; i++) { + let bestCluster = 0; + let bestScore = -Infinity; + for (let k = 0; k < K; k++) { + const score = VectorStore.cosineSimilarity(vectors[i].embedding, centroids[k]); + if (score > bestScore) { + bestScore = score; + bestCluster = k; + } + } + assignments[i] = bestCluster; + } + + // M-step + const newCentroids = Array.from({ length: K }, () => new Array(dim).fill(0)); + const counts = new Array(K).fill(0); + + for (let i = 0; i < N; i++) { + const k = assignments[i]; + counts[k]++; + const emb = vectors[i].embedding; + for (let d = 0; d < dim; d++) { + newCentroids[k][d] += emb[d]; + } + } + + for (let k = 0; k < K; k++) { + if (counts[k] > 0) { + for (let d = 0; d < dim; d++) { + newCentroids[k][d] /= counts[k]; + } + this._normalize(newCentroids[k]); + } else { + const randIdx = Math.floor(Math.random() * N); + newCentroids[k] = [...vectors[randIdx].embedding]; + } + } + + // 收敛检测 + let converged = true; + for (let k = 0; k < K; k++) { + if (VectorStore.cosineSimilarity(centroids[k], newCentroids[k]) < 0.999) { + converged = false; + break; + } + } + + for (let k = 0; k < K; k++) { + centroids[k] = newCentroids[k]; + } + + if (converged) break; + } + + // 3. 构建倒排链表 + const invertedLists = new Map(); + for (let i = 0; i < N; i++) { + const k = assignments[i]; + if (!invertedLists.has(k)) invertedLists.set(k, []); + invertedLists.get(k).push(vectors[i].id); + } + + return { centroids, invertedLists }; + } + + /** + * K-Means++ 初始化 + */ + _kmeansPPInit(vectors, K, dim) { + const centroids = []; + const first = Math.floor(Math.random() * vectors.length); + centroids.push([...vectors[first].embedding]); + + for (let k = 1; k < K; k++) { + const dists = vectors.map(v => { + let minDist = Infinity; + for (const c of centroids) { + const sim = VectorStore.cosineSimilarity(v.embedding, c); + const dist = 1 - sim; + if (dist < minDist) minDist = dist; + } + return minDist; + }); + + const total = dists.reduce((s, d) => s + d, 0); + if (total === 0) { + centroids.push([...vectors[Math.floor(Math.random() * vectors.length)].embedding]); + continue; + } + + let r = Math.random() * total; + for (let i = 0; i < vectors.length; i++) { + r -= dists[i]; + if (r <= 0) { + centroids.push([...vectors[i].embedding]); + break; + } + } + } + return centroids; + } + + /** 归一化向量（in-place） */ + _normalize(vec) { + let norm = 0; + for (let i = 0; i < vec.length; i++) norm += vec[i] * vec[i]; + norm = Math.sqrt(norm); + if (norm > 0) for (let i = 0; i < vec.length; i++) vec[i] /= norm; + } + + // ═══════════════════════════════════════════ + // 工具 + // ═══════════════════════════════════════════ + static cosineSimilarity(a, b) { if (!a || !b || a.length !== b.length) return 0; let dot = 0, normA = 0, normB = 0;