Index that uses B+ tree (#302)

Author: kevin-dp

.changeset/better-owls-judge.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/db": patch
+---
+
+Remove OrderedIndex in favor of more efficient BTree index.

packages/db/src/collection.ts

@@ -4,7 +4,7 @@ import {
   createSingleRowRefProxy,
   toExpression,
 } from "./query/builder/ref-proxy"
-import { OrderedIndex } from "./indexes/ordered-index.js"
+import { BTreeIndex } from "./indexes/btree-index.js"
 import { IndexProxy, LazyIndexWrapper } from "./indexes/lazy-index.js"
 import { ensureIndexForExpression } from "./indexes/auto-index.js"
 import { createTransaction, getActiveTransaction } from "./transactions"
@@ -1297,12 +1297,12 @@ export class CollectionImpl<
    * @returns An index proxy that provides access to the index when ready
    *
    * @example
-   * // Create a default ordered index
+   * // Create a default B+ tree index
    * const ageIndex = collection.createIndex((row) => row.age)
    *
    * // Create a ordered index with custom options
    * const ageIndex = collection.createIndex((row) => row.age, {
-   *   indexType: OrderedIndex,
+   *   indexType: BTreeIndex,
    *   options: { compareFn: customComparator },
    *   name: 'age_btree'
    * })
@@ -1316,9 +1316,7 @@ export class CollectionImpl<
    *   options: { language: 'en' }
    * })
    */
