@tanstack/db 0.3.1 → 0.4.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (162) hide show
  1. package/dist/cjs/{change-events.cjs → collection/change-events.cjs} +13 -42
  2. package/dist/cjs/collection/change-events.cjs.map +1 -0
  3. package/dist/{esm/change-events.d.ts → cjs/collection/change-events.d.cts} +6 -6
  4. package/dist/cjs/collection/changes.cjs +108 -0
  5. package/dist/cjs/collection/changes.cjs.map +1 -0
  6. package/dist/cjs/collection/changes.d.cts +53 -0
  7. package/dist/cjs/collection/events.cjs +90 -0
  8. package/dist/cjs/collection/events.cjs.map +1 -0
  9. package/dist/cjs/collection/events.d.cts +53 -0
  10. package/dist/cjs/collection/index.cjs +417 -0
  11. package/dist/cjs/collection/index.cjs.map +1 -0
  12. package/dist/{esm/collection.d.ts → cjs/collection/index.d.cts} +56 -172
  13. package/dist/cjs/collection/indexes.cjs +124 -0
  14. package/dist/cjs/collection/indexes.cjs.map +1 -0
  15. package/dist/cjs/collection/indexes.d.cts +47 -0
  16. package/dist/cjs/collection/lifecycle.cjs +150 -0
  17. package/dist/cjs/collection/lifecycle.cjs.map +1 -0
  18. package/dist/cjs/collection/lifecycle.d.cts +70 -0
  19. package/dist/cjs/collection/mutations.cjs +315 -0
  20. package/dist/cjs/collection/mutations.cjs.map +1 -0
  21. package/dist/cjs/collection/mutations.d.cts +33 -0
  22. package/dist/cjs/collection/state.cjs +597 -0
  23. package/dist/cjs/collection/state.cjs.map +1 -0
  24. package/dist/cjs/collection/state.d.cts +122 -0
  25. package/dist/cjs/collection/subscription.cjs +160 -0
  26. package/dist/cjs/collection/subscription.cjs.map +1 -0
  27. package/dist/cjs/collection/subscription.d.cts +57 -0
  28. package/dist/cjs/collection/sync.cjs +154 -0
  29. package/dist/cjs/collection/sync.cjs.map +1 -0
  30. package/dist/cjs/collection/sync.d.cts +34 -0
  31. package/dist/cjs/index.cjs +8 -8
  32. package/dist/cjs/index.d.cts +2 -2
  33. package/dist/cjs/indexes/auto-index.cjs.map +1 -1
  34. package/dist/cjs/indexes/auto-index.d.cts +1 -1
  35. package/dist/cjs/indexes/base-index.cjs.map +1 -1
  36. package/dist/cjs/indexes/base-index.d.cts +2 -2
  37. package/dist/cjs/indexes/btree-index.cjs +2 -2
  38. package/dist/cjs/indexes/btree-index.cjs.map +1 -1
  39. package/dist/cjs/indexes/btree-index.d.cts +1 -1
  40. package/dist/cjs/query/builder/index.cjs +2 -2
  41. package/dist/cjs/query/builder/index.cjs.map +1 -1
  42. package/dist/cjs/query/builder/types.d.cts +1 -1
  43. package/dist/cjs/query/compiler/index.cjs +5 -2
  44. package/dist/cjs/query/compiler/index.cjs.map +1 -1
  45. package/dist/cjs/query/compiler/index.d.cts +3 -2
  46. package/dist/cjs/query/compiler/joins.cjs +22 -24
  47. package/dist/cjs/query/compiler/joins.cjs.map +1 -1
  48. package/dist/cjs/query/compiler/joins.d.cts +3 -2
  49. package/dist/cjs/query/compiler/order-by.cjs.map +1 -1
  50. package/dist/cjs/query/compiler/order-by.d.cts +1 -1
  51. package/dist/cjs/query/ir.cjs.map +1 -1
  52. package/dist/cjs/query/ir.d.cts +1 -1
  53. package/dist/cjs/query/live/collection-config-builder.cjs +29 -12
  54. package/dist/cjs/query/live/collection-config-builder.cjs.map +1 -1
  55. package/dist/cjs/query/live/collection-config-builder.d.cts +3 -0
  56. package/dist/cjs/query/live/collection-subscriber.cjs +43 -104
  57. package/dist/cjs/query/live/collection-subscriber.cjs.map +1 -1
  58. package/dist/cjs/query/live/collection-subscriber.d.cts +4 -7
  59. package/dist/cjs/query/live-query-collection.cjs +2 -2
  60. package/dist/cjs/query/live-query-collection.cjs.map +1 -1
  61. package/dist/cjs/query/live-query-collection.d.cts +1 -1
  62. package/dist/cjs/transactions.cjs +3 -3
  63. package/dist/cjs/transactions.cjs.map +1 -1
  64. package/dist/cjs/types.d.cts +12 -10
  65. package/dist/{cjs/change-events.d.cts → esm/collection/change-events.d.ts} +6 -6
  66. package/dist/esm/{change-events.js → collection/change-events.js} +13 -42
  67. package/dist/esm/collection/change-events.js.map +1 -0
  68. package/dist/esm/collection/changes.d.ts +53 -0
  69. package/dist/esm/collection/changes.js +108 -0
  70. package/dist/esm/collection/changes.js.map +1 -0
  71. package/dist/esm/collection/events.d.ts +53 -0
  72. package/dist/esm/collection/events.js +90 -0
  73. package/dist/esm/collection/events.js.map +1 -0
  74. package/dist/{cjs/collection.d.cts → esm/collection/index.d.ts} +56 -172
  75. package/dist/esm/collection/index.js +417 -0
  76. package/dist/esm/collection/index.js.map +1 -0
  77. package/dist/esm/collection/indexes.d.ts +47 -0
  78. package/dist/esm/collection/indexes.js +124 -0
  79. package/dist/esm/collection/indexes.js.map +1 -0
  80. package/dist/esm/collection/lifecycle.d.ts +70 -0
  81. package/dist/esm/collection/lifecycle.js +150 -0
  82. package/dist/esm/collection/lifecycle.js.map +1 -0
  83. package/dist/esm/collection/mutations.d.ts +33 -0
  84. package/dist/esm/collection/mutations.js +315 -0
  85. package/dist/esm/collection/mutations.js.map +1 -0
  86. package/dist/esm/collection/state.d.ts +122 -0
  87. package/dist/esm/collection/state.js +597 -0
  88. package/dist/esm/collection/state.js.map +1 -0
  89. package/dist/esm/collection/subscription.d.ts +57 -0
  90. package/dist/esm/collection/subscription.js +160 -0
  91. package/dist/esm/collection/subscription.js.map +1 -0
  92. package/dist/esm/collection/sync.d.ts +34 -0
  93. package/dist/esm/collection/sync.js +154 -0
  94. package/dist/esm/collection/sync.js.map +1 -0
  95. package/dist/esm/index.d.ts +2 -2
  96. package/dist/esm/index.js +1 -1
  97. package/dist/esm/indexes/auto-index.d.ts +1 -1
  98. package/dist/esm/indexes/auto-index.js.map +1 -1
  99. package/dist/esm/indexes/base-index.d.ts +2 -2
  100. package/dist/esm/indexes/base-index.js.map +1 -1
  101. package/dist/esm/indexes/btree-index.d.ts +1 -1
  102. package/dist/esm/indexes/btree-index.js +2 -2
  103. package/dist/esm/indexes/btree-index.js.map +1 -1
  104. package/dist/esm/proxy.js +1 -1
  105. package/dist/esm/query/builder/index.js +1 -1
  106. package/dist/esm/query/builder/index.js.map +1 -1
  107. package/dist/esm/query/builder/types.d.ts +1 -1
  108. package/dist/esm/query/compiler/index.d.ts +3 -2
  109. package/dist/esm/query/compiler/index.js +5 -2
  110. package/dist/esm/query/compiler/index.js.map +1 -1
  111. package/dist/esm/query/compiler/joins.d.ts +3 -2
  112. package/dist/esm/query/compiler/joins.js +22 -24
  113. package/dist/esm/query/compiler/joins.js.map +1 -1
  114. package/dist/esm/query/compiler/order-by.d.ts +1 -1
  115. package/dist/esm/query/compiler/order-by.js.map +1 -1
  116. package/dist/esm/query/ir.d.ts +1 -1
  117. package/dist/esm/query/ir.js.map +1 -1
  118. package/dist/esm/query/live/collection-config-builder.d.ts +3 -0
  119. package/dist/esm/query/live/collection-config-builder.js +29 -12
  120. package/dist/esm/query/live/collection-config-builder.js.map +1 -1
  121. package/dist/esm/query/live/collection-subscriber.d.ts +4 -7
  122. package/dist/esm/query/live/collection-subscriber.js +43 -104
  123. package/dist/esm/query/live/collection-subscriber.js.map +1 -1
  124. package/dist/esm/query/live-query-collection.d.ts +1 -1
  125. package/dist/esm/query/live-query-collection.js +1 -1
  126. package/dist/esm/query/live-query-collection.js.map +1 -1
  127. package/dist/esm/transactions.js +3 -3
  128. package/dist/esm/transactions.js.map +1 -1
  129. package/dist/esm/types.d.ts +12 -10
  130. package/package.json +2 -2
  131. package/src/{change-events.ts → collection/change-events.ts} +25 -39
  132. package/src/collection/changes.ts +163 -0
  133. package/src/collection/events.ts +171 -0
  134. package/src/collection/index.ts +808 -0
  135. package/src/collection/indexes.ts +172 -0
  136. package/src/collection/lifecycle.ts +221 -0
  137. package/src/collection/mutations.ts +535 -0
  138. package/src/collection/state.ts +866 -0
  139. package/src/collection/subscription.ts +239 -0
  140. package/src/collection/sync.ts +235 -0
  141. package/src/index.ts +2 -2
  142. package/src/indexes/auto-index.ts +1 -1
  143. package/src/indexes/base-index.ts +3 -3
  144. package/src/indexes/btree-index.ts +2 -2
  145. package/src/query/builder/index.ts +1 -1
  146. package/src/query/builder/types.ts +1 -1
  147. package/src/query/compiler/index.ts +7 -1
  148. package/src/query/compiler/joins.ts +28 -41
  149. package/src/query/compiler/order-by.ts +1 -1
  150. package/src/query/ir.ts +1 -1
  151. package/src/query/live/collection-config-builder.ts +48 -22
  152. package/src/query/live/collection-subscriber.ts +63 -168
  153. package/src/query/live-query-collection.ts +2 -2
  154. package/src/transactions.ts +3 -3
  155. package/src/types.ts +14 -15
  156. package/dist/cjs/change-events.cjs.map +0 -1
  157. package/dist/cjs/collection.cjs +0 -1580
  158. package/dist/cjs/collection.cjs.map +0 -1
  159. package/dist/esm/change-events.js.map +0 -1
  160. package/dist/esm/collection.js +0 -1580
  161. package/dist/esm/collection.js.map +0 -1
  162. package/src/collection.ts +0 -2488
