@tanstack/db-ivm 0.1.0 → 0.1.2

package/src/operators/tap.ts

@@ -0,0 +1,53 @@
+import { DifferenceStreamWriter, LinearUnaryOperator } from "../graph.js"
+import { StreamBuilder } from "../d2.js"
+import type { IStreamBuilder, PipedOperator } from "../types.js"
+import type { DifferenceStreamReader } from "../graph.js"
+import type { MultiSet } from "../multiset.js"
+
+/**
+ * Operator that applies a function to each element in the input stream
+ */
+export class TapOperator<T> extends LinearUnaryOperator<T, T> {
+  #f: (data: T) => void
+
+  constructor(
+    id: number,
+    inputA: DifferenceStreamReader<T>,
+    output: DifferenceStreamWriter<T>,
+    f: (data: T) => void
+  ) {
+    super(id, inputA, output)
+    this.#f = f
+  }
+
+  inner(collection: MultiSet<T>): MultiSet<T> {
+    return collection.map((data) => {
+      this.#f(data)
+      return data
+    })
+  }
+}
+
+/**
+ * Invokes a function for each element 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
+ * @returns The input stream
+ */
+export function tap<T>(f: (data: T) => void): PipedOperator<T, T> {
+  return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
+    const output = new StreamBuilder<T>(
+      stream.graph,
+      new DifferenceStreamWriter<T>()
+    )
+    const operator = new TapOperator<T>(
+      stream.graph.getNextOperatorId(),
+      stream.connectReader(),
+      output.writer,
+      f
+    )
+    stream.graph.addOperator(operator)
+    stream.graph.addStream(output.connectReader())
+    return output
+  }
+}

package/src/operators/topKWithFractionalIndex.ts

@@ -2,14 +2,14 @@ import { generateKeyBetween } from "fractional-indexing"
 import { DifferenceStreamWriter, UnaryOperator } from "../graph.js"
 import { StreamBuilder } from "../d2.js"
 import { MultiSet } from "../multiset.js"
-import { Index } from "../indexes.js"
 import { binarySearch, globalObjectIdGenerator } from "../utils.js"
 import type { DifferenceStreamReader } from "../graph.js"
-import type { IStreamBuilder, KeyValue, PipedOperator } from "../types.js"
+import type { IStreamBuilder, PipedOperator } from "../types.js"
 
 export interface TopKWithFractionalIndexOptions {
   limit?: number
   offset?: number
+  setSizeCallback?: (getSize: () => number) => void
 }
 
 export type TopKChanges<V> = {
@@ -24,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>
 }
@@ -50,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 }
 
