Deduplicate loadSubset requests in Query collection

Author: kevin-dp

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

@@ -1,17 +1,12 @@
 import { QueryObserver, hashKey } from "@tanstack/query-core"
+import { DeduplicatedLoadSubset } from "@tanstack/db"
 import {
   GetKeyRequiredError,
   QueryClientRequiredError,
   QueryFnRequiredError,
   QueryKeyRequiredError,
 } from "./errors"
 import { createWriteUtils } from "./manual-sync"
-import type {
-  QueryClient,
-  QueryFunctionContext,
-  QueryKey,
-  QueryObserverOptions,
-} from "@tanstack/query-core"
 import type {
   BaseCollectionConfig,
   ChangeMessage,
@@ -23,6 +18,12 @@ import type {
   UpdateMutationFnParams,
   UtilsRecord,
 } from "@tanstack/db"
+import type {
+  QueryClient,
+  QueryFunctionContext,
+  QueryKey,
+  QueryObserverOptions,
+} from "@tanstack/query-core"
 import type { StandardSchemaV1 } from "@standard-schema/spec"
 
 // Re-export for external use
@@ -476,7 +477,8 @@ export function queryCollectionOptions(
     let syncStarted = false
 
     const createQueryFromOpts = (
-      opts: LoadSubsetOptions
+      opts: LoadSubsetOptions,
+      queryFunction: typeof queryFn = queryFn
     ): true | Promise<void> => {
       // Push the predicates down to the queryKey and queryFn
       const key = typeof queryKey === `function` ? queryKey(opts) : queryKey
@@ -519,7 +521,7 @@ export function queryCollectionOptions(
         any
       > = {
         queryKey: key,
-        queryFn: queryFn,
+        queryFn: queryFunction,
         meta: extendedMeta,
         enabled: enabled,
         refetchInterval: refetchInterval,
@@ -667,6 +669,21 @@ export function queryCollectionOptions(
       return handleQueryResult
     }
 
+    // This function is called when a loadSubset call is deduplicated
+    // meaning that we have all the data locally available to answer the query
+    // so we execute the query locally
+    const createLocalQuery = (opts: LoadSubsetOptions) => {
+      const queryFn = ({ meta }: QueryFunctionContext<any>) => {
+        const inserts = collection.currentStateAsChanges(
+          meta!.loadSubsetOptions as LoadSubsetOptions
+        )!
+        const data = inserts.map(({ value }) => value)
+        return Promise.resolve(data)
+      }
+
+      createQueryFromOpts(opts, queryFn)
+    }
+
     const isSubscribed = (hashedQueryKey: string) => {
       return unsubscribes.has(hashedQueryKey)
     }
@@ -796,8 +813,20 @@ export function queryCollectionOptions(
       )
     }
 
+    // Create deduplicated loadSubset wrapper for non-eager modes
+    // This prevents redundant snapshot requests when multiple concurrent
+    // live queries request overlapping or subset predicates
+    const loadSubsetDedupe =
+      syncMode === `eager`
+        ? undefined
+        : new DeduplicatedLoadSubset(createQueryFromOpts, createLocalQuery)
+
+    // TODO: run the tests, probably some will fail bc different requests are made now
+    //       so fix the test expectations
+    //       then also add all the new dedup tests that Sam also added to the Electric collection
+
     return {
-      loadSubset: syncMode === `eager` ? undefined : createQueryFromOpts,
+      loadSubset: loadSubsetDedupe?.loadSubset,
       cleanup,
     }
   }

packages/query-db-collection/tests/query.test.ts

@@ -2,6 +2,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
 import { QueryClient } from "@tanstack/query-core"
 import { createCollection } from "@tanstack/db"
 import { queryCollectionOptions } from "../src/query"
+import { Func, PropRef, Value } from "../../db/src/query/ir"
 import type { QueryFunctionContext } from "@tanstack/query-core"
 import type {
   CollectionImpl,
@@ -18,6 +19,12 @@ interface TestItem {
   value?: number
 }
 
+interface CategorisedItem {
+  id: string
+  name: string
+  category: string
+}
+
 const getKey = (item: TestItem) => item.id
 
 // Helper to advance timers and allow microtasks to flush
@@ -2462,10 +2469,10 @@ describe(`QueryCollection`, () => {
       const queryFn = vi.fn().mockImplementation((context) => {
         const { meta } = context
         const loadSubsetOptions = meta?.loadSubsetOptions ?? {}
-        const { where, orderBy, limit } = loadSubsetOptions
+        const { where } = loadSubsetOptions
 
-        // Query 1: items 1, 2, 3 (no predicates)
-        if (!where && !orderBy && !limit) {
+        // Query 1: items 1, 2, 3 (where: { category: 'A' })
+        if (where?.category === `A`) {
           return Promise.resolve([
             { id: `1`, name: `Item 1` },
             { id: `2`, name: `Item 2` },
@@ -2517,7 +2524,7 @@ describe(`QueryCollection`, () => {
       expect(collection.size).toBe(0)
 
       // Load query 1 with no predicates (items 1, 2, 3)
-      await collection._sync.loadSubset({})
+      await collection._sync.loadSubset({ where: { category: `A` } } as any)
 
       // Wait for query 1 data to load
       await vi.waitFor(() => {
@@ -2551,7 +2558,10 @@ describe(`QueryCollection`, () => {
 
       // GC query 1 (no predicates) - should only remove item 1 (unique to query 1)
       // Items 2 and 3 should remain because they're shared with other queries
-      queryClient.removeQueries({ queryKey: queryKey({}), exact: true })
+      queryClient.removeQueries({
+        queryKey: queryKey({ where: { category: `A` } }),
+        exact: true,
+      })
 
       await vi.waitFor(() => {
         expect(collection.size).toBe(4) // Should have items 2, 3, 4, 5
@@ -2608,13 +2618,13 @@ describe(`QueryCollection`, () => {
       const queryFn = vi.fn().mockImplementation(() => {
         // All queries return the same data regardless of predicates
         return Promise.resolve([
-          { id: `1`, name: `Item 1` },
-          { id: `2`, name: `Item 2` },
-          { id: `3`, name: `Item 3` },
+          { id: `1`, name: `Item 1`, category: `A` },
+          { id: `2`, name: `Item 2`, category: `A` },
+          { id: `3`, name: `Item 3`, category: `A` },
         ])
       })
 
-      const config: QueryCollectionConfig<TestItem> = {
+      const config: QueryCollectionConfig<CategorisedItem> = {
         id: `identical-test`,
         queryClient,
         queryKey: (ctx) => {
@@ -2644,15 +2654,29 @@ describe(`QueryCollection`, () => {
       })
 
       // Add query 2 with different predicates (but returns same data)
-      await collection._sync.loadSubset({ where: { category: `A` } } as any)
+      const whereClause1 = new Func(`eq`, [
+        new PropRef([`category`]),
+        new Value(`A`),
+      ])
+
+      await collection._sync.loadSubset({
+        where: whereClause1,
+      })
 
       // Wait for query 2 data to load
       await vi.waitFor(() => {
         expect(collection.size).toBe(3) // Same data, no new items
       })
 
       // Add query 3 with different predicates (but returns same data)
-      await collection._sync.loadSubset({ where: { category: `B` } } as any)
+      const whereClause2 = new Func(`or`, [
+        new Func(`eq`, [new PropRef([`category`]), new Value(`A`)]),
+        new Func(`eq`, [new PropRef([`category`]), new Value(`B`)]),
+      ])
+
+      await collection._sync.loadSubset({
+        where: whereClause2,
+      })
 
       // Wait for query 3 data to load
       await vi.waitFor(() => {
@@ -2676,7 +2700,7 @@ describe(`QueryCollection`, () => {
 
       // GC query 2 - should still not remove any items (all items are shared with query 3)
       queryClient.removeQueries({
-        queryKey: (config.queryKey as any)({ where: { category: `A` } }),
+        queryKey: (config.queryKey as any)({ where: whereClause1 }),
         exact: true,
       })
 
@@ -2691,7 +2715,7 @@ describe(`QueryCollection`, () => {
 
       // GC query 3 - should remove all items (no more queries reference them)
       queryClient.removeQueries({
-        queryKey: (config.queryKey as any)({ where: { category: `B` } }),
+        queryKey: (config.queryKey as any)({ where: whereClause2 }),
         exact: true,
       })
 
@@ -2712,18 +2736,13 @@ describe(`QueryCollection`, () => {
       const queryFn = vi.fn().mockImplementation((context) => {
         const { meta } = context
         const loadSubsetOptions = meta?.loadSubsetOptions || {}
-        const { where, orderBy, limit } = loadSubsetOptions
+        const { where } = loadSubsetOptions
 
-        // Query 1: empty array (no predicates)
-        if (!where && !orderBy && !limit) {
-          return Promise.resolve([])
-        }
-
-        // Query 2: some items (where: { category: 'A' })
-        if (where?.category === `A`) {
+        // Query 2: some items (where: { category: 'B' })
+        if (where?.name === `eq` && where?.args[1].value === `B`) {
           return Promise.resolve([
-            { id: `1`, name: `Item 1` },
-            { id: `2`, name: `Item 2` },
+            { id: `1`, name: `Item 1`, category: `B` },
+            { id: `2`, name: `Item 2`, category: `B` },
           ])
         }
 
@@ -2752,15 +2771,23 @@ describe(`QueryCollection`, () => {
       expect(collection.size).toBe(0)
 
       // Load query 1 with no predicates (returns empty array)
-      await collection._sync.loadSubset({})
+      const whereClause1 = new Func(`eq`, [
+        new PropRef([`category`]),
+        new Value(`A`),
+      ])
+      await collection._sync.loadSubset({ where: whereClause1 })
 
       // Wait for query 1 data to load (still empty)
       await vi.waitFor(() => {
         expect(collection.size).toBe(0) // Empty query
       })
 
       // Add query 2 with different predicates (items 1, 2)
-      await collection._sync.loadSubset({ where: { category: `A` } } as any)
+      const whereClause2 = new Func(`eq`, [
+        new PropRef([`category`]),
+        new Value(`B`),
+      ])
+      await collection._sync.loadSubset({ where: whereClause2 } as any)
 
       // Wait for query 2 data to load
       await vi.waitFor(() => {
@@ -2786,7 +2813,7 @@ describe(`QueryCollection`, () => {
 
       // GC non-empty query 2 - should remove its items
       queryClient.removeQueries({
-        queryKey: (config.queryKey as any)({ where: { category: `A` } }),
+        queryKey: (config.queryKey as any)({ where: whereClause2 }),
         exact: true,
       })
 
@@ -2805,29 +2832,29 @@ describe(`QueryCollection`, () => {
       const queryFn = vi.fn().mockImplementation((context) => {
         const { meta } = context
         const loadSubsetOptions = meta?.loadSubsetOptions || {}
-        const { where, orderBy, limit } = loadSubsetOptions
+        const { where } = loadSubsetOptions
 
         // Query 1: items 1, 2 (no predicates)
-        if (!where && !orderBy && !limit) {
+        if (where?.name === `eq` && where?.args[1].value === `C`) {
           return Promise.resolve([
-            { id: `1`, name: `Item 1` },
-            { id: `2`, name: `Item 2` },
+            { id: `1`, name: `Item 1`, type: `C` },
+            { id: `2`, name: `Item 2`, type: `C` },
           ])
         }
 
         // Query 2: items 2, 3 (where: { type: 'A' })
-        if (where?.type === `A`) {
+        if (where?.name === `eq` && where?.args[1].value === `A`) {
           return Promise.resolve([
-            { id: `2`, name: `Item 2` },
-            { id: `3`, name: `Item 3` },
+            { id: `2`, name: `Item 2`, type: `A` },
+            { id: `3`, name: `Item 3`, type: `A` },
           ])
         }
 
         // Query 3: items 3, 4 (where: { type: 'B' })
-        if (where?.type === `B`) {
+        if (where?.name === `eq` && where?.args[1].value === `B`) {
           return Promise.resolve([
-            { id: `3`, name: `Item 3` },
-            { id: `4`, name: `Item 4` },
+            { id: `3`, name: `Item 3`, type: `B` },
+            { id: `4`, name: `Item 4`, type: `B` },
           ])
         }
 
@@ -2856,23 +2883,35 @@ describe(`QueryCollection`, () => {
       expect(collection.size).toBe(0)
 
       // Load query 1 with no predicates (items 1, 2)
-      await collection._sync.loadSubset({})
+      const whereClause1 = new Func(`eq`, [
+        new PropRef([`type`]),
+        new Value(`C`),
+      ])
+      await collection._sync.loadSubset({ where: whereClause1 })
 
       // Wait for query 1 data to load
       await vi.waitFor(() => {
         expect(collection.size).toBe(2)
       })
 
       // Add query 2 with different predicates (items 2, 3)
-      await collection._sync.loadSubset({ where: { type: `A` } } as any)
+      const whereClause2 = new Func(`eq`, [
+        new PropRef([`type`]),
+        new Value(`A`),
+      ])
+      await collection._sync.loadSubset({ where: whereClause2 })
 
       // Wait for query 2 data to load
       await vi.waitFor(() => {
         expect(collection.size).toBe(3) // Should have items 1, 2, 3
       })
 
       // Add query 3 with different predicates
-      await collection._sync.loadSubset({ where: { type: `B` } } as any)
+      const whereClause3 = new Func(`eq`, [
+        new PropRef([`type`]),
+        new Value(`B`),
+      ])
+      await collection._sync.loadSubset({ where: whereClause3 } as any)
 
       // Wait for query 3 data to load
       await vi.waitFor(() => {
@@ -2881,15 +2920,15 @@ describe(`QueryCollection`, () => {
 
       // GC all queries concurrently
       queryClient.removeQueries({
-        queryKey: (config.queryKey as any)({}),
+        queryKey: (config.queryKey as any)({ where: whereClause1 }),
         exact: true,
       })
       queryClient.removeQueries({
-        queryKey: (config.queryKey as any)({ where: { type: `A` } }),
+        queryKey: (config.queryKey as any)({ where: whereClause2 }),
         exact: true,
       })
       queryClient.removeQueries({
-        queryKey: (config.queryKey as any)({ where: { type: `B` } }),
+        queryKey: (config.queryKey as any)({ where: whereClause3 }),
         exact: true,
       })