perf: IVF 倒排索引优化向量搜索 - K-Means 聚类 + nprobe 近邻检索

- N < 200 自动降级暴力搜索 - N ≥ 200 K-Means++ 聚类构建 IVF 索引 - nprobe=5 探测最近 5 个聚类，搜索量降 75% - 索引持久化到 IndexedDB，数据变化自动重建 - 内存向量缓存 + Map O(1) 查找
2026-04-05 01:47:57 +08:00
parent 4c52650881
commit e97e2e4078
1 changed files with 312 additions and 27 deletions
@@ -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<vecId, vectorData> }
+        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;