@@ -158,31 +166,36 @@ class TopKArray<V> implements TopK<V> {
  * This operator maintains fractional indices for sorted elements
  * and only updates indices when elements move position
  */
-export class TopKWithFractionalIndexOperator<K, V1> extends UnaryOperator<
-  [K, V1],
-  [K, IndexedValue<V1>]
+export class TopKWithFractionalIndexOperator<K, T> extends UnaryOperator<
+  [K, T],
+  [K, IndexedValue<T>]
 > {
-  #index = new Index<K, V1>()
+  #index: Map<K, number> = new Map() // maps keys to their multiplicity
 
   /**
    * topK data structure that supports insertions and deletions
    * and returns changes to the topK.
    */
-  #topK: TopK<TaggedValue<V1>>
+  #topK: TopK<TaggedValue<K, T>>
+
+  #limit: number
 
   constructor(
     id: number,
-    inputA: DifferenceStreamReader<[K, V1]>,
-    output: DifferenceStreamWriter<[K, [V1, string]]>,
-    comparator: (a: V1, b: V1) => number,
+    inputA: DifferenceStreamReader<[K, T]>,
+    output: DifferenceStreamWriter<[K, IndexedValue<T>]>,
+    comparator: (a: T, b: T) => number,
     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<V1>, b: TaggedValue<V1>) => {
+    const compareTaggedValues = (
+      a: TaggedValue<K, T>,
+      b: TaggedValue<K, T>
+    ) => {
       // First compare on the value
-      const valueComparison = comparator(untagValue(a), untagValue(b))
+      const valueComparison = comparator(getVal(a), getVal(b))
       if (valueComparison !== 0) {
         return valueComparison
       }
@@ -191,19 +204,20 @@ export class TopKWithFractionalIndexOperator<K, V1> 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(
     offset: number,
     limit: number,
-    comparator: (a: TaggedValue<V1>, b: TaggedValue<V1>) => number
-  ): TopK<TaggedValue<V1>> {
+    comparator: (a: TaggedValue<K, T>, b: TaggedValue<K, T>) => number
+  ): TopK<TaggedValue<K, T>> {
     return new TopKArray(offset, limit, comparator)
   }
 
   run(): void {
-    const result: Array<[[K, [V1, string]], number]> = []
+    const result: Array<[[K, IndexedValue<T>], number]> = []
     for (const message of this.inputMessages()) {
       for (const [item, multiplicity] of message.getInner()) {
         const [key, value] = item
@@ -218,27 +232,25 @@ export class TopKWithFractionalIndexOperator<K, V1> extends UnaryOperator<
 
   processElement(
     key: K,
-    value: V1,
+    value: T,
     multiplicity: number,
-    result: Array<[[K, [V1, string]], number]>
+    result: Array<[[K, IndexedValue<T>], number]>
   ): void {
-    const oldMultiplicity = this.#index.getMultiplicity(key, value)
-    this.#index.addValue(key, [value, multiplicity])
-    const newMultiplicity = this.#index.getMultiplicity(key, value)
+    const { oldMultiplicity, newMultiplicity } = this.addKey(key, multiplicity)
 
-    let res: TopKChanges<TaggedValue<V1>> = {
+    let res: TopKChanges<TaggedValue<K, T>> = {
       moveIn: null,
       moveOut: null,
     }
     if (oldMultiplicity <= 0 && newMultiplicity > 0) {
       // The value was invisible but should now be visible
       // Need to insert it into the array of sorted values
-      const taggedValue = tagValue(value)
+      const taggedValue = tagValue(key, value)
       res = this.#topK.insert(taggedValue)
     } else if (oldMultiplicity > 0 && newMultiplicity <= 0) {
       // The value was visible but should now be invisible
       // Need to remove it from the array of sorted values
-      const taggedValue = tagValue(value)
+      const taggedValue = tagValue(key, value)
       res = this.#topK.delete(taggedValue)
     } else {
       // The value was invisible and it remains invisible
@@ -247,26 +259,45 @@ export class TopKWithFractionalIndexOperator<K, V1> extends UnaryOperator<
     }
 
     if (res.moveIn) {
-      const valueWithoutTieBreaker = mapValue(res.moveIn, untagValue)
-      result.push([[key, valueWithoutTieBreaker], 1])
+      const index = getIndex(res.moveIn)
+      const taggedValue = getValue(res.moveIn)
+      const k = getKey(taggedValue)
+      const val = getVal(taggedValue)
+      result.push([[k, [val, index]], 1])
     }
 
     if (res.moveOut) {
-      const valueWithoutTieBreaker = mapValue(res.moveOut, untagValue)
-      result.push([[key, valueWithoutTieBreaker], -1])
+      const index = getIndex(res.moveOut)
+      const taggedValue = getValue(res.moveOut)
+      const k = getKey(taggedValue)
+      const val = getVal(taggedValue)
+      result.push([[k, [val, index]], -1])
     }
 
     return
   }
+
+  private getMultiplicity(key: K): number {
+    return this.#index.get(key) ?? 0
+  }
+
+  private addKey(
+    key: K,
+    multiplicity: number
+  ): { oldMultiplicity: number; newMultiplicity: number } {
+    const oldMultiplicity = this.getMultiplicity(key)
+    const newMultiplicity = oldMultiplicity + multiplicity
+    if (newMultiplicity === 0) {
+      this.#index.delete(key)
+    } else {
+      this.#index.set(key, newMultiplicity)
+    }
+    return { oldMultiplicity, newMultiplicity }
+  }
 }
 
 /**
  * Limits the number of results based on a comparator, with optional offset.
- * This works on a keyed stream, where the key is the first element of the tuple.
- * The ordering is within a key group, i.e. elements are sorted within a key group
- * and the limit + offset is applied to that sorted group.
- * To order the entire stream, key by the same value for all elements such as null.
- *
  * Uses fractional indexing to minimize the number of changes when elements move positions.
  * Each element is assigned a fractional index that is lexicographically sortable.
  * When elements move, only the indices of the moved elements are updated, not all elements.
@@ -275,26 +306,22 @@ export class TopKWithFractionalIndexOperator<K, V1> extends UnaryOperator<
  * @param options - An optional object containing limit and offset properties
  * @returns A piped operator that orders the elements and limits the number of results
  */
-export function topKWithFractionalIndex<
-  KType extends T extends KeyValue<infer K, infer _V> ? K : never,
-  V1Type extends T extends KeyValue<KType, infer V> ? V : never,
-  T,
->(
-  comparator: (a: V1Type, b: V1Type) => number,
+export function topKWithFractionalIndex<KType, T>(
+  comparator: (a: T, b: T) => number,
   options?: TopKWithFractionalIndexOptions
-): PipedOperator<T, KeyValue<KType, [V1Type, string]>> {
+): PipedOperator<[KType, T], [KType, IndexedValue<T>]> {
   const opts = options || {}
 
   return (
-    stream: IStreamBuilder<T>
-  ): IStreamBuilder<KeyValue<KType, [V1Type, string]>> => {
-    const output = new StreamBuilder<KeyValue<KType, [V1Type, string]>>(
+    stream: IStreamBuilder<[KType, T]>
+  ): IStreamBuilder<[KType, IndexedValue<T>]> => {
+    const output = new StreamBuilder<[KType, IndexedValue<T>]>(
       stream.graph,
-      new DifferenceStreamWriter<KeyValue<KType, [V1Type, string]>>()
+      new DifferenceStreamWriter<[KType, IndexedValue<T>]>()
     )
-    const operator = new TopKWithFractionalIndexOperator<KType, V1Type>(
+    const operator = new TopKWithFractionalIndexOperator<KType, T>(
       stream.graph.getNextOperatorId(),
-      stream.connectReader() as DifferenceStreamReader<KeyValue<KType, V1Type>>,
+      stream.connectReader(),
       output.writer,
       comparator,
       opts
@@ -324,24 +351,21 @@ export function getIndex<V>(indexedVal: IndexedValue<V>): FractionalIndex {
   return indexedVal[1]
 }
 
-function mapValue<V, W>(
-  indexedVal: IndexedValue<V>,
-  f: (value: V) => W
-): IndexedValue<W> {
-  return [f(getValue(indexedVal)), getIndex(indexedVal)]
-}
-
 export type Tag = number
-export type TaggedValue<V> = [V, Tag]
+export type TaggedValue<K, V> = [K, V, Tag]
 
-function tagValue<V>(value: V): TaggedValue<V> {
-  return [value, globalObjectIdGenerator.getId(value)]
+function tagValue<K, V>(key: K, value: V): TaggedValue<K, V> {
+  return [key, value, globalObjectIdGenerator.getId(key)]
 }
 
-function untagValue<V>(tieBreakerTaggedValue: TaggedValue<V>): V {
+function getKey<K, V>(tieBreakerTaggedValue: TaggedValue<K, V>): K {
   return tieBreakerTaggedValue[0]
 }
 
-function getTag<V>(tieBreakerTaggedValue: TaggedValue<V>): Tag {
+function getVal<K, V>(tieBreakerTaggedValue: TaggedValue<K, V>): V {
   return tieBreakerTaggedValue[1]
 }
+
+function getTag<K, V>(tieBreakerTaggedValue: TaggedValue<K, V>): Tag {
+  return tieBreakerTaggedValue[2]
+}

package/src/operators/topKWithFractionalIndexBTree.ts

@@ -7,8 +7,7 @@ import {
   getValue,
   indexedValue,
 } from "./topKWithFractionalIndex.js"
-import type { IStreamBuilder, KeyValue, PipedOperator } from "../types.js"
-import type { DifferenceStreamReader } from "../graph.js"
+import type { IStreamBuilder, PipedOperator } from "../types.js"
 import type {
   IndexedValue,
   TaggedValue,
@@ -76,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.
@@ -238,13 +244,13 @@ class TopKTree<V> implements TopK<V> {
  */
 export class TopKWithFractionalIndexBTreeOperator<
   K,
-  V1,
-> extends TopKWithFractionalIndexOperator<K, V1> {
+  T,
+> extends TopKWithFractionalIndexOperator<K, T> {
   protected override createTopK(
     offset: number,
     limit: number,
-    comparator: (a: TaggedValue<V1>, b: TaggedValue<V1>) => number
-  ): TopK<TaggedValue<V1>> {
+    comparator: (a: TaggedValue<K, T>, b: TaggedValue<K, T>) => number
+  ): TopK<TaggedValue<K, T>> {
     if (BTree === undefined) {
       throw new Error(
         `B+ tree not loaded. You need to call loadBTree() before using TopKWithFractionalIndexBTreeOperator.`
@@ -269,14 +275,10 @@ export class TopKWithFractionalIndexBTreeOperator<
  * @param options - An optional object containing limit and offset properties
  * @returns A piped operator that orders the elements and limits the number of results
  */
-export function topKWithFractionalIndexBTree<
-  KType extends T extends KeyValue<infer K, infer _V> ? K : never,
-  V1Type extends T extends KeyValue<KType, infer V> ? V : never,
-  T,
->(
-  comparator: (a: V1Type, b: V1Type) => number,
+export function topKWithFractionalIndexBTree<KType, T>(
+  comparator: (a: T, b: T) => number,
   options?: TopKWithFractionalIndexOptions
-): PipedOperator<T, KeyValue<KType, [V1Type, string]>> {
+): PipedOperator<[KType, T], [KType, IndexedValue<T>]> {
   const opts = options || {}
 
   if (BTree === undefined) {
@@ -286,15 +288,15 @@ export function topKWithFractionalIndexBTree<
   }
 
   return (
-    stream: IStreamBuilder<T>
-  ): IStreamBuilder<KeyValue<KType, [V1Type, string]>> => {
-    const output = new StreamBuilder<KeyValue<KType, [V1Type, string]>>(
+    stream: IStreamBuilder<[KType, T]>
+  ): IStreamBuilder<[KType, IndexedValue<T>]> => {
+    const output = new StreamBuilder<[KType, IndexedValue<T>]>(
       stream.graph,
-      new DifferenceStreamWriter<KeyValue<KType, [V1Type, string]>>()
+      new DifferenceStreamWriter<[KType, IndexedValue<T>]>()
     )
-    const operator = new TopKWithFractionalIndexOperator<KType, V1Type>(
+    const operator = new TopKWithFractionalIndexBTreeOperator<KType, T>(
       stream.graph.getNextOperatorId(),
-      stream.connectReader() as DifferenceStreamReader<KeyValue<KType, V1Type>>,
+      stream.connectReader(),
       output.writer,
       comparator,
       opts