eigen-db 4.3.0 → 5.0.0

package/src/lib/vector-db.ts

@@ -17,7 +17,7 @@ import { VectorCapacityExceededError } from "./errors";
 import { decodeLexicon, encodeLexicon } from "./lexicon";
 import { MemoryManager } from "./memory-manager";
 import type { ResultItem } from "./result-set";
-import { iterableResults, topKResults } from "./result-set";
+import { iterableResults, queryResults } from "./result-set";
 import { getSimdWasmBinary } from "./simd-binary";
 import type { StorageProvider } from "./storage";
 import { InMemoryStorageProvider } from "./storage";
@@ -38,7 +38,7 @@ const STREAM_CHUNK_SIZE = 65536;
 export class VectorDB {
   private readonly memoryManager: MemoryManager;
   private readonly storage: StorageProvider;
-  private readonly dimensions: number;
+  private readonly _dimensions: number;
   private readonly shouldNormalize: boolean;
   private wasmExports: WasmExports | null;
 
@@ -62,7 +62,7 @@ export class VectorDB {
   ) {
     this.memoryManager = memoryManager;
     this.storage = storage;
-    this.dimensions = dimensions;
+    this._dimensions = dimensions;
     this.shouldNormalize = shouldNormalize;
     this.wasmExports = wasmExports;
     this.keyToSlot = keyToSlot;
@@ -122,6 +122,81 @@ export class VectorDB {
     return this.keyToSlot.size;
   }
 
+  /** Number of dimensions per vector */
+  get dimensions(): number {
+    return this._dimensions;
+  }
+
+  /**
+   * Check whether a key exists in the database.
+   * Uses the internal key-to-slot map for O(1) lookup.
+   */
+  has(key: string): boolean {
+    this.assertOpen();
+    return this.keyToSlot.has(key);
+  }
+
+  /**
+   * Delete an entry by key. Returns true if the key existed, false otherwise.
+   * Uses swap-and-pop to avoid gaps in the underlying vector array.
+   */
+  delete(key: string): boolean {
+    this.assertOpen();
+
+    const slot = this.keyToSlot.get(key);
+    if (slot === undefined) return false;
+
+    const lastSlot = this.memoryManager.vectorCount - 1;
+
+    if (slot !== lastSlot) {
+      // Move last vector into the deleted slot
+      const lastVector = new Float32Array(this.memoryManager.readVector(lastSlot));
+      this.memoryManager.writeVector(slot, lastVector);
+
+      // Update mappings for the moved key
+      const movedKey = this.slotToKey[lastSlot];
+      this.keyToSlot.set(movedKey, slot);
+      this.slotToKey[slot] = movedKey;
+    }
+
+    // Remove the deleted key and shrink
+    this.keyToSlot.delete(key);
+    this.slotToKey.length = lastSlot;
+    this.memoryManager.setVectorCount(lastSlot);
+
+    return true;
+  }
+
+  /**
+   * Returns an iterable of all keys in the database.
+   */
+  keys(): IterableIterator<string> {
+    this.assertOpen();
+    return this.keyToSlot.keys();
+  }
+
+  /**
+   * Returns an iterable of [key, value] pairs.
+   * Values are returned as plain number array copies.
+   */
+  entries(): IterableIterator<[string, number[]]> {
+    this.assertOpen();
+    const keyToSlot = this.keyToSlot;
+    const mm = this.memoryManager;
+    return (function* () {
+      for (const [key, slot] of keyToSlot) {
+        yield [key, Array.from(mm.readVector(slot))] as [string, number[]];
+      }
+    })();
+  }
+
+  /**
+   * Implements the iterable protocol. Same as entries().
+   */
+  [Symbol.iterator](): IterableIterator<[string, number[]]> {
+    return this.entries();
+  }
+
   /**
    * Set a key-value pair. If the key already exists, its vector is overwritten (last-write-wins).
    * The value is a number[] or Float32Array of length equal to the configured dimensions.
@@ -129,8 +204,8 @@ export class VectorDB {
   set(key: string, value: VectorInput, options?: SetOptions): void {
     this.assertOpen();
 
-    if (value.length !== this.dimensions) {
-      throw new Error(`Vector dimension mismatch: expected ${this.dimensions}, got ${value.length}`);
+    if (value.length !== this._dimensions) {
+      throw new Error(`Vector dimension mismatch: expected ${this._dimensions}, got ${value.length}`);
     }
 
     // Convert to Float32Array (also clones to avoid mutating caller's array)
@@ -206,16 +281,18 @@ export class VectorDB {
   query(value: VectorInput, options?: QueryOptions): ResultItem[] | Iterable<ResultItem> {
     this.assertOpen();
 
-    const k = options?.topK ?? Infinity;
+    const limit = options?.limit ?? Infinity;
+    const order = options?.order ?? "descend";
     const minSimilarity = options?.minSimilarity;
+    const maxSimilarity = options?.maxSimilarity;
     const iterable = options && "iterable" in options && options.iterable;
 
     if (this.size === 0) {
       return [];
     }
 
-    if (value.length !== this.dimensions) {
-      throw new Error(`Query vector dimension mismatch: expected ${this.dimensions}, got ${value.length}`);
+    if (value.length !== this._dimensions) {
+      throw new Error(`Query vector dimension mismatch: expected ${this._dimensions}, got ${value.length}`);
     }
 
     // Convert to Float32Array and optionally normalize the query vector
@@ -242,21 +319,21 @@ export class VectorDB {
         this.memoryManager.dbOffset,
         scoresOffset,
         totalVectors,
-        this.dimensions,
+        this._dimensions,
       );
     } else {
       const queryView = new Float32Array(
         this.memoryManager.memory.buffer,
         this.memoryManager.queryOffset,
-        this.dimensions,
+        this._dimensions,
       );
       const dbView = new Float32Array(
         this.memoryManager.memory.buffer,
         this.memoryManager.dbOffset,
-        totalVectors * this.dimensions,
+        totalVectors * this._dimensions,
       );
       const scoresView = new Float32Array(this.memoryManager.memory.buffer, scoresOffset, totalVectors);
-      searchAll(queryView, dbView, scoresView, totalVectors, this.dimensions);
+      searchAll(queryView, dbView, scoresView, totalVectors, this._dimensions);
     }
 
     // Read scores (make a copy so the buffer can be reused)
@@ -268,10 +345,12 @@ export class VectorDB {
       return slotToKey[slotIndex];
     };
 
+    const resultOptions = { limit, order, minSimilarity, maxSimilarity } as const;
+
     if (iterable) {
-      return iterableResults(scores, resolveKey, k, minSimilarity);
+      return iterableResults(scores, resolveKey, resultOptions);
     }
-    return topKResults(scores, resolveKey, k, minSimilarity);
+    return queryResults(scores, resolveKey, resultOptions);
   }
 
   /**
@@ -283,12 +362,12 @@ export class VectorDB {
     const totalVectors = this.memoryManager.vectorCount;
 
     // Serialize vectors from WASM memory
-    const vectorBytes = new Uint8Array(totalVectors * this.dimensions * 4);
+    const vectorBytes = new Uint8Array(totalVectors * this._dimensions * 4);
     if (totalVectors > 0) {
       const src = new Uint8Array(
         this.memoryManager.memory.buffer,
         this.memoryManager.dbOffset,
-        totalVectors * this.dimensions * 4,
+        totalVectors * this._dimensions * 4,
       );
       vectorBytes.set(src);
     }
@@ -335,7 +414,7 @@ export class VectorDB {
     this.assertOpen();
 
     const totalVectors = this.memoryManager.vectorCount;
-    const vectorDataLen = totalVectors * this.dimensions * 4;
+    const vectorDataLen = totalVectors * this._dimensions * 4;
 
     // Encode keys (typically much smaller than vectors)
     const keysBytes = encodeLexicon(this.slotToKey);
@@ -346,7 +425,7 @@ export class VectorDB {
     const headerView = new DataView(header);
     headerView.setUint32(0, EXPORT_MAGIC, true);
     headerView.setUint32(4, EXPORT_VERSION, true);
-    headerView.setUint32(8, this.dimensions, true);
+    headerView.setUint32(8, this._dimensions, true);
     headerView.setUint32(12, totalVectors, true);
     headerView.setUint32(16, vectorDataLen, true);
     headerView.setUint32(20, keysDataLen, true);
@@ -415,8 +494,8 @@ export class VectorDB {
     }
 
     const dimensions = headerView.getUint32(8, true);
-    if (dimensions !== this.dimensions) {
-      throw new Error(`Import dimension mismatch: expected ${this.dimensions}, got ${dimensions}`);
+    if (dimensions !== this._dimensions) {
+      throw new Error(`Import dimension mismatch: expected ${this._dimensions}, got ${dimensions}`);
     }
 
     const vectorCount = headerView.getUint32(12, true);