verso-db 0.1.1

package/src/Collection.ts

@@ -0,0 +1,695 @@
+import { HNSWIndex, DistanceMetric } from './HNSWIndex';
+import { mkdir } from 'fs/promises';
+import * as path from 'path';
+
+/**
+ * Configuration for adding vectors to a collection.
+ */
+export interface AddConfig {
+  /** Unique string identifiers for each vector */
+  ids: string[];
+  /** Array of vectors to add (must match collection dimension) */
+  vectors: number[][];
+  /** Optional metadata for each vector (same length as ids/vectors) */
+  metadata?: Array<Record<string, any>>;
+}
+
+/**
+ * Configuration for querying a collection.
+ *
+ * @example
+ * ```typescript
+ * // Simple query
+ * const results = await collection.query({
+ *   queryVector: [0.1, 0.2, 0.3],
+ *   k: 10
+ * });
+ *
+ * // Query with metadata filter
+ * const results = await collection.query({
+ *   queryVector: [0.1, 0.2, 0.3],
+ *   k: 10,
+ *   filter: {
+ *     category: 'science',           // Exact match
+ *     year: { $gte: 2020 },          // Greater than or equal
+ *     status: { $in: ['active', 'pending'] }  // In array
+ *   }
+ * });
+ * ```
+ */
+export interface QueryConfig {
+  /** The query vector (must match collection dimension) */
+  queryVector: number[];
+  /** Number of nearest neighbors to return */
+  k: number;
+  /**
+   * Optional metadata filter. Supports MongoDB-style operators:
+   * - Simple equality: `{ field: value }`
+   * - `$gt`: Greater than `{ field: { $gt: 5 } }`
+   * - `$gte`: Greater than or equal `{ field: { $gte: 5 } }`
+   * - `$lt`: Less than `{ field: { $lt: 10 } }`
+   * - `$lte`: Less than or equal `{ field: { $lte: 10 } }`
+   * - `$ne`: Not equal `{ field: { $ne: 'excluded' } }`
+   * - `$in`: In array `{ field: { $in: ['a', 'b', 'c'] } }`
+   * - `$nin`: Not in array `{ field: { $nin: ['x', 'y'] } }`
+   */
+  filter?: Record<string, any>;
+  /** Search effort parameter (higher = better recall, slower). Default: max(k*2, 50) */
+  efSearch?: number;
+}
+
+export interface QueryResult {
+  ids: string[];
+  distances: number[];
+  metadata: Array<Record<string, any>>;
+}
+
+export class Collection {
+  private name: string;
+  private dimension: number;
+  private metric: DistanceMetric;
+  private M: number; // HNSW M parameter
+  private efConstruction: number; // HNSW efConstruction parameter
+  private indexPath: string;
+  private metaPath: string;
+  private deletedPath: string; // Path for persisting deleted IDs
+  private hnsw: HNSWIndex;
+  private idMap: Map<string, number>; // Maps string IDs to internal numeric IDs
+  private idReverseMap: Map<number, string>; // Maps internal numeric IDs to string IDs
+  private metadata: Map<number, Record<string, any>>; // Stores metadata for each vector
+  private deletedIds: Set<number>; // Tombstone markers for deleted vectors
+
+  constructor(name: string, config: { dimension: number; metric?: DistanceMetric; M?: number; efConstruction?: number }, collectionPath: string) {
+    this.name = name;
+    this.dimension = config.dimension;
+    this.metric = config.metric || 'cosine';
+    this.M = config.M || 16;
+    this.efConstruction = config.efConstruction || 200;
+    this.indexPath = path.join(collectionPath, `${name}.hnsw`);
+    this.metaPath = path.join(collectionPath, `${name}.meta`);
+    this.deletedPath = path.join(collectionPath, `${name}.deleted`);
+
+    // Initialize HNSW index with the specified parameters
+    this.hnsw = new HNSWIndex(
+      config.dimension,
+      this.metric,
+      this.M,
+      this.efConstruction
+    );
+
+    this.idMap = new Map();
+    this.idReverseMap = new Map();
+    this.metadata = new Map();
+    this.deletedIds = new Set();
+  }
+
+  async init(): Promise<void> {
+    // Load existing data if files exist
+    await this.loadFromDisk();
+  }
+
+  private async loadFromDisk(): Promise<void> {
+    // Check if index files exist and load the binary HNSW index using Bun APIs
+    const indexFile = Bun.file(this.indexPath);
+    if (await indexFile.exists()) {
+      try {
+        this.hnsw = await HNSWIndex.loadFromFile(this.indexPath);
+      } catch (e) {
+        console.warn(`Failed to load HNSW index from ${this.indexPath}:`, e);
+      }
+    }
+
+    // Load metadata from file using Bun APIs
+    const metaFile = Bun.file(this.metaPath);
+    if (await metaFile.exists()) {
+      try {
+        const metaContent = await metaFile.text();
+        // Process lines without creating intermediate filter array
+        const lines = metaContent.split('\n');
+        for (let i = 0; i < lines.length; i++) {
+          const line = lines[i];
+          if (line.length === 0 || line.trim().length === 0) continue;
+
+          try {
+            const parts = line.split('\t');
+            const id = parts[0];
+            const internalId = parseInt(parts[1], 10);
+            if (!isNaN(internalId)) {
+              this.idMap.set(id, internalId);
+              this.idReverseMap.set(internalId, id);
+
+              if (parts.length > 2) {
+                // Join remaining parts for metadata (avoiding spread operator)
+                const metaStr = parts.slice(2).join('\t');
+                const metadata = JSON.parse(metaStr);
+                this.metadata.set(internalId, metadata);
+              }
+            }
+          } catch (e) {
+            // Skip malformed lines
+            console.warn(`Skipping malformed line in metadata file: ${line}`);
+          }
+        }
+      } catch (e) {
+        // Ignore file read errors
+      }
+    }
+
+    // Load deleted IDs from file
+    const deletedFile = Bun.file(this.deletedPath);
+    if (await deletedFile.exists()) {
+      try {
+        const deletedContent = await deletedFile.text();
+        const deletedArray = JSON.parse(deletedContent);
+        if (Array.isArray(deletedArray)) {
+          this.deletedIds = new Set(deletedArray);
+        }
+      } catch (e) {
+        // Ignore file read errors, start with empty set
+      }
+    }
+  }
+
+  async add(config: AddConfig): Promise<void> {
+    if (config.vectors.length !== config.ids.length) {
+      throw new Error('Number of vectors must match number of IDs');
+    }
+
+    if (config.metadata && config.metadata.length !== config.ids.length) {
+      throw new Error('Number of metadata entries must match number of IDs');
+    }
+
+    // Check for duplicate IDs before adding any vectors
+    const duplicates: string[] = [];
+    for (const id of config.ids) {
+      if (this.idMap.has(id)) {
+        duplicates.push(id);
+      }
+    }
+    if (duplicates.length > 0) {
+      throw new Error(
+        `Cannot add: IDs already exist: ${duplicates.slice(0, 5).join(', ')}${duplicates.length > 5 ? ` and ${duplicates.length - 5} more` : ''}. ` +
+        `Use upsert() to update existing vectors.`
+      );
+    }
+
+    // Prepare bulk insert data - single pass validation and transformation
+    const points: Array<{ id: number; vector: Float32Array }> = new Array(config.ids.length);
+    const startId = this.idMap.size;
+
+    for (let i = 0; i < config.ids.length; i++) {
+      const id = config.ids[i];
+      const vector = config.vectors[i];
+      const metadata = config.metadata ? config.metadata[i] : undefined;
+
+      // Validate vector dimension
+      if (vector.length !== this.dimension) {
+        throw new Error(`Vector at index ${i} has dimension ${vector.length}, expected ${this.dimension}`);
+      }
+
+      // Generate a numeric ID for internal use
+      const numericId = startId + i;
+
+      // Update mappings
+      this.idMap.set(id, numericId);
+      this.idReverseMap.set(numericId, id);
+
+      // Store metadata if provided
+      if (metadata) {
+        this.metadata.set(numericId, metadata);
+      }
+
+      // Prepare point for bulk insert
+      points[i] = { id: numericId, vector: new Float32Array(vector) };
+    }
+
+    // Use bulk insert for better performance
+    await this.hnsw.addPointsBulk(points);
+  }
+
+  async query(config: QueryConfig): Promise<QueryResult> {
+    const { queryVector, k, filter, efSearch } = config;
+
+    if (queryVector.length !== this.dimension) {
+      throw new Error(`Query vector has dimension ${queryVector.length}, expected ${this.dimension}`);
+    }
+
+    // Perform the search
+    const ef = efSearch || Math.max(k * 2, 50); // Default ef to 2*k or 50
+    const results = this.hnsw.searchKNN(new Float32Array(queryVector), k * 2, ef); // Request more results to handle duplicates
+
+    // Apply filtering if provided (check for non-empty filter without creating array)
+    let filteredResults = results;
+    if (filter) {
+      let hasFilter = false;
+      for (const _ in filter) { hasFilter = true; break; }
+      if (hasFilter) {
+        filteredResults = results.filter(result => {
+          const metadata = this.metadata.get(result.id) || {};
+          return this.matchesFilter(metadata, filter);
+        });
+      }
+    }
+
+    // Deduplicate results and filter out deleted vectors
+    const seenIds = new Set<number>();
+    const deduplicatedResults: Array<{ id: number; distance: number }> = [];
+
+    for (const result of filteredResults) {
+      // Skip deleted vectors (tombstone check)
+      if (this.deletedIds.has(result.id)) continue;
+
+      if (!seenIds.has(result.id)) {
+        seenIds.add(result.id);
+        deduplicatedResults.push(result);
+      }
+    }
+
+    // Sort by distance to ensure proper ordering (closest first)
+    deduplicatedResults.sort((a, b) => a.distance - b.distance);
+
+    // Limit to k results
+    const resultCount = Math.min(deduplicatedResults.length, k);
+
+    // Pre-allocate output arrays for better performance
+    const ids = new Array<string>(resultCount);
+    const distances = new Array<number>(resultCount);
+    const metadata = new Array<Record<string, any>>(resultCount);
+
+    let outIdx = 0;
+    for (let i = 0; i < resultCount; i++) {
+      const result = deduplicatedResults[i];
+      const id = this.idReverseMap.get(result.id);
+      if (id) {
+        ids[outIdx] = id;
+        distances[outIdx] = result.distance;
+        metadata[outIdx] = this.metadata.get(result.id) || {};
+        outIdx++;
+      }
+    }
+
+    // Trim arrays if any results were filtered (shouldn't happen often)
+    if (outIdx < resultCount) {
+      ids.length = outIdx;
+      distances.length = outIdx;
+      metadata.length = outIdx;
+    }
+
+    return { ids, distances, metadata };
+  }
+
+  /**
+   * Batch query for multiple vectors at once.
+   * More efficient than calling query() multiple times.
+   *
+   * @param configs Array of query configurations
+   * @returns Array of query results, one per query
+   */
+  async queryBatch(configs: QueryConfig[]): Promise<QueryResult[]> {
+    if (configs.length === 0) return [];
+
+    // Validate all queries have correct dimension
+    for (let i = 0; i < configs.length; i++) {
+      if (configs[i].queryVector.length !== this.dimension) {
+        throw new Error(`Query ${i} has dimension ${configs[i].queryVector.length}, expected ${this.dimension}`);
+      }
+    }
+
+    // Use batch search on the HNSW index - single loop instead of 3 map calls
+    const queries = new Array<Float32Array>(configs.length);
+    let k = 0;
+    let efSearch = 0;
+    for (let i = 0; i < configs.length; i++) {
+      const c = configs[i];
+      queries[i] = new Float32Array(c.queryVector);
+      if (c.k > k) k = c.k;
+      const ef = c.efSearch || Math.max(c.k * 2, 50);
+      if (ef > efSearch) efSearch = ef;
+    }
+
+    const batchResults = this.hnsw.searchKNNBatch(queries, k * 2, efSearch);
+
+    // Process results for each query
+    const results: QueryResult[] = [];
+
+    for (let q = 0; q < configs.length; q++) {
+      const config = configs[q];
+      const rawResults = batchResults[q];
+
+      // Apply filtering if provided (check for non-empty filter without creating array)
+      let filteredResults = rawResults;
+      if (config.filter) {
+        let hasFilter = false;
+        for (const _ in config.filter) { hasFilter = true; break; }
+        if (hasFilter) {
+          filteredResults = rawResults.filter(result => {
+            const metadata = this.metadata.get(result.id) || {};
+            return this.matchesFilter(metadata, config.filter!);
+          });
+        }
+      }
+
+      // Deduplicate and filter out deleted vectors
+      const seenIds = new Set<number>();
+      const deduplicatedResults: Array<{ id: number; distance: number }> = [];
+      for (const result of filteredResults) {
+        // Skip deleted vectors (tombstone check)
+        if (this.deletedIds.has(result.id)) continue;
+
+        if (!seenIds.has(result.id)) {
+          seenIds.add(result.id);
+          deduplicatedResults.push(result);
+        }
+      }
+
+      // Sort and limit
+      deduplicatedResults.sort((a, b) => a.distance - b.distance);
+      const resultCount = Math.min(deduplicatedResults.length, config.k);
+
+      // Pre-allocate output arrays
+      const ids = new Array<string>(resultCount);
+      const distances = new Array<number>(resultCount);
+      const metadata = new Array<Record<string, any>>(resultCount);
+
+      let outIdx = 0;
+      for (let i = 0; i < resultCount; i++) {
+        const result = deduplicatedResults[i];
+        const id = this.idReverseMap.get(result.id);
+        if (id) {
+          ids[outIdx] = id;
+          distances[outIdx] = result.distance;
+          metadata[outIdx] = this.metadata.get(result.id) || {};
+          outIdx++;
+        }
+      }
+
+      // Trim if needed
+      if (outIdx < resultCount) {
+        ids.length = outIdx;
+        distances.length = outIdx;
+        metadata.length = outIdx;
+      }
+
+      results.push({ ids, distances, metadata });
+    }
+
+    return results;
+  }
+
+  /**
+   * Brute-force KNN search for validation and correctness checking
+   * This checks all vectors and returns the true k nearest neighbors
+   */
+  async queryBruteForce(config: QueryConfig): Promise<QueryResult> {
+    const { queryVector, k, filter } = config;
+
+    if (queryVector.length !== this.dimension) {
+      throw new Error(`Query vector has dimension ${queryVector.length}, expected ${this.dimension}`);
+    }
+
+    // Get all vectors from the HNSW index
+    const allNodes: Array<{ id: number; vector: Float32Array }> = [];
+    for (const [numericId, vector] of this.hnsw.getAllVectors()) {
+      allNodes.push({ id: numericId, vector });
+    }
+
+    // Compute distances to all vectors
+    const distances: Array<{ id: number; distance: number }> = [];
+    for (const node of allNodes) {
+      const distance = this.hnsw.calculateDistance(new Float32Array(queryVector), node.vector);
+      distances.push({ id: node.id, distance });
+    }
+
+    // Sort by distance and get top k
+    distances.sort((a, b) => a.distance - b.distance);
+    const topK = distances.slice(0, k);
+
+    // Apply filtering if provided
+    let filteredResults = topK;
+    if (filter && Object.keys(filter).length > 0) {
+      filteredResults = [];
+      for (const result of topK) {
+        const metadata = this.metadata.get(result.id) || {};
+        if (this.matchesFilter(metadata, filter)) {
+          filteredResults.push(result);
+        }
+        if (filteredResults.length >= k) break;
+      }
+    }
+
+    // Prepare the output
+    const ids: string[] = [];
+    const distancesOut: number[] = [];
+    const metadata: Array<Record<string, any>> = [];
+
+    for (const result of filteredResults) {
+      const id = this.idReverseMap.get(result.id);
+      if (id) {
+        ids.push(id);
+        distancesOut.push(result.distance);
+        metadata.push(this.metadata.get(result.id) || {});
+      }
+    }
+
+    return {
+      ids,
+      distances: distancesOut,
+      metadata
+    };
+  }
+
+  /**
+   * Upsert vectors (insert or update).
+   *
+   * NOTE: HNSW indices don't support efficient in-place updates.
+   * Duplicate IDs will throw an error. For updates, rebuild the index
+   * without the old vectors and add the new ones.
+   *
+   * @throws Error if any ID already exists
+   */
+  async upsert(config: AddConfig): Promise<void> {
+    // Check for duplicates and throw clear error
+    const duplicates: string[] = [];
+    for (const id of config.ids) {
+      if (this.idMap.has(id)) {
+        duplicates.push(id);
+      }
+    }
+
+    if (duplicates.length > 0) {
+      throw new Error(
+        `Cannot upsert: IDs already exist: ${duplicates.join(', ')}. ` +
+        `HNSW indices don't support efficient updates. ` +
+        `To update vectors, create a new index without the old vectors.`
+      );
+    }
+
+    // No duplicates - proceed with add
+    await this.add(config);
+  }
+
+  /**
+   * Returns the number of active (non-deleted) vectors in the collection.
+   */
+  count(): number {
+    return this.idMap.size - this.deletedIds.size;
+  }
+
+  /**
+   * Returns the total number of vectors including deleted (tombstoned) ones.
+   * Use this to determine when compaction might be beneficial.
+   */
+  countWithDeleted(): number {
+    return this.idMap.size;
+  }
+
+  /**
+   * Returns the number of deleted (tombstoned) vectors.
+   */
+  deletedCount(): number {
+    return this.deletedIds.size;
+  }
+
+  /**
+   * Mark a vector as deleted (tombstone deletion).
+   * The vector remains in the index but is excluded from search results.
+   * Use compact() to permanently remove deleted vectors and reclaim space.
+   *
+   * @param id The string ID of the vector to delete
+   * @returns true if the vector was deleted, false if it didn't exist or was already deleted
+   */
+  delete(id: string): boolean {
+    const numericId = this.idMap.get(id);
+    if (numericId === undefined) return false;
+    if (this.deletedIds.has(numericId)) return false;
+
+    this.deletedIds.add(numericId);
+    return true;
+  }
+
+  /**
+   * Mark multiple vectors as deleted (tombstone deletion).
+   *
+   * @param ids Array of string IDs to delete
+   * @returns Number of vectors that were successfully deleted
+   */
+  deleteBatch(ids: string[]): number {
+    let deleted = 0;
+    for (const id of ids) {
+      if (this.delete(id)) deleted++;
+    }
+    return deleted;
+  }
+
+  /**
+   * Check if a vector exists and is not deleted.
+   */
+  has(id: string): boolean {
+    const numericId = this.idMap.get(id);
+    if (numericId === undefined) return false;
+    return !this.deletedIds.has(numericId);
+  }
+
+  /**
+   * Check if a vector was deleted (tombstoned).
+   */
+  isDeleted(id: string): boolean {
+    const numericId = this.idMap.get(id);
+    if (numericId === undefined) return false;
+    return this.deletedIds.has(numericId);
+  }
+
+  async saveToDisk(): Promise<void> {
+    // Ensure the directory exists before saving using Bun-compatible async API
+    const dirPath = path.dirname(this.indexPath);
+    await mkdir(dirPath, { recursive: true }).catch(() => {});
+
+    // Save the HNSW index to binary file (now async)
+    await this.hnsw.saveToFile(this.indexPath);
+
+    // Save the metadata to disk using Bun APIs
+    const metaLines: string[] = [];
+    for (const [numericId, id] of this.idReverseMap) {
+      const meta = this.metadata.get(numericId);
+      const metaStr = meta ? JSON.stringify(meta) : '{}';
+      metaLines.push(`${id}\t${numericId}\t${metaStr}`);
+    }
+    await Bun.write(this.metaPath, metaLines.join('\n'));
+
+    // Save deleted IDs to disk (only if there are any)
+    if (this.deletedIds.size > 0) {
+      await Bun.write(this.deletedPath, JSON.stringify([...this.deletedIds]));
+    } else {
+      // Remove the deleted file if no deletions (clean state)
+      const deletedFile = Bun.file(this.deletedPath);
+      if (await deletedFile.exists()) {
+        await Bun.write(this.deletedPath, '[]');
+      }
+    }
+  }
+
+  private matchesFilter(metadata: Record<string, any>, filter: Record<string, any>): boolean {
+    // Use for-in loop to avoid Object.entries() array allocation
+    for (const key in filter) {
+      const value = filter[key];
+      const metaValue = metadata[key];
+      if (typeof value === 'object' && value !== null) {
+        // Handle complex filters like { $gt: 25 }
+        if (value.$gt !== undefined && metaValue <= value.$gt) return false;
+        if (value.$lt !== undefined && metaValue >= value.$lt) return false;
+        if (value.$gte !== undefined && metaValue < value.$gte) return false;
+        if (value.$lte !== undefined && metaValue > value.$lte) return false;
+        if (value.$ne !== undefined && metaValue === value.$ne) return false;
+        if (value.$in !== undefined && !value.$in.includes(metaValue)) return false;
+        if (value.$nin !== undefined && value.$nin.includes(metaValue)) return false;
+      } else {
+        // Handle simple equality check
+        if (metaValue !== value) return false;
+      }
+    }
+    return true;
+  }
+
+  /**
+   * Compact the collection by rebuilding the index without deleted vectors.
+   * This permanently removes tombstoned vectors and reclaims space.
+   *
+   * @returns Number of vectors removed during compaction
+   */
+  async compact(): Promise<number> {
+    if (this.deletedIds.size === 0) return 0;
+
+    const removedCount = this.deletedIds.size;
+
+    // Collect all non-deleted vectors
+    const activeVectors: Array<{ id: string; numericId: number; vector: Float32Array; meta?: Record<string, any> }> = [];
+
+    for (const [numericId, vector] of this.hnsw.getAllVectors()) {
+      if (!this.deletedIds.has(numericId)) {
+        const stringId = this.idReverseMap.get(numericId);
+        if (stringId) {
+          activeVectors.push({
+            id: stringId,
+            numericId,
+            vector,
+            meta: this.metadata.get(numericId)
+          });
+        }
+      }
+    }
+
+    // Destroy old index
+    this.hnsw.destroy();
+
+    // Create new index with same parameters as the original
+    this.hnsw = new HNSWIndex(
+      this.dimension,
+      this.metric,
+      this.M,
+      this.efConstruction
+    );
+
+    // Clear all mappings
+    this.idMap.clear();
+    this.idReverseMap.clear();
+    this.metadata.clear();
+    this.deletedIds.clear();
+
+    // Prepare bulk insert data and re-populate mappings
+    const points: Array<{ id: number; vector: Float32Array }> = new Array(activeVectors.length);
+
+    for (let i = 0; i < activeVectors.length; i++) {
+      const { id, vector, meta } = activeVectors[i];
+      const newNumericId = i;
+
+      this.idMap.set(id, newNumericId);
+      this.idReverseMap.set(newNumericId, id);
+      if (meta) {
+        this.metadata.set(newNumericId, meta);
+      }
+
+      points[i] = { id: newNumericId, vector };
+    }
+
+    // Use bulk insert for better performance
+    await this.hnsw.addPointsBulk(points);
+
+    return removedCount;
+  }
+
+  async destroy(): Promise<void> {
+    // Save to disk before destroying
+    await this.saveToDisk();
+
+    // Destroy the HNSW index to free memory
+    if (this.hnsw && typeof this.hnsw.destroy === 'function') {
+      this.hnsw.destroy();
+    }
+
+    // Clear all maps to free memory
+    this.idMap.clear();
+    this.idReverseMap.clear();
+    this.metadata.clear();
+    this.deletedIds.clear();
+  }
+}