@@ -16,8 +16,9 @@ import {
16
16
  UnsupportedJoinSourceTypeError,
17
17
  UnsupportedJoinTypeError,
18
18
  } from "../../errors.js"
19
- import { findIndexForField } from "../../utils/index-optimization.js"
20
19
  import { ensureIndexForField } from "../../indexes/auto-index.js"
20
+ import { PropRef } from "../ir.js"
21
+ import { inArray } from "../builder/functions.js"
21
22
  import { compileExpression } from "./evaluators.js"
22
23
  import { compileQuery, followRef } from "./index.js"
23
24
  import type { OrderByOptimizationInfo } from "./order-by.js"
@@ -25,19 +26,18 @@ import type {
25
26
  BasicExpression,
26
27
  CollectionRef,
27
28
  JoinClause,
28
- PropRef,
29
29
  QueryIR,
30
30
  QueryRef,
31
31
  } from "../ir.js"
32
32
  import type { IStreamBuilder, JoinType } from "@tanstack/db-ivm"
33
- import type { Collection } from "../../collection.js"
33
+ import type { Collection } from "../../collection/index.js"
34
34
  import type {
35
35
  KeyedStream,
36
36
  NamespacedAndKeyedStream,
37
37
  NamespacedRow,
38
38
  } from "../../types.js"