-  public createIndex<
-    TResolver extends IndexResolver<TKey> = typeof OrderedIndex,
-  >(
+  public createIndex<TResolver extends IndexResolver<TKey> = typeof BTreeIndex>(
     indexCallback: (row: SingleRowRefProxy<T>) => any,
     config: IndexOptions<TResolver> = {}
   ): IndexProxy<TKey> {
@@ -1329,8 +1327,8 @@ export class CollectionImpl<
     const indexExpression = indexCallback(singleRowRefProxy)
     const expression = toExpression(indexExpression)
 
-    // Default to OrderedIndex if no type specified
-    const resolver = config.indexType ?? (OrderedIndex as unknown as TResolver)
+    // Default to BTreeIndex if no type specified
+    const resolver = config.indexType ?? (BTreeIndex as unknown as TResolver)
 
     // Create lazy wrapper
     const lazyIndex = new LazyIndexWrapper<TKey>(
@@ -1344,13 +1342,13 @@ export class CollectionImpl<
 
     this.lazyIndexes.set(indexId, lazyIndex)
 
-    // For OrderedIndex, resolve immediately and synchronously
-    if ((resolver as unknown) === OrderedIndex) {
+    // For BTreeIndex, resolve immediately and synchronously
+    if ((resolver as unknown) === BTreeIndex) {
       try {
         const resolvedIndex = lazyIndex.getResolved()
         this.resolvedIndexes.set(indexId, resolvedIndex)
       } catch (error) {
-        console.warn(`Failed to resolve OrderedIndex:`, error)
+        console.warn(`Failed to resolve BTreeIndex:`, error)
       }
     } else if (typeof resolver === `function` && resolver.prototype) {
       // Other synchronous constructors - resolve immediately

packages/db/src/index.ts

@@ -12,7 +12,7 @@ export * from "./errors"
 
 // Index system exports
 export * from "./indexes/base-index.js"
-export * from "./indexes/ordered-index.js"
+export * from "./indexes/btree-index.js"
 export * from "./indexes/lazy-index.js"
 export { type IndexOptions } from "./indexes/index-options.js"
 

packages/db/src/indexes/auto-index.ts

@@ -1,4 +1,4 @@
-import { OrderedIndex } from "./ordered-index"
+import { BTreeIndex } from "./btree-index"
 import type { BasicExpression } from "../query/ir"
 import type { CollectionImpl } from "../collection"
 
@@ -46,7 +46,7 @@ export function ensureIndexForExpression<
     try {
       collection.createIndex((row) => (row as any)[fieldName], {
         name: `auto_${fieldName}`,
-        indexType: OrderedIndex,
+        indexType: BTreeIndex,
       })
     } catch (error) {
       console.warn(

packages/db/src/indexes/ordered-index.ts renamed to packages/db/src/indexes/btree-index.ts

@@ -1,12 +1,13 @@
 import { ascComparator } from "../utils/comparison.js"
-import { findInsertPosition } from "../utils/array-utils.js"
+import { BTree } from "../utils/btree.js"
 import { BaseIndex } from "./base-index.js"
+import type { BasicExpression } from "../query/ir.js"
 import type { IndexOperation } from "./base-index.js"
 
 /**
  * Options for Ordered index
  */
-export interface OrderedIndexOptions {
+export interface BTreeIndexOptions {
   compareFn?: (a: any, b: any) => number
 }
 
@@ -21,10 +22,10 @@ export interface RangeQueryOptions {
 }
 
 /**
- * Ordered index for sorted data with range queries
+ * B+Tree index for sorted data with range queries
  * This maintains items in sorted order and provides efficient range operations
  */
-export class OrderedIndex<
+export class BTreeIndex<
   TKey extends string | number = string | number,
 > extends BaseIndex<TKey> {
   public readonly supportedOperations = new Set<IndexOperation>([
@@ -37,15 +38,26 @@ export class OrderedIndex<
   ])
 
   // Internal data structures - private to hide implementation details
-  private orderedEntries: Array<[any, Set<TKey>]> = []
-  private valueMap = new Map<any, Set<TKey>>()
+  // The `orderedEntries` B+ tree is used for efficient range queries
+  // The `valueMap` is used for O(1) lookups of PKs by indexed value
+  private orderedEntries: BTree<any, undefined> // we don't associate values with the keys of the B+ tree (the keys are indexed values)
+  private valueMap = new Map<any, Set<TKey>>() // instead we store a mapping of indexed values to a set of PKs
   private indexedKeys = new Set<TKey>()
   private compareFn: (a: any, b: any) => number = ascComparator
 
-  protected initialize(options?: OrderedIndexOptions): void {
+  constructor(
+    id: number,
+    expression: BasicExpression,
+    name?: string,
+    options?: any
+  ) {
+    super(id, expression, name, options)
     this.compareFn = options?.compareFn ?? ascComparator
+    this.orderedEntries = new BTree(this.compareFn)
   }
 
+  protected initialize(_options?: BTreeIndexOptions): void {}
+
   /**
    * Adds a value to the index
    */
@@ -67,14 +79,7 @@ export class OrderedIndex<
       // Create new set for this value
       const keySet = new Set<TKey>([key])
       this.valueMap.set(indexedValue, keySet)
-
-      // Find correct position in ordered entries using binary search
-      const insertIndex = findInsertPosition(
-        this.orderedEntries,
-        indexedValue,
-        this.compareFn
-      )
-      this.orderedEntries.splice(insertIndex, 0, [indexedValue, keySet])
+      this.orderedEntries.set(indexedValue, undefined)
     }
 
     this.indexedKeys.add(key)
@@ -104,13 +109,8 @@ export class OrderedIndex<
       if (keySet.size === 0) {
         this.valueMap.delete(indexedValue)
 
-        // Find and remove from ordered entries
-        const index = this.orderedEntries.findIndex(
-          ([value]) => this.compareFn(value, indexedValue) === 0
-        )
-        if (index !== -1) {
-          this.orderedEntries.splice(index, 1)
-        }
+        // Remove from ordered entries
+        this.orderedEntries.delete(indexedValue)
       }
     }
 
@@ -141,7 +141,7 @@ export class OrderedIndex<
    * Clears all data from the index
    */
   clear(): void {
-    this.orderedEntries = []
+    this.orderedEntries.clear()
     this.valueMap.clear()
     this.indexedKeys.clear()
     this.updateTimestamp()
@@ -175,7 +175,7 @@ export class OrderedIndex<
         result = this.inArrayLookup(value)
         break
       default:
-        throw new Error(`Operation ${operation} not supported by OrderedIndex`)
+        throw new Error(`Operation ${operation} not supported by BTreeIndex`)
     }
 
     this.trackLookup(startTime)
@@ -206,70 +206,26 @@ export class OrderedIndex<
     const { from, to, fromInclusive = true, toInclusive = true } = options
     const result = new Set<TKey>()
 
-    if (this.orderedEntries.length === 0) {
-      return result
-    }
-
-    // Find start position
-    let startIndex = 0
-    if (from !== undefined) {
-      const fromInsertIndex = findInsertPosition(
-        this.orderedEntries,
-        from,
-        this.compareFn
-      )
-
-      if (fromInclusive) {
-        // Include values equal to 'from'
-        startIndex = fromInsertIndex
-      } else {
-        // Exclude values equal to 'from'
-        startIndex = fromInsertIndex
-        // Skip the value if it exists at this position
-        if (
-          startIndex < this.orderedEntries.length &&
-          this.compareFn(this.orderedEntries[startIndex]![0], from) === 0
-        ) {
-          startIndex++
+    const fromKey = from ?? this.orderedEntries.minKey()
+    const toKey = to ?? this.orderedEntries.maxKey()
+
+    this.orderedEntries.forRange(
+      fromKey,
+      toKey,
+      toInclusive,
+      (indexedValue, _) => {
+        if (!fromInclusive && this.compareFn(indexedValue, from) === 0) {
+          // the B+ tree `forRange` method does not support exclusive lower bounds
+          // so we need to exclude it manually
+          return
         }
-      }
-    }
-
-    // Find end position
-    let endIndex = this.orderedEntries.length
-    if (to !== undefined) {
-      const toInsertIndex = findInsertPosition(
-        this.orderedEntries,
-        to,
-        this.compareFn
-      )
 
-      if (toInclusive) {
-        // Include values equal to 'to'
-        endIndex = toInsertIndex
-        // Include the value if it exists at this position
-        if (
-          toInsertIndex < this.orderedEntries.length &&
-          this.compareFn(this.orderedEntries[toInsertIndex]![0], to) === 0
-        ) {
-          endIndex = toInsertIndex + 1
+        const keys = this.valueMap.get(indexedValue)
+        if (keys) {
+          keys.forEach((key) => result.add(key))
         }
-      } else {
-        // Exclude values equal to 'to'
-        endIndex = toInsertIndex
       }
-    }
-
-    // Ensure startIndex doesn't exceed endIndex
-    if (startIndex >= endIndex) {
-      return result
-    }
-
-    // Collect keys from the range
-    for (let i = startIndex; i < endIndex; i++) {
-      const keys = this.orderedEntries[i]![1]
-      keys.forEach((key) => result.add(key))
-    }
+    )
 
     return result
   }
@@ -297,6 +253,8 @@ export class OrderedIndex<
 
   get orderedEntriesArray(): Array<[any, Set<TKey>]> {
     return this.orderedEntries
+      .keysArray()
+      .map((key) => [key, this.valueMap.get(key) ?? new Set()])
   }
 
   get valueMapData(): Map<any, Set<TKey>> {