optimise key loading into query graph (#526)

Author: samwillis

.changeset/afraid-coats-type.md

@@ -0,0 +1,6 @@
+---
+"@tanstack/db-ivm": patch
+"@tanstack/db": patch
+---
+
+optimise key loading into query graph

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

@@ -5,36 +5,34 @@ import type { DifferenceStreamReader } from "../graph.js"
 import type { MultiSet } from "../multiset.js"
 
 /**
- * Operator that applies a function to each element in the input stream
+ * Operator that applies a function to each multi-set in the input stream
  */
 export class TapOperator<T> extends LinearUnaryOperator<T, T> {
-  #f: (data: T) => void
+  #f: (data: MultiSet<T>) => void
 
   constructor(
     id: number,
     inputA: DifferenceStreamReader<T>,
     output: DifferenceStreamWriter<T>,
-    f: (data: T) => void
+    f: (data: MultiSet<T>) => void
   ) {
     super(id, inputA, output)
     this.#f = f
   }
 
   inner(collection: MultiSet<T>): MultiSet<T> {
-    return collection.map((data) => {
-      this.#f(data)
-      return data
-    })
+    this.#f(collection)
+    return collection
   }
 }
 
 /**
- * Invokes a function for each element in the input stream.
+ * Invokes a function for each multi-set in the input stream.
  * This operator doesn't modify the stream and is used to perform side effects.
- * @param f - The function to invoke on each element
+ * @param f - The function to invoke on each multi-set
  * @returns The input stream
  */
-export function tap<T>(f: (data: T) => void): PipedOperator<T, T> {
+export function tap<T>(f: (data: MultiSet<T>) => void): PipedOperator<T, T> {
   return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
     const output = new StreamBuilder<T>(
       stream.graph,

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

@@ -243,7 +243,7 @@ function processJoin(
       const activePipelineWithLoading: IStreamBuilder<
         [key: unknown, [originalKey: string, namespacedRow: NamespacedRow]]
       > = activePipeline.pipe(
-        tap(([joinKey, _]) => {
+        tap((data) => {
           if (deoptimized) {
             return
           }
@@ -270,10 +270,11 @@ function processJoin(
 
           const { loadKeys, loadInitialState } = collectionCallbacks
 
-          if (index && index.supports(`eq`)) {
+          if (index && index.supports(`in`)) {
             // Use the index to fetch the PKs of the rows in the lazy collection
             // that match this row from the active collection based on the value of the joinKey
-            const matchingKeys = index.lookup(`eq`, joinKey)
+            const joinKeys = data.getInner().map(([[joinKey]]) => joinKey)
+            const matchingKeys = index.lookup(`in`, joinKeys)
             // Inform the lazy collection that those keys need to be loaded
             loadKeys(matchingKeys)
           } else {

packages/db/src/query/live/collection-subscriber.ts

@@ -138,16 +138,20 @@ export class CollectionSubscriber<
     keys: Iterable<string | number>,
     filterFn: (item: object) => boolean
   ) {
+    const changes: Array<ChangeMessage<any, string | number>> = []
     for (const key of keys) {
       // Only load the key once
       if (this.sentKeys.has(key)) continue
 
       const value = this.collection.get(key)
       if (value !== undefined && filterFn(value)) {
         this.sentKeys.add(key)
-        this.sendChangesToPipeline([{ type: `insert`, key, value }])
+        changes.push({ type: `insert`, key, value })
       }
     }
+    if (changes.length > 0) {
+      this.sendChangesToPipeline(changes)
+    }
   }
 
   private subscribeToAllChanges(

packages/db/tests/collection-auto-index.test.ts

@@ -532,9 +532,9 @@ describe(`Collection Auto-Indexing`, () => {
     expect(tracker.stats.queriesExecuted).toEqual([
       {
         type: `index`,
-        operation: `eq`,
+        operation: `in`,
         field: `id2`,
-        value: `other2`,
+        value: [`other2`],
       },
     ])
 
@@ -651,9 +651,9 @@ describe(`Collection Auto-Indexing`, () => {
     expect(tracker.stats.queriesExecuted).toEqual([
       {
         type: `index`,
-        operation: `eq`,
+        operation: `in`,
         field: `id2`,
-        value: `other2`,
+        value: [`other2`],
       },
     ])
 

packages/db/tests/query/indexes.test.ts

@@ -780,9 +780,9 @@ describe(`Query Index Optimization`, () => {
         expect(tracker1.stats.queriesExecuted).toEqual([
           {
             type: `index`,
-            operation: `eq`,
+            operation: `in`,
             field: `id`,
-            value: `1`,
+            value: [`1`],
           },
         ])
       } finally {
@@ -983,28 +983,16 @@ describe(`Query Index Optimization`, () => {
         expect(tracker2.stats.queriesExecuted).toEqual([
           {
             type: `index`,
-            operation: `eq`,
-            field: `id2`,
-            value: `1`,
-          },
-          {
-            type: `index`,
-            operation: `eq`,
-            field: `id2`,
-            value: `3`,
-          },
-          {
-            type: `index`,
-            operation: `eq`,
+            operation: `in`,
             field: `id2`,
-            value: `5`,
+            value: [`1`, `3`, `5`],
           },
         ])
 
         expectIndexUsage(combinedStats, {
           shouldUseIndex: true,
           shouldUseFullScan: false,
-          indexCallCount: 4,
+          indexCallCount: 2,
           fullScanCallCount: 0,
         })
       } finally {
@@ -1187,9 +1175,9 @@ describe(`Query Index Optimization`, () => {
         expect(tracker1.stats.queriesExecuted).toEqual([
           {
             type: `index`,
-            operation: `eq`,
+            operation: `in`,
             field: `id`,
-            value: `1`,
+            value: [`1`],
           },
         ])
       } finally {