39
39
  import type { QueryCache, QueryMapping } from "./types.js"
40
- import type { BaseIndex } from "../../indexes/base-index.js"
40
+ import type { CollectionSubscription } from "../../collection/subscription.js"
41
41
 
42
42
  export type LoadKeysFn = (key: Set<string | number>) => void
43
43
  export type LazyCollectionCallbacks = {
@@ -58,6 +58,7 @@ export function processJoins(
58
58
  cache: QueryCache,
59
59
  queryMapping: QueryMapping,
60
60
  collections: Record<string, Collection>,
61
+ subscriptions: Record<string, CollectionSubscription>,
61
62
  callbacks: Record<string, LazyCollectionCallbacks>,
62
63
  lazyCollections: Set<string>,
63
64
  optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
@@ -76,6 +77,7 @@ export function processJoins(
76
77
  cache,
77
78
  queryMapping,
78
79
  collections,
80
+ subscriptions,
79
81
  callbacks,
80
82
  lazyCollections,
81
83
  optimizableOrderByCollections,
@@ -99,6 +101,7 @@ function processJoin(
99
101
  cache: QueryCache,
100
102
  queryMapping: QueryMapping,
101
103
  collections: Record<string, Collection>,
104
+ subscriptions: Record<string, CollectionSubscription>,
102
105
  callbacks: Record<string, LazyCollectionCallbacks>,
103
106
  lazyCollections: Set<string>,
104
107
  optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
@@ -113,6 +116,7 @@ function processJoin(
113
116
  joinClause.from,
114
117
  allInputs,
115
118
  collections,
119
+ subscriptions,
116
120
  callbacks,
117
121
  lazyCollections,
118
122
  optimizableOrderByCollections,
@@ -215,8 +219,6 @@ function processJoin(
215
219
  const activePipeline =
216
220
  activeCollection === `main` ? mainPipeline : joinedPipeline
217
221
 
218
- let index: BaseIndex<string | number> | undefined
219
-
220
222
  const lazyCollectionJoinExpr =
221
223
  activeCollection === `main`
222
224
  ? (joinedExpr as PropRef)
@@ -238,50 +240,33 @@ function processJoin(
238
240
  )
239
241
  }
240
242
 
241
- let deoptimized = false
242
-
243
243
  const activePipelineWithLoading: IStreamBuilder<
244
244
  [key: unknown, [originalKey: string, namespacedRow: NamespacedRow]]
245
245
  > = activePipeline.pipe(
246
246
  tap((data) => {
247
- if (deoptimized) {
248
- return
249
- }
247
+ const lazyCollectionSubscription = subscriptions[lazyCollection.id]
250
248
 
251
- // Find the index for the path we join on
252
- // we need to find the index inside the map operator
253
- // because the indexes are only available after the initial sync
254
- // so we can't fetch it during compilation
255
- index ??= findIndexForField(
256
- followRefCollection.indexes,
257
- followRefResult.path
258
- )
259
-
260
- // The `callbacks` object is passed by the liveQueryCollection to the compiler.
261
- // It contains a function to lazy load keys for each lazy collection
262
- // as well as a function to switch back to a regular collection
263
- // (useful when there's no index for available for lazily loading the collection)
264
- const collectionCallbacks = callbacks[lazyCollection.id]
265
- if (!collectionCallbacks) {
249
+ if (!lazyCollectionSubscription) {
266
250
  throw new Error(
267
- `Internal error: callbacks for collection are missing in join pipeline. Make sure the live query collection sets them before running the pipeline.`
251
+ `Internal error: subscription for collection is missing in join pipeline. Make sure the live query collection sets the subscription before running the pipeline.`
268
252
  )
269
253
  }
270
254
 
271
- const { loadKeys, loadInitialState } = collectionCallbacks
272
-
273
- if (index && index.supports(`in`)) {
274
- // Use the index to fetch the PKs of the rows in the lazy collection
275
- // that match this row from the active collection based on the value of the joinKey
276
- const joinKeys = data.getInner().map(([[joinKey]]) => joinKey)
277
- const matchingKeys = index.lookup(`in`, joinKeys)
278
- // Inform the lazy collection that those keys need to be loaded
279
- loadKeys(matchingKeys)
280
- } else {
281
- // We can't optimize the join because there is no index for the join key
282
- // on the lazy collection, so we load the initial state
283
- deoptimized = true
284
- loadInitialState()
255
+ if (lazyCollectionSubscription.hasLoadedInitialState()) {
256
+ // Entire state was already loaded because we deoptimized the join
257
+ return
258
+ }
259
+
260
+ const joinKeys = data.getInner().map(([[joinKey]]) => joinKey)
261
+ const lazyJoinRef = new PropRef(followRefResult.path)
262
+ const loaded = lazyCollectionSubscription.requestSnapshot({
263
+ where: inArray(lazyJoinRef, joinKeys),
264
+ optimizedOnly: true,
265
+ })
266
+
267
+ if (!loaded) {
268
+ // Snapshot wasn't sent because it could not be loaded from the indexes
269
+ lazyCollectionSubscription.requestSnapshot()
285
270
  }
286
271
  })
287
272
  )
@@ -397,6 +382,7 @@ function processJoinSource(
397
382
  from: CollectionRef | QueryRef,
398
383
  allInputs: Record<string, KeyedStream>,
399
384
  collections: Record<string, Collection>,
385
+ subscriptions: Record<string, CollectionSubscription>,
400
386
  callbacks: Record<string, LazyCollectionCallbacks>,
401
387
  lazyCollections: Set<string>,
402
388
  optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
@@ -420,6 +406,7 @@ function processJoinSource(
420
406
  originalQuery,
421
407
  allInputs,
422
408
  collections,
409
+ subscriptions,
423
410
  callbacks,
424
411
  lazyCollections,
425
412
  optimizableOrderByCollections,
@@ -11,7 +11,7 @@ import type { OrderByClause, QueryIR, Select } from "../ir.js"
11
11
  import type { NamespacedAndKeyedStream, NamespacedRow } from "../../types.js"
12
12
  import type { IStreamBuilder, KeyValue } from "@tanstack/db-ivm"
13
13
  import type { BaseIndex } from "../../indexes/base-index.js"
14
- import type { Collection } from "../../collection.js"
14
+ import type { Collection } from "../../collection/index.js"
15
15
 
16
16
  export type OrderByOptimizationInfo = {
17
17
  offset: number
package/src/query/ir.ts CHANGED
@@ -3,7 +3,7 @@ This is the intermediate representation of the query.
3
3
  */
4
4
 
5
5
  import type { CompareOptions } from "./builder/types"
6
- import type { CollectionImpl } from "../collection"
6
+ import type { CollectionImpl } from "../collection/index.js"
7
7
  import type { NamespacedRow } from "../types"
8
8
 
9
9
  export interface QueryIR {
@@ -2,9 +2,10 @@ import { D2, output } from "@tanstack/db-ivm"
2
2
  import { compileQuery } from "../compiler/index.js"
3
3
  import { buildQuery, getQueryIR } from "../builder/index.js"
4
4
  import { CollectionSubscriber } from "./collection-subscriber.js"
5
+ import type { CollectionSubscription } from "../../collection/subscription.js"
5
6
  import type { RootStreamBuilder } from "@tanstack/db-ivm"
6
7
  import type { OrderByOptimizationInfo } from "../compiler/order-by.js"
7
- import type { Collection } from "../../collection.js"
8
+ import type { Collection } from "../../collection/index.js"
8
9
  import type {
9
10
  CollectionConfig,
10
11
  KeyedStream,
@@ -41,6 +42,8 @@ export class CollectionConfigBuilder<
41
42
 
42
43
  private readonly compare?: (val1: TResult, val2: TResult) => number
43
44
 
45
+ private isGraphRunning = false
46
+
44
47
  private graphCache: D2 | undefined
45
48
  private inputsCache: Record<string, RootStreamBuilder<unknown>> | undefined
46
49
  private pipelineCache: ResultStream | undefined
@@ -48,6 +51,8 @@ export class CollectionConfigBuilder<
48
51
  | Map<string, BasicExpression<boolean>>
49
52
  | undefined
50
53
 
54
+ // Map of collection ID to subscription
55
+ readonly subscriptions: Record<string, CollectionSubscription> = {}
51
56
  // Map of collection IDs to functions that load keys for that lazy collection
52
57
  lazyCollectionsCallbacks: Record<string, LazyCollectionCallbacks> = {}
53
58
  // Set of collection IDs that are lazy collections
@@ -104,25 +109,41 @@ export class CollectionConfigBuilder<
104
109
  syncState: FullSyncState,
105
110
  callback?: () => boolean
106
111
  ) {
107
- const { begin, commit, markReady } = config
112
+ if (this.isGraphRunning) {
113
+ // no nested runs of the graph
114
+ // which is possible if the `callback`
115
+ // would call `maybeRunGraph` e.g. after it has loaded some more data
116
+ return
117
+ }
108
118
 
109
- // We only run the graph if all the collections are ready
110
- if (
111
- this.allCollectionsReadyOrInitialCommit() &&
112
- syncState.subscribedToAllCollections
113
- ) {
114
- syncState.graph.run()
115
- const ready = callback?.() ?? true
116
- // On the initial run, we may need to do an empty commit to ensure that
117
- // the collection is initialized
118
- if (syncState.messagesCount === 0) {
119
- begin()
120
- commit()
121
- }
122
- // Mark the collection as ready after the first successful run
123
- if (ready && this.allCollectionsReady()) {
124
- markReady()
119
+ this.isGraphRunning = true
120
+
121
+ try {
122
+ const { begin, commit, markReady } = config
123
+
124
+ // We only run the graph if all the collections are ready
125
+ if (
126
+ this.allCollectionsReadyOrInitialCommit() &&
127
+ syncState.subscribedToAllCollections
128
+ ) {
129
+ while (syncState.graph.pendingWork()) {
130
+ syncState.graph.run()
131
+ callback?.()
132
+ }
133
+
134
+ // On the initial run, we may need to do an empty commit to ensure that
135
+ // the collection is initialized
136
+ if (syncState.messagesCount === 0) {
137
+ begin()
138
+ commit()
139
+ }
140
+ // Mark the collection as ready after the first successful run
141
+ if (this.allCollectionsReady()) {
142
+ markReady()
143
+ }
125
144
  }
145
+ } finally {
146
+ this.isGraphRunning = false
126
147
  }
127
148
  }
128
149
 
@@ -189,6 +210,7 @@ export class CollectionConfigBuilder<
189
210
  this.query,
190
211
  this.inputsCache as Record<string, KeyedStream>,
191
212
  this.collections,
213
+ this.subscriptions,
192
214
  this.lazyCollectionsCallbacks,
193
215
  this.lazyCollections,
194
216
  this.optimizableOrderByCollections
@@ -320,17 +342,21 @@ export class CollectionConfigBuilder<
320
342
  syncState,
321
343
  this
322
344
  )
323
- collectionSubscriber.subscribe()
324
345
 
325
- const loadMore =
326
- collectionSubscriber.loadMoreIfNeeded.bind(collectionSubscriber)
346
+ const subscription = collectionSubscriber.subscribe()
347
+ this.subscriptions[collectionId] = subscription
348
+
349
+ const loadMore = collectionSubscriber.loadMoreIfNeeded.bind(
350
+ collectionSubscriber,
351
+ subscription
352
+ )
327
353
 
328
354
  return loadMore
329
355
  }
330
356
  )
331
357
 
332
358
  const loadMoreDataCallback = () => {
333
- loaders.map((loader) => loader()) // .every((doneLoading) => doneLoading)
359
+ loaders.map((loader) => loader())
334
360
  return true
335
361
  }
336
362
 
@@ -1,21 +1,18 @@
1
1
  import { MultiSet } from "@tanstack/db-ivm"
2
- import { createFilterFunctionFromExpression } from "../../change-events.js"
3
2
  import { convertToBasicExpression } from "../compiler/expressions.js"
4
3
  import type { FullSyncState } from "./types.js"
5
4
  import type { MultiSetArray, RootStreamBuilder } from "@tanstack/db-ivm"
6
- import type { Collection } from "../../collection.js"
5
+ import type { Collection } from "../../collection/index.js"
7
6
  import type { ChangeMessage, SyncConfig } from "../../types.js"
8
7
  import type { Context, GetResult } from "../builder/types.js"
9
8
  import type { BasicExpression } from "../ir.js"
10
9
  import type { CollectionConfigBuilder } from "./collection-config-builder.js"
10
+ import type { CollectionSubscription } from "../../collection/subscription.js"
11
11
 
12
12
  export class CollectionSubscriber<
13
13
  TContext extends Context,
14
14
  TResult extends object = GetResult<TContext>,
15
15
  > {
16
- // Keep track of the keys we've sent (needed for join and orderBy optimizations)
17
- private sentKeys = new Set<string | number>()
18
-
19
16
  // Keep track of the biggest value we've sent so far (needed for orderBy optimization)
20
17
  private biggest: any = undefined
21
18
 
@@ -27,7 +24,7 @@ export class CollectionSubscriber<
27
24
  private collectionConfigBuilder: CollectionConfigBuilder<TContext, TResult>
28
25
  ) {}
29
26
 
30
- subscribe() {
27
+ subscribe(): CollectionSubscription {
31
28
  const collectionAlias = findCollectionAlias(
32
29
  this.collectionId,
33
30
  this.collectionConfigBuilder.query
@@ -43,7 +40,7 @@ export class CollectionSubscriber<
43
40
 
44
41
  if (whereExpression) {
45
42
  // Use index optimization for this collection
46
- this.subscribeToChanges(whereExpression)
43
+ return this.subscribeToChanges(whereExpression)
47
44
  } else {
48
45
  // This should not happen - if we have a whereClause but can't create whereExpression,
49
46
  // it indicates a bug in our optimization logic
@@ -54,25 +51,34 @@ export class CollectionSubscriber<
54
51
  }
55
52
  } else {
56
53
  // No WHERE clause for this collection, use regular subscription
57
- this.subscribeToChanges()
54
+ return this.subscribeToChanges()
58
55
  }
59
56
  }
60
57
 
61
58
  private subscribeToChanges(whereExpression?: BasicExpression<boolean>) {
62
- let unsubscribe: () => void
63
- if (this.collectionConfigBuilder.lazyCollections.has(this.collectionId)) {
64
- unsubscribe = this.subscribeToMatchingChanges(whereExpression)
65
- } else if (
59
+ let subscription: CollectionSubscription
60
+ if (
66
61
  Object.hasOwn(
67
62
  this.collectionConfigBuilder.optimizableOrderByCollections,
68
63
  this.collectionId
69
64
  )
70
65
  ) {
71
- unsubscribe = this.subscribeToOrderedChanges(whereExpression)
66
+ subscription = this.subscribeToOrderedChanges(whereExpression)
72
67
  } else {
73
- unsubscribe = this.subscribeToAllChanges(whereExpression)
68
+ // If the collection is lazy then we should not include the initial state
69
+ const includeInitialState =
70
+ !this.collectionConfigBuilder.lazyCollections.has(this.collectionId)
71
+
72
+ subscription = this.subscribeToMatchingChanges(
73
+ whereExpression,
74
+ includeInitialState
75
+ )
76
+ }
77
+ const unsubscribe = () => {
78
+ subscription.unsubscribe()
74
79
  }
75
80
  this.syncState.unsubscribeCallbacks.add(unsubscribe)
81
+ return subscription
76
82
  }
77
83
 
78
84
  private sendChangesToPipeline(
@@ -101,153 +107,38 @@ export class CollectionSubscriber<
101
107
  )
102
108
  }
103
109
 
104
- // Wraps the sendChangesToPipeline function
105
- // in order to turn `update`s into `insert`s
106
- // for keys that have not been sent to the pipeline yet
107
- // and filter out deletes for keys that have not been sent
108
- private sendVisibleChangesToPipeline = (
109
- changes: Array<ChangeMessage<any, string | number>>,
110
- loadedInitialState: boolean
111
- ) => {
112
- if (loadedInitialState) {
113
- // There was no index for the join key
114
- // so we loaded the initial state
115
- // so we can safely assume that the pipeline has seen all keys
116
- return this.sendChangesToPipeline(changes)
117
- }
118
-
119
- const newChanges = []
120
- for (const change of changes) {
121
- let newChange = change
122
- if (!this.sentKeys.has(change.key)) {
123
- if (change.type === `update`) {
124
- newChange = { ...change, type: `insert` }
125
- } else if (change.type === `delete`) {
126
- // filter out deletes for keys that have not been sent
127
- continue
128
- }
129
- this.sentKeys.add(change.key)
130
- }
131
- newChanges.push(newChange)
132
- }
133
-
134
- return this.sendChangesToPipeline(newChanges)
135
- }
136
-
137
- private loadKeys(
138
- keys: Iterable<string | number>,
139
- filterFn: (item: object) => boolean
140
- ) {
141
- const changes: Array<ChangeMessage<any, string | number>> = []
142
- for (const key of keys) {
143
- // Only load the key once
144
- if (this.sentKeys.has(key)) continue
145
-
146
- const value = this.collection.get(key)
147
- if (value !== undefined && filterFn(value)) {
148
- this.sentKeys.add(key)
149
- changes.push({ type: `insert`, key, value })
150
- }
151
- }
152
- if (changes.length > 0) {
153
- this.sendChangesToPipeline(changes)
154
- }
155
- }
156
-
157
- private subscribeToAllChanges(
158
- whereExpression: BasicExpression<boolean> | undefined
159
- ) {
160
- const sendChangesToPipeline = this.sendChangesToPipeline.bind(this)
161
- const unsubscribe = this.collection.subscribeChanges(
162
- sendChangesToPipeline,
163
- {
164
- includeInitialState: true,
165
- ...(whereExpression ? { whereExpression } : undefined),
166
- }
167
- )
168
- return unsubscribe
169
- }
170
-
171
110
  private subscribeToMatchingChanges(
172
- whereExpression: BasicExpression<boolean> | undefined
111
+ whereExpression: BasicExpression<boolean> | undefined,
112
+ includeInitialState: boolean = false
173
113
  ) {
174
- // Flag to indicate we have send to whole initial state of the collection
175
- // to the pipeline, this is set when there are no indexes that can be used
176
- // to filter the changes and so the whole state was requested from the collection
177
- let loadedInitialState = false
178
-
179
- // Flag to indicate that we have started sending changes to the pipeline.
180
- // This is set to true by either the first call to `loadKeys` or when the
181
- // query requests the whole initial state in `loadInitialState`.
182
- // Until that point we filter out all changes from subscription to the collection.
183
- let sendChanges = false
184
-
185
- const sendVisibleChanges = (
114
+ const sendChanges = (
186
115
  changes: Array<ChangeMessage<any, string | number>>
187
116
  ) => {
188
- // We filter out changes when sendChanges is false to ensure that we don't send
189
- // any changes from the live subscription until the join operator requests either
190
- // the initial state or its first key. This is needed otherwise it could receive
191
- // changes which are then later subsumed by the initial state (and that would
192
- // lead to weird bugs due to the data being received twice).
193
- this.sendVisibleChangesToPipeline(
194
- sendChanges ? changes : [],
195
- loadedInitialState
196
- )
117
+ this.sendChangesToPipeline(changes)
197
118
  }
198
119
 
199
- const unsubscribe = this.collection.subscribeChanges(sendVisibleChanges, {
120
+ const subscription = this.collection.subscribeChanges(sendChanges, {
121
+ includeInitialState,
200
122
  whereExpression,
201
123
  })
202
124
 
203
- // Create a function that loads keys from the collection
204
- // into the query pipeline on demand
205
- const filterFn = whereExpression
206
- ? createFilterFunctionFromExpression(whereExpression)
207
- : () => true
208
- const loadKs = (keys: Set<string | number>) => {
209
- sendChanges = true
210
- return this.loadKeys(keys, filterFn)
211
- }
212
-
213
- // Store the functions to load keys and load initial state in the `lazyCollectionsCallbacks` map
214
- // This is used by the join operator to dynamically load matching keys from the lazy collection
215
- // or to get the full initial state of the collection if there's no index for the join key
216
- this.collectionConfigBuilder.lazyCollectionsCallbacks[this.collectionId] = {
217
- loadKeys: loadKs,
218
- loadInitialState: () => {
219
- // Make sure we only load the initial state once
220
- if (loadedInitialState) return
221
- loadedInitialState = true
222
- sendChanges = true
223
-
224
- const changes = this.collection.currentStateAsChanges({
225
- whereExpression,
226
- })
227
- this.sendChangesToPipeline(changes)
228
- },
229
- }
230
- return unsubscribe
125
+ return subscription
231
126
  }
232
127
 
233
128
  private subscribeToOrderedChanges(
234
129
  whereExpression: BasicExpression<boolean> | undefined
235
130
  ) {
236
- const { offset, limit, comparator, dataNeeded } =
131
+ const { offset, limit, comparator, dataNeeded, index } =
237
132
  this.collectionConfigBuilder.optimizableOrderByCollections[
238
133
  this.collectionId
239
134
  ]!
240
135
 
241
- // Load the first `offset + limit` values from the index
242
- // i.e. the K items from the collection that fall into the requested range: [offset, offset + limit[
243
- this.loadNextItems(offset + limit)
244
-
245
136
  const sendChangesInRange = (
246
137
  changes: Iterable<ChangeMessage<any, string | number>>
247
138
  ) => {
248
139
  // Split live updates into a delete of the old value and an insert of the new value
249
140
  // and filter out changes that are bigger than the biggest value we've sent so far
250
- // because they can't affect the topK
141
+ // because they can't affect the topK (and if later we need more data, we will dynamically load more data)
251
142
  const splittedChanges = splitUpdates(changes)
252
143
  let filteredChanges = splittedChanges
253
144
  if (dataNeeded!() === 0) {
@@ -261,25 +152,31 @@ export class CollectionSubscriber<
261
152
  this.biggest
262
153
  )
263
154
  }
264
- this.sendChangesToPipeline(
265
- filteredChanges,
266
- this.loadMoreIfNeeded.bind(this)
267
- )
155
+
156
+ this.sendChangesToPipelineWithTracking(filteredChanges, subscription)
268
157
  }
269
158
 
270
159
  // Subscribe to changes and only send changes that are smaller than the biggest value we've sent so far
271
160
  // values that are bigger don't need to be sent because they can't affect the topK
272
- const unsubscribe = this.collection.subscribeChanges(sendChangesInRange, {
161
+ const subscription = this.collection.subscribeChanges(sendChangesInRange, {
273
162
  whereExpression,
274
163
  })
275
164
 
276
- return unsubscribe
165
+ subscription.setOrderByIndex(index)
166
+
167
+ // Load the first `offset + limit` values from the index
168
+ // i.e. the K items from the collection that fall into the requested range: [offset, offset + limit[
169
+ subscription.requestLimitedSnapshot({
170
+ limit: offset + limit,
171
+ })
172
+
173
+ return subscription
277
174
  }
278
175
 
279
176
  // This function is called by maybeRunGraph
280
177
  // after each iteration of the query pipeline
281
178
  // to ensure that the orderBy operator has enough data to work with
282
- loadMoreIfNeeded() {
179
+ loadMoreIfNeeded(subscription: CollectionSubscription) {
283
180
  const orderByInfo =
284
181
  this.collectionConfigBuilder.optimizableOrderByCollections[
285
182
  this.collectionId
@@ -287,7 +184,7 @@ export class CollectionSubscriber<
287
184
 
288
185
  if (!orderByInfo) {
289
186
  // This query has no orderBy operator
290
- // so there's no data to load, just return true
187
+ // so there's no data to load
291
188
  return true
292
189
  }
293
190
 
@@ -304,32 +201,31 @@ export class CollectionSubscriber<
304
201
  // `dataNeeded` probes the orderBy operator to see if it needs more data
305
202
  // if it needs more data, it returns the number of items it needs
306
203
  const n = dataNeeded()
307
- let noMoreNextItems = false
308
204
  if (n > 0) {
309
- const loadedItems = this.loadNextItems(n)
310
- noMoreNextItems = loadedItems === 0
205
+ this.loadNextItems(n, subscription)
311
206
  }
312
-
313
- // Indicate that we're done loading data if we didn't need to load more data
314
- // or there's no more data to load
315
- return n === 0 || noMoreNextItems
207
+ return true
316
208
  }
317
209
 
318
210
  private sendChangesToPipelineWithTracking(
319
- changes: Iterable<ChangeMessage<any, string | number>>
211
+ changes: Iterable<ChangeMessage<any, string | number>>,
212
+ subscription: CollectionSubscription
320
213
  ) {
321
214
  const { comparator } =
322
215
  this.collectionConfigBuilder.optimizableOrderByCollections[
323
216
  this.collectionId
324
217
  ]!
325
218
  const trackedChanges = this.trackSentValues(changes, comparator)
326
- this.sendChangesToPipeline(trackedChanges, this.loadMoreIfNeeded.bind(this))
219
+ this.sendChangesToPipeline(
220
+ trackedChanges,
221
+ this.loadMoreIfNeeded.bind(this, subscription)
222
+ )
327
223
  }
328
224
 
329
225
  // Loads the next `n` items from the collection
330
226
  // starting from the biggest item it has sent
331
- private loadNextItems(n: number) {
332
- const { valueExtractorForRawRow, index } =
227
+ private loadNextItems(n: number, subscription: CollectionSubscription) {
228
+ const { valueExtractorForRawRow } =
333
229
  this.collectionConfigBuilder.optimizableOrderByCollections[
334
230
  this.collectionId
335
231
  ]!
@@ -338,13 +234,10 @@ export class CollectionSubscriber<
338
234
  ? valueExtractorForRawRow(biggestSentRow)
339
235
  : biggestSentRow
340
236
  // Take the `n` items after the biggest sent value
341
- const nextOrderedKeys = index.take(n, biggestSentValue)
342
- const nextInserts: Array<ChangeMessage<any, string | number>> =
343
- nextOrderedKeys.map((key) => {
344
- return { type: `insert`, key, value: this.collection.get(key) }
345
- })
346
- this.sendChangesToPipelineWithTracking(nextInserts)
347
- return nextInserts.length
237
+ subscription.requestLimitedSnapshot({
238
+ limit: n,
239
+ minValue: biggestSentValue,
240
+ })
348
241
  }
349
242
 
350
243
  private getWhereClauseFromAlias(
@@ -363,8 +256,6 @@ export class CollectionSubscriber<
363
256
  comparator: (a: any, b: any) => number
364
257
  ) {
365
258
  for (const change of changes) {
366
- this.sentKeys.add(change.key)
367
-
368
259
  if (!this.biggest) {
369
260
  this.biggest = change.value
370
261
  } else if (comparator(this.biggest, change.value) < 0) {
@@ -427,7 +318,11 @@ function sendChangesToInput(
427
318
  multiSetArray.push([[key, change.value], -1])
428
319
  }
429
320
  }
430
- input.sendData(new MultiSet(multiSetArray))
321
+
322
+ if (multiSetArray.length !== 0) {
323
+ input.sendData(new MultiSet(multiSetArray))
324
+ }
325
+
431
326
  return multiSetArray.length
432
327
  }
433
328
 
@@ -463,7 +358,7 @@ function* filterChanges<
463
358
  }
464
359
 
465
360
  /**
466
- * Filters changes to only include those that are smaller than the provided max value
361
+ * Filters changes to only include those that are smaller or equal to the provided max value
467
362
  * @param changes - Iterable of changes to filter
468
363
  * @param comparator - Comparator function to use for filtering
469
364
  * @param maxValue - Range to filter changes within (range boundaries are exclusive)