Reverse index

.changeset/lemon-poems-raise.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/db": patch
+---
+
+Optimization: reverse the index when the direction does not match.

packages/db/src/collection/subscription.ts

@@ -5,7 +5,7 @@ import {
   createFilteredCallback,
 } from "./change-events.js"
 import type { BasicExpression } from "../query/ir.js"
-import type { BaseIndex } from "../indexes/base-index.js"
+import type { IndexInterface } from "../indexes/base-index.js"
 import type { ChangeMessage } from "../types.js"
 import type { CollectionImpl } from "./index.js"
 
@@ -38,7 +38,7 @@ export class CollectionSubscription {
 
   private filteredCallback: (changes: Array<ChangeMessage<any, any>>) => void
 
-  private orderByIndex: BaseIndex<string | number> | undefined
+  private orderByIndex: IndexInterface<string | number> | undefined
 
   constructor(
     private collection: CollectionImpl<any, any, any, any, any>,
@@ -65,7 +65,7 @@ export class CollectionSubscription {
       : this.callback
   }
 
-  setOrderByIndex(index: BaseIndex<any>) {
+  setOrderByIndex(index: IndexInterface<any>) {
     this.orderByIndex = index
   }
 

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

@@ -1,8 +1,9 @@
 import { compileSingleRowExpression } from "../query/compiler/evaluators.js"
 import { comparisonFunctions } from "../query/builder/functions.js"
 import { DEFAULT_COMPARE_OPTIONS, deepEquals } from "../utils.js"
+import type { RangeQueryOptions } from "./btree-index.js"
 import type { CompareOptions } from "../query/builder/types.js"
-import type { BasicExpression } from "../query/ir.js"
+import type { BasicExpression, OrderByDirection } from "../query/ir.js"
 
 /**
  * Operations that indexes can support, imported from available comparison functions
@@ -24,12 +25,57 @@ export interface IndexStats {
   readonly lastUpdated: Date
 }
 
+export interface IndexInterface<
+  TKey extends string | number = string | number,
+> {
+  add: (key: TKey, item: any) => void
+  remove: (key: TKey, item: any) => void
+  update: (key: TKey, oldItem: any, newItem: any) => void
+
+  build: (entries: Iterable<[TKey, any]>) => void
+  clear: () => void
+
+  lookup: (operation: IndexOperation, value: any) => Set<TKey>
+
+  equalityLookup: (value: any) => Set<TKey>
+  inArrayLookup: (values: Array<any>) => Set<TKey>
+
+  rangeQuery: (options: RangeQueryOptions) => Set<TKey>
+  rangeQueryReversed: (options: RangeQueryOptions) => Set<TKey>
+
+  take: (
+    n: number,
+    from?: TKey,
+    filterFn?: (key: TKey) => boolean
+  ) => Array<TKey>
+  takeReversed: (
+    n: number,
+    from?: TKey,
+    filterFn?: (key: TKey) => boolean
+  ) => Array<TKey>
+
+  get keyCount(): number
+  get orderedEntriesArray(): Array<[any, Set<TKey>]>
+  get orderedEntriesArrayReversed(): Array<[any, Set<TKey>]>
+
+  get indexedKeysSet(): Set<TKey>
+  get valueMapData(): Map<any, Set<TKey>>
+
+  supports: (operation: IndexOperation) => boolean
+
+  matchesField: (fieldPath: Array<string>) => boolean
+  matchesCompareOptions: (compareOptions: CompareOptions) => boolean
+  matchesDirection: (direction: OrderByDirection) => boolean
+
+  getStats: () => IndexStats
+}
+
 /**
  * Base abstract class that all index types extend
  */
-export abstract class BaseIndex<
-  TKey extends string | number = string | number,
-> {
+export abstract class BaseIndex<TKey extends string | number = string | number>
+  implements IndexInterface<TKey>
+{
   public readonly id: number
   public readonly name?: string
   public readonly expression: BasicExpression
@@ -65,7 +111,20 @@ export abstract class BaseIndex<
     from?: TKey,
     filterFn?: (key: TKey) => boolean
   ): Array<TKey>
+  abstract takeReversed(
+    n: number,
+    from?: TKey,
+    filterFn?: (key: TKey) => boolean
+  ): Array<TKey>
   abstract get keyCount(): number
+  abstract equalityLookup(value: any): Set<TKey>
+  abstract inArrayLookup(values: Array<any>): Set<TKey>
+  abstract rangeQuery(options: RangeQueryOptions): Set<TKey>
+  abstract rangeQueryReversed(options: RangeQueryOptions): Set<TKey>
+  abstract get orderedEntriesArray(): Array<[any, Set<TKey>]>
+  abstract get orderedEntriesArrayReversed(): Array<[any, Set<TKey>]>
+  abstract get indexedKeysSet(): Set<TKey>
+  abstract get valueMapData(): Map<any, Set<TKey>>
 
   // Common methods
   supports(operation: IndexOperation): boolean {
@@ -80,8 +139,31 @@ export abstract class BaseIndex<
     )
   }
 
+  /**
+   * Checks if the compare options match the index's compare options.
+   * The direction is ignored because the index can be reversed if the direction is different.
+   */
   matchesCompareOptions(compareOptions: CompareOptions): boolean {
-    return deepEquals(this.compareOptions, compareOptions)
+    const thisCompareOptionsWithoutDirection = {
+      ...this.compareOptions,
+      direction: undefined,
+    }
+    const compareOptionsWithoutDirection = {
+      ...compareOptions,
+      direction: undefined,
+    }
+
+    return deepEquals(
+      thisCompareOptionsWithoutDirection,
+      compareOptionsWithoutDirection
+    )
+  }
+
+  /**
+   * Checks if the index matches the provided direction.
+   */
+  matchesDirection(direction: OrderByDirection): boolean {
+    return this.compareOptions.direction === direction
   }
 
   getStats(): IndexStats {

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

@@ -245,15 +245,26 @@ export class BTreeIndex<
   }
 
   /**
-   * Returns the next n items after the provided item or the first n items if no from item is provided.
-   * @param n - The number of items to return
-   * @param from - The item to start from (exclusive). Starts from the smallest item (inclusive) if not provided.
-   * @returns The next n items after the provided key. Returns the first n items if no from item is provided.
+   * Performs a reversed range query
    */
-  take(n: number, from?: any, filterFn?: (key: TKey) => boolean): Array<TKey> {
+  rangeQueryReversed(options: RangeQueryOptions = {}): Set<TKey> {
+    const { from, to, fromInclusive = true, toInclusive = true } = options
+    return this.rangeQuery({
+      from: to ?? this.orderedEntries.maxKey(),
+      to: from ?? this.orderedEntries.minKey(),
+      fromInclusive: toInclusive,
+      toInclusive: fromInclusive,
+    })
+  }
+
+  private takeInternal(
+    n: number,
+    nextPair: (k?: any) => [any, any] | undefined,
+    from?: any,
+    filterFn?: (key: TKey) => boolean
+  ): Array<TKey> {
     const keysInResult: Set<TKey> = new Set()
     const result: Array<TKey> = []
-    const nextPair = (k?: any) => this.orderedEntries.nextHigherPair(k)
     let pair: [any, any] | undefined
     let key = normalizeValue(from)
 
@@ -275,6 +286,32 @@ export class BTreeIndex<
     return result
   }
 
+  /**
+   * Returns the next n items after the provided item or the first n items if no from item is provided.
+   * @param n - The number of items to return
+   * @param from - The item to start from (exclusive). Starts from the smallest item (inclusive) if not provided.
+   * @returns The next n items after the provided key. Returns the first n items if no from item is provided.
+   */
+  take(n: number, from?: any, filterFn?: (key: TKey) => boolean): Array<TKey> {
+    const nextPair = (k?: any) => this.orderedEntries.nextHigherPair(k)
+    return this.takeInternal(n, nextPair, from, filterFn)
+  }
+
+  /**
+   * Returns the next n items **before** the provided item (in descending order) or the last n items if no from item is provided.
+   * @param n - The number of items to return
+   * @param from - The item to start from (exclusive). Starts from the largest item (inclusive) if not provided.
+   * @returns The next n items **before** the provided key. Returns the last n items if no from item is provided.
+   */
+  takeReversed(
+    n: number,
+    from?: any,
+    filterFn?: (key: TKey) => boolean
+  ): Array<TKey> {
+    const nextPair = (k?: any) => this.orderedEntries.nextLowerPair(k)
+    return this.takeInternal(n, nextPair, from, filterFn)
+  }
+
   /**
    * Performs an IN array lookup
    */
@@ -303,6 +340,13 @@ export class BTreeIndex<
       .map((key) => [key, this.valueMap.get(key) ?? new Set()])
   }
 
+  get orderedEntriesArrayReversed(): Array<[any, Set<TKey>]> {
+    return this.takeReversed(this.orderedEntries.size).map((key) => [
+      key,
+      this.valueMap.get(key) ?? new Set(),
+    ])
+  }
+
   get valueMapData(): Map<any, Set<TKey>> {
     return this.valueMap
   }

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

@@ -0,0 +1,120 @@
+import type { CompareOptions } from "../query/builder/types"
+import type { OrderByDirection } from "../query/ir"
+import type { IndexInterface, IndexOperation, IndexStats } from "./base-index"
+import type { RangeQueryOptions } from "./btree-index"
+
+export class ReverseIndex<TKey extends string | number>
+  implements IndexInterface<TKey>
+{
+  private originalIndex: IndexInterface<TKey>
+
+  constructor(index: IndexInterface<TKey>) {
+    this.originalIndex = index
+  }
+
+  // Define the reversed operations
+
+  lookup(operation: IndexOperation, value: any): Set<TKey> {
+    const reverseOperation =
+      operation === `gt`
+        ? `lt`
+        : operation === `gte`
+          ? `lte`
+          : operation === `lt`
+            ? `gt`
+            : operation === `lte`
+              ? `gte`
+              : operation
+    return this.originalIndex.lookup(reverseOperation, value)
+  }
+
+  rangeQuery(options: RangeQueryOptions = {}): Set<TKey> {
+    return this.originalIndex.rangeQueryReversed(options)
+  }
+
+  rangeQueryReversed(options: RangeQueryOptions = {}): Set<TKey> {
+    return this.originalIndex.rangeQuery(options)
+  }
+
+  take(n: number, from?: any, filterFn?: (key: TKey) => boolean): Array<TKey> {
+    return this.originalIndex.takeReversed(n, from, filterFn)
+  }
+
+  takeReversed(
+    n: number,
+    from?: any,
+    filterFn?: (key: TKey) => boolean
+  ): Array<TKey> {
+    return this.originalIndex.take(n, from, filterFn)
+  }
+
+  get orderedEntriesArray(): Array<[any, Set<TKey>]> {
+    return this.originalIndex.orderedEntriesArrayReversed
+  }
+
+  get orderedEntriesArrayReversed(): Array<[any, Set<TKey>]> {
+    return this.originalIndex.orderedEntriesArray
+  }
+
+  // All operations below delegate to the original index
+
+  supports(operation: IndexOperation): boolean {
+    return this.originalIndex.supports(operation)
+  }
+
+  matchesField(fieldPath: Array<string>): boolean {
+    return this.originalIndex.matchesField(fieldPath)
+  }
+
+  matchesCompareOptions(compareOptions: CompareOptions): boolean {
+    return this.originalIndex.matchesCompareOptions(compareOptions)
+  }
+
+  matchesDirection(direction: OrderByDirection): boolean {
+    return this.originalIndex.matchesDirection(direction)
+  }
+
+  getStats(): IndexStats {
+    return this.originalIndex.getStats()
+  }
+
+  add(key: TKey, item: any): void {
+    this.originalIndex.add(key, item)
+  }
+
+  remove(key: TKey, item: any): void {
+    this.originalIndex.remove(key, item)
+  }
+
+  update(key: TKey, oldItem: any, newItem: any): void {
+    this.originalIndex.update(key, oldItem, newItem)
+  }
+
+  build(entries: Iterable<[TKey, any]>): void {
+    this.originalIndex.build(entries)
+  }
+
+  clear(): void {
+    this.originalIndex.clear()
+  }
+
+  get keyCount(): number {
+    return this.originalIndex.keyCount
+  }
+
+  equalityLookup(value: any): Set<TKey> {
+    return this.originalIndex.equalityLookup(value)
+  }
+
+  inArrayLookup(values: Array<any>): Set<TKey> {
+    return this.originalIndex.inArrayLookup(values)
+  }
+
+  get indexedKeysSet(): Set<TKey> {
+    return this.originalIndex.indexedKeysSet
+  }
+
+  get valueMapData(): Map<any, Set<TKey>> {
+    return this.originalIndex.valueMapData
+  }
+}

packages/db/src/query/compiler/order-by.ts

@@ -9,7 +9,7 @@ import type { CompiledSingleRowExpression } from "./evaluators.js"
 import type { OrderByClause, QueryIR, Select } from "../ir.js"
 import type { NamespacedAndKeyedStream, NamespacedRow } from "../../types.js"
 import type { IStreamBuilder, KeyValue } from "@tanstack/db-ivm"
-import type { BaseIndex } from "../../indexes/base-index.js"
+import type { IndexInterface } from "../../indexes/base-index.js"
 import type { Collection } from "../../collection/index.js"
 
 export type OrderByOptimizationInfo = {
@@ -20,7 +20,7 @@ export type OrderByOptimizationInfo = {
     b: Record<string, unknown> | null | undefined
   ) => number
   valueExtractorForRawRow: (row: Record<string, unknown>) => any
-  index: BaseIndex<string | number>
+  index: IndexInterface<string | number>
   dataNeeded?: () => number
 }
 
@@ -151,11 +151,12 @@ export function processOrderBy(
         return compare(extractedA, extractedB)
       }
 
-      const index: BaseIndex<string | number> | undefined = findIndexForField(
-        followRefCollection.indexes,
-        followRefResult.path,
-        clause.compareOptions
-      )
+      const index: IndexInterface<string | number> | undefined =
+        findIndexForField(
+          followRefCollection.indexes,
+          followRefResult.path,
+          clause.compareOptions
+        )
 
       if (index && index.supports(`gt`)) {
         // We found an index that we can use to lazily load ordered data