@push.rocks/smartmongo 2.2.0 → 4.0.0

package/ts/tsmdb/engine/IndexEngine.ts

@@ -0,0 +1,798 @@
+import * as plugins from '../tsmdb.plugins.js';
+import type { IStorageAdapter } from '../storage/IStorageAdapter.js';
+
+// Simple B-Tree implementation for range queries
+// Since sorted-btree has ESM/CJS interop issues, we use a simple custom implementation
+class SimpleBTree<K, V> {
+  private entries: Map<string, { key: K; value: V }> = new Map();
+  private sortedKeys: K[] = [];
+  private comparator: (a: K, b: K) => number;
+
+  constructor(_unused?: undefined, comparator?: (a: K, b: K) => number) {
+    this.comparator = comparator || ((a: K, b: K) => {
+      if (a < b) return -1;
+      if (a > b) return 1;
+      return 0;
+    });
+  }
+
+  private keyToString(key: K): string {
+    return JSON.stringify(key);
+  }
+
+  set(key: K, value: V): boolean {
+    const keyStr = this.keyToString(key);
+    const existed = this.entries.has(keyStr);
+    this.entries.set(keyStr, { key, value });
+
+    if (!existed) {
+      // Insert in sorted order
+      const idx = this.sortedKeys.findIndex(k => this.comparator(k, key) > 0);
+      if (idx === -1) {
+        this.sortedKeys.push(key);
+      } else {
+        this.sortedKeys.splice(idx, 0, key);
+      }
+    }
+    return !existed;
+  }
+
+  get(key: K): V | undefined {
+    const entry = this.entries.get(this.keyToString(key));
+    return entry?.value;
+  }
+
+  delete(key: K): boolean {
+    const keyStr = this.keyToString(key);
+    if (this.entries.has(keyStr)) {
+      this.entries.delete(keyStr);
+      const idx = this.sortedKeys.findIndex(k => this.comparator(k, key) === 0);
+      if (idx !== -1) {
+        this.sortedKeys.splice(idx, 1);
+      }
+      return true;
+    }
+    return false;
+  }
+
+  forRange(
+    lowKey: K | undefined,
+    highKey: K | undefined,
+    lowInclusive: boolean,
+    highInclusive: boolean,
+    callback: (value: V, key: K) => void
+  ): void {
+    for (const key of this.sortedKeys) {
+      // Check low bound
+      if (lowKey !== undefined) {
+        const cmp = this.comparator(key, lowKey);
+        if (cmp < 0) continue;
+        if (cmp === 0 && !lowInclusive) continue;
+      }
+
+      // Check high bound
+      if (highKey !== undefined) {
+        const cmp = this.comparator(key, highKey);
+        if (cmp > 0) break;
+        if (cmp === 0 && !highInclusive) break;
+      }
+
+      const entry = this.entries.get(this.keyToString(key));
+      if (entry) {
+        callback(entry.value, key);
+      }
+    }
+  }
+}
+import type {
+  Document,
+  IStoredDocument,
+  IIndexSpecification,
+  IIndexInfo,
+  ICreateIndexOptions,
+} from '../types/interfaces.js';
+import { TsmdbDuplicateKeyError, TsmdbIndexError } from '../errors/TsmdbErrors.js';
+import { QueryEngine } from './QueryEngine.js';
+
+/**
+ * Comparator for B-Tree that handles mixed types consistently
+ */
+function indexKeyComparator(a: any, b: any): number {
+  // Handle null/undefined
+  if (a === null || a === undefined) {
+    if (b === null || b === undefined) return 0;
+    return -1;
+  }
+  if (b === null || b === undefined) return 1;
+
+  // Handle arrays (compound keys)
+  if (Array.isArray(a) && Array.isArray(b)) {
+    for (let i = 0; i < Math.max(a.length, b.length); i++) {
+      const cmp = indexKeyComparator(a[i], b[i]);
+      if (cmp !== 0) return cmp;
+    }
+    return 0;
+  }
+
+  // Handle ObjectId
+  if (a instanceof plugins.bson.ObjectId && b instanceof plugins.bson.ObjectId) {
+    return a.toHexString().localeCompare(b.toHexString());
+  }
+
+  // Handle Date
+  if (a instanceof Date && b instanceof Date) {
+    return a.getTime() - b.getTime();
+  }
+
+  // Handle different types - use type ordering (null < number < string < object)
+  const typeOrder = (v: any): number => {
+    if (v === null || v === undefined) return 0;
+    if (typeof v === 'number') return 1;
+    if (typeof v === 'string') return 2;
+    if (typeof v === 'boolean') return 3;
+    if (v instanceof Date) return 4;
+    if (v instanceof plugins.bson.ObjectId) return 5;
+    return 6;
+  };
+
+  const typeA = typeOrder(a);
+  const typeB = typeOrder(b);
+  if (typeA !== typeB) return typeA - typeB;
+
+  // Same type comparison
+  if (typeof a === 'number') return a - b;
+  if (typeof a === 'string') return a.localeCompare(b);
+  if (typeof a === 'boolean') return (a ? 1 : 0) - (b ? 1 : 0);
+
+  // Fallback to string comparison
+  return String(a).localeCompare(String(b));
+}
+
+/**
+ * Index data structure using B-Tree for range queries
+ */
+interface IIndexData {
+  name: string;
+  key: Record<string, 1 | -1 | string>;
+  unique: boolean;
+  sparse: boolean;
+  expireAfterSeconds?: number;
+  // B-Tree for ordered index lookups (supports range queries)
+  btree: SimpleBTree<any, Set<string>>;
+  // Hash map for fast equality lookups
+  hashMap: Map<string, Set<string>>;
+}
+
+/**
+ * Index engine for managing indexes and query optimization
+ */
+export class IndexEngine {
+  private dbName: string;
+  private collName: string;
+  private storage: IStorageAdapter;
+  private indexes: Map<string, IIndexData> = new Map();
+  private initialized = false;
+
+  constructor(dbName: string, collName: string, storage: IStorageAdapter) {
+    this.dbName = dbName;
+    this.collName = collName;
+    this.storage = storage;
+  }
+
+  /**
+   * Initialize indexes from storage
+   */
+  async initialize(): Promise<void> {
+    if (this.initialized) return;
+
+    const storedIndexes = await this.storage.getIndexes(this.dbName, this.collName);
+    const documents = await this.storage.findAll(this.dbName, this.collName);
+
+    for (const indexSpec of storedIndexes) {
+      const indexData: IIndexData = {
+        name: indexSpec.name,
+        key: indexSpec.key,
+        unique: indexSpec.unique || false,
+        sparse: indexSpec.sparse || false,
+        expireAfterSeconds: indexSpec.expireAfterSeconds,
+        btree: new SimpleBTree<any, Set<string>>(undefined, indexKeyComparator),
+        hashMap: new Map(),
+      };
+
+      // Build index entries
+      for (const doc of documents) {
+        const keyValue = this.extractKeyValue(doc, indexSpec.key);
+        if (keyValue !== null || !indexData.sparse) {
+          const keyStr = JSON.stringify(keyValue);
+
+          // Add to hash map
+          if (!indexData.hashMap.has(keyStr)) {
+            indexData.hashMap.set(keyStr, new Set());
+          }
+          indexData.hashMap.get(keyStr)!.add(doc._id.toHexString());
+
+          // Add to B-tree
+          const existing = indexData.btree.get(keyValue);
+          if (existing) {
+            existing.add(doc._id.toHexString());
+          } else {
+            indexData.btree.set(keyValue, new Set([doc._id.toHexString()]));
+          }
+        }
+      }
+
+      this.indexes.set(indexSpec.name, indexData);
+    }
+
+    this.initialized = true;
+  }
+
+  /**
+   * Create a new index
+   */
+  async createIndex(
+    key: Record<string, 1 | -1 | 'text' | '2dsphere'>,
+    options?: ICreateIndexOptions
+  ): Promise<string> {
+    await this.initialize();
+
+    // Generate index name if not provided
+    const name = options?.name || this.generateIndexName(key);
+
+    // Check if index already exists
+    if (this.indexes.has(name)) {
+      return name;
+    }
+
+    // Create index data structure
+    const indexData: IIndexData = {
+      name,
+      key: key as Record<string, 1 | -1 | string>,
+      unique: options?.unique || false,
+      sparse: options?.sparse || false,
+      expireAfterSeconds: options?.expireAfterSeconds,
+      btree: new SimpleBTree<any, Set<string>>(undefined, indexKeyComparator),
+      hashMap: new Map(),
+    };
+
+    // Build index from existing documents
+    const documents = await this.storage.findAll(this.dbName, this.collName);
+
+    for (const doc of documents) {
+      const keyValue = this.extractKeyValue(doc, key);
+
+      if (keyValue === null && indexData.sparse) {
+        continue;
+      }
+
+      const keyStr = JSON.stringify(keyValue);
+
+      if (indexData.unique && indexData.hashMap.has(keyStr)) {
+        throw new TsmdbDuplicateKeyError(
+          `E11000 duplicate key error index: ${this.dbName}.${this.collName}.$${name}`,
+          key as Record<string, 1>,
+          keyValue
+        );
+      }
+
+      // Add to hash map
+      if (!indexData.hashMap.has(keyStr)) {
+        indexData.hashMap.set(keyStr, new Set());
+      }
+      indexData.hashMap.get(keyStr)!.add(doc._id.toHexString());
+
+      // Add to B-tree
+      const existing = indexData.btree.get(keyValue);
+      if (existing) {
+        existing.add(doc._id.toHexString());
+      } else {
+        indexData.btree.set(keyValue, new Set([doc._id.toHexString()]));
+      }
+    }
+
+    // Store index
+    this.indexes.set(name, indexData);
+    await this.storage.saveIndex(this.dbName, this.collName, name, {
+      key,
+      unique: options?.unique,
+      sparse: options?.sparse,
+      expireAfterSeconds: options?.expireAfterSeconds,
+    });
+
+    return name;
+  }
+
+  /**
+   * Drop an index
+   */
+  async dropIndex(name: string): Promise<void> {
+    await this.initialize();
+
+    if (name === '_id_') {
+      throw new TsmdbIndexError('cannot drop _id index');
+    }
+
+    if (!this.indexes.has(name)) {
+      throw new TsmdbIndexError(`index not found: ${name}`);
+    }
+
+    this.indexes.delete(name);
+    await this.storage.dropIndex(this.dbName, this.collName, name);
+  }
+
+  /**
+   * Drop all indexes except _id
+   */
+  async dropAllIndexes(): Promise<void> {
+    await this.initialize();
+
+    const names = Array.from(this.indexes.keys()).filter(n => n !== '_id_');
+    for (const name of names) {
+      this.indexes.delete(name);
+      await this.storage.dropIndex(this.dbName, this.collName, name);
+    }
+  }
+
+  /**
+   * List all indexes
+   */
+  async listIndexes(): Promise<IIndexInfo[]> {
+    await this.initialize();
+
+    return Array.from(this.indexes.values()).map(idx => ({
+      v: 2,
+      key: idx.key,
+      name: idx.name,
+      unique: idx.unique || undefined,
+      sparse: idx.sparse || undefined,
+      expireAfterSeconds: idx.expireAfterSeconds,
+    }));
+  }
+
+  /**
+   * Check if an index exists
+   */
+  async indexExists(name: string): Promise<boolean> {
+    await this.initialize();
+    return this.indexes.has(name);
+  }
+
+  /**
+   * Update index entries after document insert
+   */
+  async onInsert(doc: IStoredDocument): Promise<void> {
+    await this.initialize();
+
+    for (const [name, indexData] of this.indexes) {
+      const keyValue = this.extractKeyValue(doc, indexData.key);
+
+      if (keyValue === null && indexData.sparse) {
+        continue;
+      }
+
+      const keyStr = JSON.stringify(keyValue);
+
+      // Check unique constraint
+      if (indexData.unique) {
+        const existing = indexData.hashMap.get(keyStr);
+        if (existing && existing.size > 0) {
+          throw new TsmdbDuplicateKeyError(
+            `E11000 duplicate key error collection: ${this.dbName}.${this.collName} index: ${name}`,
+            indexData.key as Record<string, 1>,
+            keyValue
+          );
+        }
+      }
+
+      // Add to hash map
+      if (!indexData.hashMap.has(keyStr)) {
+        indexData.hashMap.set(keyStr, new Set());
+      }
+      indexData.hashMap.get(keyStr)!.add(doc._id.toHexString());
+
+      // Add to B-tree
+      const btreeSet = indexData.btree.get(keyValue);
+      if (btreeSet) {
+        btreeSet.add(doc._id.toHexString());
+      } else {
+        indexData.btree.set(keyValue, new Set([doc._id.toHexString()]));
+      }
+    }
+  }
+
+  /**
+   * Update index entries after document update
+   */
+  async onUpdate(oldDoc: IStoredDocument, newDoc: IStoredDocument): Promise<void> {
+    await this.initialize();
+
+    for (const [name, indexData] of this.indexes) {
+      const oldKeyValue = this.extractKeyValue(oldDoc, indexData.key);
+      const newKeyValue = this.extractKeyValue(newDoc, indexData.key);
+      const oldKeyStr = JSON.stringify(oldKeyValue);
+      const newKeyStr = JSON.stringify(newKeyValue);
+
+      // Remove old entry if key changed
+      if (oldKeyStr !== newKeyStr) {
+        if (oldKeyValue !== null || !indexData.sparse) {
+          // Remove from hash map
+          const oldHashSet = indexData.hashMap.get(oldKeyStr);
+          if (oldHashSet) {
+            oldHashSet.delete(oldDoc._id.toHexString());
+            if (oldHashSet.size === 0) {
+              indexData.hashMap.delete(oldKeyStr);
+            }
+          }
+
+          // Remove from B-tree
+          const oldBtreeSet = indexData.btree.get(oldKeyValue);
+          if (oldBtreeSet) {
+            oldBtreeSet.delete(oldDoc._id.toHexString());
+            if (oldBtreeSet.size === 0) {
+              indexData.btree.delete(oldKeyValue);
+            }
+          }
+        }
+
+        // Add new entry
+        if (newKeyValue !== null || !indexData.sparse) {
+          // Check unique constraint
+          if (indexData.unique) {
+            const existing = indexData.hashMap.get(newKeyStr);
+            if (existing && existing.size > 0) {
+              throw new TsmdbDuplicateKeyError(
+                `E11000 duplicate key error collection: ${this.dbName}.${this.collName} index: ${name}`,
+                indexData.key as Record<string, 1>,
+                newKeyValue
+              );
+            }
+          }
+
+          // Add to hash map
+          if (!indexData.hashMap.has(newKeyStr)) {
+            indexData.hashMap.set(newKeyStr, new Set());
+          }
+          indexData.hashMap.get(newKeyStr)!.add(newDoc._id.toHexString());
+
+          // Add to B-tree
+          const newBtreeSet = indexData.btree.get(newKeyValue);
+          if (newBtreeSet) {
+            newBtreeSet.add(newDoc._id.toHexString());
+          } else {
+            indexData.btree.set(newKeyValue, new Set([newDoc._id.toHexString()]));
+          }
+        }
+      }
+    }
+  }
+
+  /**
+   * Update index entries after document delete
+   */
+  async onDelete(doc: IStoredDocument): Promise<void> {
+    await this.initialize();
+
+    for (const indexData of this.indexes.values()) {
+      const keyValue = this.extractKeyValue(doc, indexData.key);
+
+      if (keyValue === null && indexData.sparse) {
+        continue;
+      }
+
+      const keyStr = JSON.stringify(keyValue);
+
+      // Remove from hash map
+      const hashSet = indexData.hashMap.get(keyStr);
+      if (hashSet) {
+        hashSet.delete(doc._id.toHexString());
+        if (hashSet.size === 0) {
+          indexData.hashMap.delete(keyStr);
+        }
+      }
+
+      // Remove from B-tree
+      const btreeSet = indexData.btree.get(keyValue);
+      if (btreeSet) {
+        btreeSet.delete(doc._id.toHexString());
+        if (btreeSet.size === 0) {
+          indexData.btree.delete(keyValue);
+        }
+      }
+    }
+  }
+
+  /**
+   * Find the best index for a query
+   */
+  selectIndex(filter: Document): { name: string; data: IIndexData } | null {
+    if (!filter || Object.keys(filter).length === 0) {
+      return null;
+    }
+
+    // Get filter fields and operators
+    const filterInfo = this.analyzeFilter(filter);
+
+    // Score each index
+    let bestIndex: { name: string; data: IIndexData } | null = null;
+    let bestScore = 0;
+
+    for (const [name, indexData] of this.indexes) {
+      const indexFields = Object.keys(indexData.key);
+      let score = 0;
+
+      // Count how many index fields can be used
+      for (const field of indexFields) {
+        const info = filterInfo.get(field);
+        if (!info) break;
+
+        // Equality is best
+        if (info.equality) {
+          score += 2;
+        } else if (info.range) {
+          // Range queries can use B-tree
+          score += 1;
+        } else if (info.in) {
+          score += 1.5;
+        } else {
+          break;
+        }
+      }
+
+      // Prefer unique indexes
+      if (indexData.unique && score > 0) {
+        score += 0.5;
+      }
+
+      if (score > bestScore) {
+        bestScore = score;
+        bestIndex = { name, data: indexData };
+      }
+    }
+
+    return bestIndex;
+  }
+
+  /**
+   * Analyze filter to extract field operators
+   */
+  private analyzeFilter(filter: Document): Map<string, { equality: boolean; range: boolean; in: boolean; ops: Record<string, any> }> {
+    const result = new Map<string, { equality: boolean; range: boolean; in: boolean; ops: Record<string, any> }>();
+
+    for (const [key, value] of Object.entries(filter)) {
+      if (key.startsWith('$')) continue;
+
+      const info = { equality: false, range: false, in: false, ops: {} as Record<string, any> };
+
+      if (typeof value !== 'object' || value === null || value instanceof plugins.bson.ObjectId || value instanceof Date) {
+        info.equality = true;
+        info.ops['$eq'] = value;
+      } else {
+        const ops = value as Record<string, any>;
+        if (ops.$eq !== undefined) {
+          info.equality = true;
+          info.ops['$eq'] = ops.$eq;
+        }
+        if (ops.$in !== undefined) {
+          info.in = true;
+          info.ops['$in'] = ops.$in;
+        }
+        if (ops.$gt !== undefined || ops.$gte !== undefined || ops.$lt !== undefined || ops.$lte !== undefined) {
+          info.range = true;
+          if (ops.$gt !== undefined) info.ops['$gt'] = ops.$gt;
+          if (ops.$gte !== undefined) info.ops['$gte'] = ops.$gte;
+          if (ops.$lt !== undefined) info.ops['$lt'] = ops.$lt;
+          if (ops.$lte !== undefined) info.ops['$lte'] = ops.$lte;
+        }
+      }
+
+      result.set(key, info);
+    }
+
+    return result;
+  }
+
+  /**
+   * Use index to find candidate document IDs (supports range queries with B-tree)
+   */
+  async findCandidateIds(filter: Document): Promise<Set<string> | null> {
+    await this.initialize();
+
+    const index = this.selectIndex(filter);
+    if (!index) return null;
+
+    const filterInfo = this.analyzeFilter(filter);
+    const indexFields = Object.keys(index.data.key);
+
+    // For single-field indexes with range queries, use B-tree
+    if (indexFields.length === 1) {
+      const field = indexFields[0];
+      const info = filterInfo.get(field);
+
+      if (info) {
+        // Handle equality using hash map (faster)
+        if (info.equality) {
+          const keyStr = JSON.stringify(info.ops['$eq']);
+          return index.data.hashMap.get(keyStr) || new Set();
+        }
+
+        // Handle $in using hash map
+        if (info.in) {
+          const results = new Set<string>();
+          for (const val of info.ops['$in']) {
+            const keyStr = JSON.stringify(val);
+            const ids = index.data.hashMap.get(keyStr);
+            if (ids) {
+              for (const id of ids) {
+                results.add(id);
+              }
+            }
+          }
+          return results;
+        }
+
+        // Handle range queries using B-tree
+        if (info.range) {
+          return this.findRangeCandidates(index.data, info.ops);
+        }
+      }
+    } else {
+      // For compound indexes, use hash map with partial key matching
+      const equalityValues: Record<string, any> = {};
+
+      for (const field of indexFields) {
+        const info = filterInfo.get(field);
+        if (!info) break;
+
+        if (info.equality) {
+          equalityValues[field] = info.ops['$eq'];
+        } else if (info.in) {
+          // Handle $in with multiple lookups
+          const results = new Set<string>();
+          for (const val of info.ops['$in']) {
+            equalityValues[field] = val;
+            const keyStr = JSON.stringify(this.buildKeyValue(equalityValues, index.data.key));
+            const ids = index.data.hashMap.get(keyStr);
+            if (ids) {
+              for (const id of ids) {
+                results.add(id);
+              }
+            }
+          }
+          return results;
+        } else {
+          break; // Non-equality/in operator, stop here
+        }
+      }
+
+      if (Object.keys(equalityValues).length > 0) {
+        const keyStr = JSON.stringify(this.buildKeyValue(equalityValues, index.data.key));
+        return index.data.hashMap.get(keyStr) || new Set();
+      }
+    }
+
+    return null;
+  }
+
+  /**
+   * Find candidates using B-tree range scan
+   */
+  private findRangeCandidates(indexData: IIndexData, ops: Record<string, any>): Set<string> {
+    const results = new Set<string>();
+
+    let lowKey: any = undefined;
+    let highKey: any = undefined;
+    let lowInclusive = true;
+    let highInclusive = true;
+
+    if (ops['$gt'] !== undefined) {
+      lowKey = ops['$gt'];
+      lowInclusive = false;
+    }
+    if (ops['$gte'] !== undefined) {
+      lowKey = ops['$gte'];
+      lowInclusive = true;
+    }
+    if (ops['$lt'] !== undefined) {
+      highKey = ops['$lt'];
+      highInclusive = false;
+    }
+    if (ops['$lte'] !== undefined) {
+      highKey = ops['$lte'];
+      highInclusive = true;
+    }
+
+    // Use B-tree range iteration
+    indexData.btree.forRange(lowKey, highKey, lowInclusive, highInclusive, (value, key) => {
+      if (value) {
+        for (const id of value) {
+          results.add(id);
+        }
+      }
+    });
+
+    return results;
+  }
+
+  // ============================================================================
+  // Helper Methods
+  // ============================================================================
+
+  private generateIndexName(key: Record<string, any>): string {
+    return Object.entries(key)
+      .map(([field, dir]) => `${field}_${dir}`)
+      .join('_');
+  }
+
+  private extractKeyValue(doc: Document, key: Record<string, any>): any {
+    const values: any[] = [];
+
+    for (const field of Object.keys(key)) {
+      const value = QueryEngine.getNestedValue(doc, field);
+      values.push(value === undefined ? null : value);
+    }
+
+    // For single-field index, return the value directly
+    if (values.length === 1) {
+      return values[0];
+    }
+
+    return values;
+  }
+
+  private buildKeyValue(values: Record<string, any>, key: Record<string, any>): any {
+    const result: any[] = [];
+
+    for (const field of Object.keys(key)) {
+      result.push(values[field] !== undefined ? values[field] : null);
+    }
+
+    if (result.length === 1) {
+      return result[0];
+    }
+
+    return result;
+  }
+
+  private getFilterFields(filter: Document, prefix = ''): string[] {
+    const fields: string[] = [];
+
+    for (const [key, value] of Object.entries(filter)) {
+      if (key.startsWith('$')) {
+        // Logical operator
+        if (key === '$and' || key === '$or' || key === '$nor') {
+          for (const subFilter of value as Document[]) {
+            fields.push(...this.getFilterFields(subFilter, prefix));
+          }
+        }
+      } else {
+        const fullKey = prefix ? `${prefix}.${key}` : key;
+        fields.push(fullKey);
+
+        // Check for nested filters
+        if (typeof value === 'object' && value !== null && !Array.isArray(value)) {
+          const subKeys = Object.keys(value);
+          if (subKeys.length > 0 && !subKeys[0].startsWith('$')) {
+            fields.push(...this.getFilterFields(value, fullKey));
+          }
+        }
+      }
+    }
+
+    return fields;
+  }
+
+  private getFilterValue(filter: Document, field: string): any {
+    // Handle dot notation
+    const parts = field.split('.');
+    let current: any = filter;
+
+    for (const part of parts) {
+      if (current === null || current === undefined) {
+        return undefined;
+      }
+      current = current[part];
+    }
+
+    return current;
+  }
+}