@tanstack/db 0.1.3 → 0.1.4

package/src/query/live-query-collection.ts

@@ -1,8 +1,10 @@
 import { D2, MultiSet, output } from "@tanstack/db-ivm"
 import { createCollection } from "../collection.js"
+import { createFilterFunctionFromExpression } from "../change-events.js"
 import { compileQuery } from "./compiler/index.js"
 import { buildQuery, getQueryIR } from "./builder/index.js"
 import { convertToBasicExpression } from "./compiler/expressions.js"
+import type { OrderByOptimizationInfo } from "./compiler/order-by.js"
 import type { InitialQueryBuilder, QueryBuilder } from "./builder/index.js"
 import type { Collection } from "../collection.js"
 import type {
@@ -16,6 +18,7 @@ import type {
 import type { Context, GetResult } from "./builder/types.js"
 import type { MultiSetArray, RootStreamBuilder } from "@tanstack/db-ivm"
 import type { BasicExpression } from "./ir.js"
+import type { LazyCollectionCallbacks } from "./compiler/joins.js"
 
 // Global counter for auto-generated collection IDs
 let liveQueryCollectionCounter = 0
@@ -163,6 +166,12 @@ export function liveQueryCollectionOptions<
   const collections = extractCollectionsFromQuery(query)
 
   const allCollectionsReady = () => {
+    return Object.values(collections).every((collection) =>
+      collection.isReady()
+    )
+  }
+
+  const allCollectionsReadyOrInitialCommit = () => {
     return Object.values(collections).every(
       (collection) =>
         collection.status === `ready` || collection.status === `initialCommit`
@@ -176,6 +185,14 @@ export function liveQueryCollectionOptions<
     | Map<string, BasicExpression<boolean>>
     | undefined
 
+  // Map of collection IDs to functions that load keys for that lazy collection
+  const lazyCollectionsCallbacks: Record<string, LazyCollectionCallbacks> = {}
+  // Set of collection IDs that are lazy collections
+  const lazyCollections = new Set<string>()
+  // Set of collection IDs that include an optimizable ORDER BY clause
+  const optimizableOrderByCollections: Record<string, OrderByOptimizationInfo> =
+    {}
+
   const compileBasePipeline = () => {
     graphCache = new D2()
     inputsCache = Object.fromEntries(
@@ -189,7 +206,14 @@ export function liveQueryCollectionOptions<
     ;({
       pipeline: pipelineCache,
       collectionWhereClauses: collectionWhereClausesCache,
-    } = compileQuery(query, inputsCache as Record<string, KeyedStream>))
+    } = compileQuery(
+      query,
+      inputsCache as Record<string, KeyedStream>,
+      collections,
+      lazyCollectionsCallbacks,
+      lazyCollections,
+      optimizableOrderByCollections
+    ))
   }
 
   const maybeCompileBasePipeline = () => {
@@ -292,10 +316,24 @@ export function liveQueryCollectionOptions<
 
       graph.finalize()
 
-      const maybeRunGraph = () => {
+      let subscribedToAllCollections = false
+
+      // The callback function is called after the graph has run.
+      // This gives the callback a chance to load more data if needed,
+      // that's used to optimize orderBy operators that set a limit,
+      // in order to load some more data if we still don't have enough rows after the pipeline has run.
+      // That can happend because even though we load N rows, the pipeline might filter some of these rows out
+      // causing the orderBy operator to receive less than N rows or even no rows at all.
+      // So this callback would notice that it doesn't have enough rows and load some more.
+      // The callback returns a boolean, when it's true it's done loading data and we can mark the collection as ready.
+      const maybeRunGraph = (callback?: () => boolean) => {
         // We only run the graph if all the collections are ready
-        if (allCollectionsReady()) {
+        if (
+          allCollectionsReadyOrInitialCommit() &&
+          subscribedToAllCollections
+        ) {
           graph.run()
+          const ready = callback?.() ?? true
           // On the initial run, we may need to do an empty commit to ensure that
           // the collection is initialized
           if (messagesCount === 0) {
@@ -303,7 +341,9 @@ export function liveQueryCollectionOptions<
             commit()
           }
           // Mark the collection as ready after the first successful run
-          markReady()
+          if (ready && allCollectionsReady()) {
+            markReady()
+          }
         }
       }
 
@@ -319,6 +359,240 @@ export function liveQueryCollectionOptions<
             ? collectionWhereClausesCache.get(collectionAlias)
             : undefined
 
+        const sendChangesToPipeline = (
+          changes: Iterable<ChangeMessage<any, string | number>>,
+          callback?: () => boolean
+        ) => {
+          sendChangesToInput(input, changes, collection.config.getKey)
+          maybeRunGraph(callback)
+        }
+
+        // Wraps the sendChangesToPipeline function
+        // in order to turn `update`s into `insert`s
+        // for keys that have not been sent to the pipeline yet
+        // and filter out deletes for keys that have not been sent
+        const sendVisibleChangesToPipeline = (
+          changes: Array<ChangeMessage<any, string | number>>,
+          loadedInitialState: boolean,
+          sentKeys: Set<string | number>
+        ) => {
+          if (loadedInitialState) {
+            // There was no index for the join key
+            // so we loaded the initial state
+            // so we can safely assume that the pipeline has seen all keys
+            return sendChangesToPipeline(changes)
+          }
+
+          const newChanges = []
+          for (const change of changes) {
+            let newChange = change
+            if (!sentKeys.has(change.key)) {
+              if (change.type === `update`) {
+                newChange = { ...change, type: `insert` }
+              } else if (change.type === `delete`) {
+                // filter out deletes for keys that have not been sent
+                continue
+              }
+            }
+            newChanges.push(newChange)
+          }
+
+          return sendChangesToPipeline(newChanges)
+        }
+
+        const loadKeys = (
+          keys: Iterable<string | number>,
+          sentKeys: Set<string | number>,
+          filterFn: (item: object) => boolean
+        ) => {
+          for (const key of keys) {
+            // Only load the key once
+            if (sentKeys.has(key)) continue
+
+            const value = collection.get(key)
+            if (value !== undefined && filterFn(value)) {
+              sentKeys.add(key)
+              sendChangesToPipeline([{ type: `insert`, key, value }])
+            }
+          }
+        }
+
+        const subscribeToAllChanges = (
+          whereExpression: BasicExpression<boolean> | undefined
+        ) => {
+          const unsubscribe = collection.subscribeChanges(
+            sendChangesToPipeline,
+            {
+              includeInitialState: true,
+              ...(whereExpression ? { whereExpression } : undefined),
+            }
+          )
+          return unsubscribe
+        }
+
+        // Subscribes to all changes but without the initial state
+        // such that we can load keys from the initial state on demand
+        // based on the matching keys from the main collection in the join
+        const subscribeToMatchingChanges = (
+          whereExpression: BasicExpression<boolean> | undefined
+        ) => {
+          let loadedInitialState = false
+          const sentKeys = new Set<string | number>()
+
+          const sendVisibleChanges = (
+            changes: Array<ChangeMessage<any, string | number>>
+          ) => {
+            sendVisibleChangesToPipeline(changes, loadedInitialState, sentKeys)
+          }
+
+          const unsubscribe = collection.subscribeChanges(sendVisibleChanges, {
+            whereExpression,
+          })
+
+          // Create a function that loads keys from the collection
+          // into the query pipeline on demand
+          const filterFn = whereExpression
+            ? createFilterFunctionFromExpression(whereExpression)
+            : () => true
+          const loadKs = (keys: Set<string | number>) => {
+            return loadKeys(keys, sentKeys, filterFn)
+          }
+
+          // Store the functions to load keys and load initial state in the `lazyCollectionsCallbacks` map
+          // This is used by the join operator to dynamically load matching keys from the lazy collection
+          // or to get the full initial state of the collection if there's no index for the join key
+          lazyCollectionsCallbacks[collectionId] = {
+            loadKeys: loadKs,
+            loadInitialState: () => {
+              // Make sure we only load the initial state once
+              if (loadedInitialState) return
+              loadedInitialState = true
+
+              const changes = collection.currentStateAsChanges({
+                whereExpression,
+              })
+              sendChangesToPipeline(changes)
+            },
+          }
+          return unsubscribe
+        }
+
+        const subscribeToOrderedChanges = (
+          whereExpression: BasicExpression<boolean> | undefined
+        ) => {
+          const {
+            offset,
+            limit,
+            comparator,
+            index,
+            dataNeeded,
+            valueExtractorForRawRow,
+          } = optimizableOrderByCollections[collectionId]!
+
+          if (!dataNeeded) {
+            // This should never happen because the topK operator should always set the size callback
+            // which in turn should lead to the orderBy operator setting the dataNeeded callback
+            throw new Error(
+              `Missing dataNeeded callback for collection ${collectionId}`
+            )
+          }
+
+          // This function is called by maybeRunGraph
+          // after each iteration of the query pipeline
+          // to ensure that the orderBy operator has enough data to work with
+          const loadMoreIfNeeded = () => {
+            // `dataNeeded` probes the orderBy operator to see if it needs more data
+            // if it needs more data, it returns the number of items it needs
+            const n = dataNeeded()
+            if (n > 0) {
+              loadNextItems(n)
+            }
+
+            // Indicate that we're done loading data if we didn't need to load more data
+            return n === 0
+          }
+
+          // Keep track of the keys we've sent
+          // and also the biggest value we've sent so far
+          const sentValuesInfo: {
+            sentKeys: Set<string | number>
+            biggest: any
+          } = {
+            sentKeys: new Set<string | number>(),
+            biggest: undefined,
+          }
+
+          const sendChangesToPipelineWithTracking = (
+            changes: Iterable<ChangeMessage<any, string | number>>
+          ) => {
+            const trackedChanges = trackSentValues(
+              changes,
+              comparator,
+              sentValuesInfo
+            )
+            sendChangesToPipeline(trackedChanges, loadMoreIfNeeded)
+          }
+
+          // Loads the next `n` items from the collection
+          // starting from the biggest item it has sent
+          const loadNextItems = (n: number) => {
+            const biggestSentRow = sentValuesInfo.biggest
+            const biggestSentValue = biggestSentRow
+              ? valueExtractorForRawRow(biggestSentRow)
+              : biggestSentRow
+            // Take the `n` items after the biggest sent value
+            const nextOrderedKeys = index.take(n, biggestSentValue)
+            const nextInserts: Array<ChangeMessage<any, string | number>> =
+              nextOrderedKeys.map((key) => {
+                return { type: `insert`, key, value: collection.get(key) }
+              })
+            sendChangesToPipelineWithTracking(nextInserts)
+          }
+
+          // Load the first `offset + limit` values from the index
+          // i.e. the K items from the collection that fall into the requested range: [offset, offset + limit[
+          loadNextItems(offset + limit)
+
+          const sendChangesInRange = (
+            changes: Iterable<ChangeMessage<any, string | number>>
+          ) => {
+            // Split live updates into a delete of the old value and an insert of the new value
+            // and filter out changes that are bigger than the biggest value we've sent so far
+            // because they can't affect the topK
+            const splittedChanges = splitUpdates(changes)
+            const filteredChanges = filterChangesSmallerOrEqualToMax(
+              splittedChanges,
+              comparator,
+              sentValuesInfo.biggest
+            )
+            sendChangesToPipeline(filteredChanges, loadMoreIfNeeded)
+          }
+
+          // Subscribe to changes and only send changes that are smaller than the biggest value we've sent so far
+          // values that are bigger don't need to be sent because they can't affect the topK
+          const unsubscribe = collection.subscribeChanges(sendChangesInRange, {
+            whereExpression,
+          })
+
+          return unsubscribe
+        }
+
+        const subscribeToChanges = (
+          whereExpression?: BasicExpression<boolean>
+        ) => {
+          let unsubscribe: () => void
+          if (lazyCollections.has(collectionId)) {
+            unsubscribe = subscribeToMatchingChanges(whereExpression)
+          } else if (
+            Object.hasOwn(optimizableOrderByCollections, collectionId)
+          ) {
+            unsubscribe = subscribeToOrderedChanges(whereExpression)
+          } else {
+            unsubscribe = subscribeToAllChanges(whereExpression)
+          }
+          unsubscribeCallbacks.add(unsubscribe)
+        }
+
         if (whereClause) {
           // Convert WHERE clause to BasicExpression format for collection subscription
           const whereExpression = convertToBasicExpression(
@@ -328,17 +602,7 @@ export function liveQueryCollectionOptions<
 
           if (whereExpression) {
             // Use index optimization for this collection
-            const subscription = collection.subscribeChanges(
-              (changes) => {
-                sendChangesToInput(input, changes, collection.config.getKey)
-                maybeRunGraph()
-              },
-              {
-                includeInitialState: true,
-                whereExpression: whereExpression,
-              }
-            )
-            unsubscribeCallbacks.add(subscription)
+            subscribeToChanges(whereExpression)
           } else {
             // This should not happen - if we have a whereClause but can't create whereExpression,
             // it indicates a bug in our optimization logic
@@ -349,17 +613,12 @@ export function liveQueryCollectionOptions<
           }
         } else {
           // No WHERE clause for this collection, use regular subscription
-          const subscription = collection.subscribeChanges(
-            (changes) => {
-              sendChangesToInput(input, changes, collection.config.getKey)
-              maybeRunGraph()
-            },
-            { includeInitialState: true }
-          )
-          unsubscribeCallbacks.add(subscription)
+          subscribeToChanges()
         }
       })
 
+      subscribedToAllCollections = true
+
       // Initial run
       maybeRunGraph()
 
@@ -494,7 +753,7 @@ function bridgeToCreateCollection<
  */
 function sendChangesToInput(
   input: RootStreamBuilder<unknown>,
-  changes: Array<ChangeMessage>,
+  changes: Iterable<ChangeMessage>,
   getKey: (item: ChangeMessage[`value`]) => any
 ) {
   const multiSetArray: MultiSetArray<unknown> = []
@@ -593,3 +852,72 @@ function findCollectionAlias(
 
   return undefined
 }
+
+function* trackSentValues(
+  changes: Iterable<ChangeMessage<any, string | number>>,
+  comparator: (a: any, b: any) => number,
+  tracker: { sentKeys: Set<string | number>; biggest: any }
+) {
+  for (const change of changes) {
+    tracker.sentKeys.add(change.key)
+
+    if (!tracker.biggest) {
+      tracker.biggest = change.value
+    } else if (comparator(tracker.biggest, change.value) < 0) {
+      tracker.biggest = change.value
+    }
+
+    yield change
+  }
+}
+
+/** Splits updates into a delete of the old value and an insert of the new value */
+function* splitUpdates<
+  T extends object = Record<string, unknown>,
+  TKey extends string | number = string | number,
+>(
+  changes: Iterable<ChangeMessage<T, TKey>>
+): Generator<ChangeMessage<T, TKey>> {
+  for (const change of changes) {
+    if (change.type === `update`) {
+      yield { type: `delete`, key: change.key, value: change.previousValue! }
+      yield { type: `insert`, key: change.key, value: change.value }
+    } else {
+      yield change
+    }
+  }
+}
+
+function* filterChanges<
+  T extends object = Record<string, unknown>,
+  TKey extends string | number = string | number,
+>(
+  changes: Iterable<ChangeMessage<T, TKey>>,
+  f: (change: ChangeMessage<T, TKey>) => boolean
+): Generator<ChangeMessage<T, TKey>> {
+  for (const change of changes) {
+    if (f(change)) {
+      yield change
+    }
+  }
+}
+
+/**
+ * Filters changes to only include those that are smaller than the provided max value
+ * @param changes - Iterable of changes to filter
+ * @param comparator - Comparator function to use for filtering
+ * @param maxValue - Range to filter changes within (range boundaries are exclusive)
+ * @returns Iterable of changes that fall within the range
+ */
+function* filterChangesSmallerOrEqualToMax<
+  T extends object = Record<string, unknown>,
+  TKey extends string | number = string | number,
+>(
+  changes: Iterable<ChangeMessage<T, TKey>>,
+  comparator: (a: any, b: any) => number,
+  maxValue: any
+): Generator<ChangeMessage<T, TKey>> {
+  yield* filterChanges(changes, (change) => {
+    return !maxValue || comparator(change.value, maxValue) <= 0
+  })
+}

package/src/types.ts

@@ -207,6 +207,7 @@ export interface SyncConfig<
     write: (message: Omit<ChangeMessage<T>, `key`>) => void
     commit: () => void
     markReady: () => void
+    truncate: () => void
   }) => void
 
   /**

package/src/utils/btree.ts

@@ -254,6 +254,14 @@ export class BTree<K = any, V = any> {
     )
   }
 
+  /** Returns the next key larger than the specified key, or undefined if there is none.
+   *  Also, nextHigherKey(undefined) returns the lowest key.
+   */
+  nextHigherKey(key: K | undefined): K | undefined {
+    const p = this.nextHigherPair(key, ReusedArray as [K, V])
+    return p && p[0]
+  }
+
   /** Returns the next pair whose key is smaller than the specified key (or undefined if there is none).
    *  If key === undefined, this function returns the highest pair.
    * @param key The key to search for.
@@ -268,6 +276,14 @@ export class BTree<K = any, V = any> {
     return this._root.getPairOrNextLower(key, this._compare, false, reusedArray)
   }
 
+  /** Returns the next key smaller than the specified key, or undefined if there is none.
+   *  Also, nextLowerKey(undefined) returns the highest key.
+   */
+  nextLowerKey(key: K | undefined): K | undefined {
+    const p = this.nextLowerPair(key, ReusedArray as [K, V])
+    return p && p[0]
+  }
+
   /** Adds all pairs from a list of key-value pairs.
    * @param pairs Pairs to add to this tree. If there are duplicate keys,
    *        later pairs currently overwrite earlier ones (e.g. [[0,1],[0,7]]
@@ -1001,6 +1017,7 @@ const EmptyLeaf = (function () {
   n.isShared = true
   return n
 })()
+const ReusedArray: Array<any> = [] // assumed thread-local
 
 function check(fact: boolean, ...args: Array<any>) {
   if (!fact) {