Optimized order by that lazily loads more data when needed (#410)

Author: kevin-dp

.changeset/common-clubs-add.md

@@ -0,0 +1,6 @@
+---
+"@tanstack/db-ivm": patch
+"@tanstack/db": patch
+---
+
+Optimize order by to lazily load ordered data if a range index is available on the field that is being ordered on.

packages/db-ivm/src/operators/orderBy.ts

@@ -12,6 +12,10 @@ export interface OrderByOptions<Ve> {
   offset?: number
 }
 
+type OrderByWithFractionalIndexOptions<Ve> = OrderByOptions<Ve> & {
+  setSizeCallback?: (getSize: () => number) => void
+}
+
 /**
  * Orders the elements and limits the number of results, with optional offset
  * This requires a keyed stream, and uses the `topK` operator to order all the elements.
@@ -136,13 +140,14 @@ export function orderByWithFractionalIndexBase<
   valueExtractor: (
     value: T extends KeyValue<unknown, infer V> ? V : never
   ) => Ve,
-  options?: OrderByOptions<Ve>
+  options?: OrderByWithFractionalIndexOptions<Ve>
 ) {
   type KeyType = T extends KeyValue<infer K, unknown> ? K : never
   type ValueType = T extends KeyValue<unknown, infer V> ? V : never
 
   const limit = options?.limit ?? Infinity
   const offset = options?.offset ?? 0
+  const setSizeCallback = options?.setSizeCallback
   const comparator =
     options?.comparator ??
     ((a, b) => {
@@ -162,6 +167,7 @@ export function orderByWithFractionalIndexBase<
         {
           limit,
           offset,
+          setSizeCallback,
         }
       ),
       consolidate()
@@ -185,7 +191,7 @@ export function orderByWithFractionalIndex<
   valueExtractor: (
     value: T extends KeyValue<unknown, infer V> ? V : never
   ) => Ve,
-  options?: OrderByOptions<Ve>
+  options?: OrderByWithFractionalIndexOptions<Ve>
 ) {
   return orderByWithFractionalIndexBase(
     topKWithFractionalIndex,

packages/db-ivm/src/operators/topKWithFractionalIndex.ts

@@ -9,6 +9,7 @@ import type { IStreamBuilder, PipedOperator } from "../types.js"
 export interface TopKWithFractionalIndexOptions {
   limit?: number
   offset?: number
+  setSizeCallback?: (getSize: () => number) => void
 }
 
 export type TopKChanges<V> = {
@@ -23,6 +24,7 @@ export type TopKChanges<V> = {
  * and returns changes to the topK.
  */
 export interface TopK<V> {
+  size: number
   insert: (value: V) => TopKChanges<V>
   delete: (value: V) => TopKChanges<V>
 }
@@ -49,6 +51,13 @@ class TopKArray<V> implements TopK<V> {
     this.#comparator = comparator
   }
 
+  get size(): number {
+    const offset = this.#topKStart
+    const limit = this.#topKEnd - this.#topKStart
+    const available = this.#sortedValues.length - offset
+    return Math.max(0, Math.min(limit, available))
+  }
+
   insert(value: V): TopKChanges<V> {
     const result: TopKChanges<V> = { moveIn: null, moveOut: null }
 
@@ -169,6 +178,8 @@ export class TopKWithFractionalIndexOperator<K, T> extends UnaryOperator<
    */
   #topK: TopK<TaggedValue<K, T>>
 
+  #limit: number
+
   constructor(
     id: number,
     inputA: DifferenceStreamReader<[K, T]>,
@@ -177,7 +188,7 @@ export class TopKWithFractionalIndexOperator<K, T> extends UnaryOperator<
     options: TopKWithFractionalIndexOptions
   ) {
     super(id, inputA, output)
-    const limit = options.limit ?? Infinity
+    this.#limit = options.limit ?? Infinity
     const offset = options.offset ?? 0
     const compareTaggedValues = (
       a: TaggedValue<K, T>,
@@ -193,7 +204,8 @@ export class TopKWithFractionalIndexOperator<K, T> extends UnaryOperator<
       const tieBreakerB = getTag(b)
       return tieBreakerA - tieBreakerB
     }
-    this.#topK = this.createTopK(offset, limit, compareTaggedValues)
+    this.#topK = this.createTopK(offset, this.#limit, compareTaggedValues)
+    options.setSizeCallback?.(() => this.#topK.size)
   }
 
   protected createTopK(

packages/db-ivm/src/operators/topKWithFractionalIndexBTree.ts

@@ -75,6 +75,13 @@ class TopKTree<V> implements TopK<V> {
     this.#tree = new BTree(undefined, comparator)
   }
 
+  get size(): number {
+    const offset = this.#topKStart
+    const limit = this.#topKEnd - this.#topKStart
+    const available = this.#tree.size - offset
+    return Math.max(0, Math.min(limit, available))
+  }
+
   /**
    * Insert a *new* value.
    * Ignores the value if it is already present.

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

@@ -29,7 +29,8 @@ export function ensureIndexForField<
 >(
   fieldName: string,
   fieldPath: Array<string>,
-  collection: CollectionImpl<T, TKey, any, any, any>
+  collection: CollectionImpl<T, TKey, any, any, any>,
+  compareFn?: (a: any, b: any) => number
 ) {
   if (!shouldAutoIndex(collection)) {
     return
@@ -49,6 +50,7 @@ export function ensureIndexForField<
     collection.createIndex((row) => (row as any)[fieldName], {
       name: `auto_${fieldName}`,
       indexType: BTreeIndex,
+      options: compareFn ? { compareFn } : {},
     })
   } catch (error) {
     console.warn(`Failed to create auto-index for field "${fieldName}":`, error)

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

@@ -1,6 +1,6 @@
 import { compileSingleRowExpression } from "../query/compiler/evaluators.js"
 import { comparisonFunctions } from "../query/builder/functions.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
@@ -56,6 +56,11 @@ export abstract class BaseIndex<
   abstract build(entries: Iterable<[TKey, any]>): void
   abstract clear(): void
   abstract lookup(operation: IndexOperation, value: any): Set<TKey>
+  abstract take(
+    n: number,
+    direction?: OrderByDirection,
+    from?: TKey
+  ): Array<TKey>
   abstract get keyCount(): number
 
   // Common methods

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

@@ -230,6 +230,35 @@ 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): Array<TKey> {
+    const keysInResult: Set<TKey> = new Set()
+    const result: Array<TKey> = []
+    const nextKey = (k?: any) => this.orderedEntries.nextHigherKey(k)
+    let key = from
+
+    while ((key = nextKey(key)) && result.length < n) {
+      const keys = this.valueMap.get(key)
+      if (keys) {
+        const it = keys.values()
+        let ks: TKey | undefined
+        while (result.length < n && (ks = it.next().value)) {
+          if (!keysInResult.has(ks)) {
+            result.push(ks)
+            keysInResult.add(ks)
+          }
+        }
+      }
+    }
+
+    return result
+  }
+
   /**
    * Performs an IN array lookup
    */

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

@@ -8,7 +8,7 @@ export interface IndexOptions<TResolver extends IndexResolver = IndexResolver> {
   indexType?: TResolver
   options?: TResolver extends IndexConstructor<any>
     ? TResolver extends new (
-        id: string,
+        id: number,
         expr: any,
         name?: string,
         options?: infer O
@@ -17,7 +17,7 @@ export interface IndexOptions<TResolver extends IndexResolver = IndexResolver> {
       : never
     : TResolver extends () => Promise<infer TCtor>
       ? TCtor extends new (
-          id: string,
+          id: number,
           expr: any,
           name?: string,
           options?: infer O

packages/db/src/query/compiler/evaluators.ts

@@ -20,9 +20,12 @@ export type CompiledSingleRowExpression = (item: Record<string, unknown>) => any
  * Compiles an expression into an optimized evaluator function.
  * This eliminates branching during evaluation by pre-compiling the expression structure.
  */
-export function compileExpression(expr: BasicExpression): CompiledExpression {
-  const compiledFn = compileExpressionInternal(expr, false)
-  return compiledFn as CompiledExpression
+export function compileExpression(
+  expr: BasicExpression,
+  isSingleRow: boolean = false
+): CompiledExpression | CompiledSingleRowExpression {
+  const compiledFn = compileExpressionInternal(expr, isSingleRow)
+  return compiledFn
 }
 
 /**

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

@@ -166,7 +166,9 @@ export function processGroupBy(
   const mapping = validateAndCreateMapping(groupByClause, selectClause)
 
   // Pre-compile groupBy expressions
-  const compiledGroupByExpressions = groupByClause.map(compileExpression)
+  const compiledGroupByExpressions = groupByClause.map((e) =>
+    compileExpression(e)
+  )
 
   // Create a key extractor function using simple __key_X format
   const keyExtractor = ([, row]: [