@tanstack/db-ivm 0.0.0

package/src/operators/map.ts

@@ -0,0 +1,48 @@
+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 MapOperator<T, U> extends LinearUnaryOperator<T, U> {
+  #f: (data: T) => U
+
+  constructor(
+    id: number,
+    inputA: DifferenceStreamReader<T>,
+    output: DifferenceStreamWriter<U>,
+    f: (data: T) => U
+  ) {
+    super(id, inputA, output)
+    this.#f = f
+  }
+
+  inner(collection: MultiSet<T>): MultiSet<U> {
+    return collection.map(this.#f)
+  }
+}
+
+/**
+ * Applies a function to each element in the input stream
+ * @param f - The function to apply to each element
+ */
+export function map<T, O>(f: (data: T) => O): PipedOperator<T, O> {
+  return (stream: IStreamBuilder<T>): IStreamBuilder<O> => {
+    const output = new StreamBuilder<O>(
+      stream.graph,
+      new DifferenceStreamWriter<O>()
+    )
+    const operator = new MapOperator<T, O>(
+      stream.graph.getNextOperatorId(),
+      stream.connectReader(),
+      output.writer,
+      f
+    )
+    stream.graph.addOperator(operator)
+    stream.graph.addStream(output.connectReader())
+    return output
+  }
+}

package/src/operators/negate.ts

@@ -0,0 +1,33 @@
+import { DifferenceStreamWriter, LinearUnaryOperator } from "../graph.js"
+import { StreamBuilder } from "../d2.js"
+import type { IStreamBuilder, PipedOperator } from "../types.js"
+import type { MultiSet } from "../multiset.js"
+
+/**
+ * Operator that negates the multiplicities in the input stream
+ */
+export class NegateOperator<T> extends LinearUnaryOperator<T, T> {
+  inner(collection: MultiSet<T>): MultiSet<T> {
+    return collection.negate()
+  }
+}
+
+/**
+ * Negates the multiplicities in the input stream
+ */
+export function negate<T>(): PipedOperator<T, T> {
+  return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
+    const output = new StreamBuilder<T>(
+      stream.graph,
+      new DifferenceStreamWriter<T>()
+    )
+    const operator = new NegateOperator<T>(
+      stream.graph.getNextOperatorId(),
+      stream.connectReader(),
+      output.writer
+    )
+    stream.graph.addOperator(operator)
+    stream.graph.addStream(output.connectReader())
+    return output
+  }
+}

package/src/operators/orderBy.ts

@@ -0,0 +1,213 @@
+import { topK, topKWithIndex } from "./topK.js"
+import { topKWithFractionalIndex } from "./topKWithFractionalIndex.js"
+import { map } from "./map.js"
+import { innerJoin } from "./join.js"
+import { consolidate } from "./consolidate.js"
+import type { KeyValue } from "../types.js"
+import type { IStreamBuilder } from "../types"
+
+export interface OrderByOptions<Ve> {
+  comparator?: (a: Ve, b: Ve) => number
+  limit?: number
+  offset?: number
+}
+
+/**
+ * Orders the elements and limits the number of results, with optional offset
+ * This requires a keyed stream, and uses the `topK` operator to order all the elements.
+ *
+ * @param valueExtractor - A function that extracts the value to order by from the element
+ * @param options - An optional object containing comparator, limit and offset properties
+ * @returns A piped operator that orders the elements and limits the number of results
+ */
+export function orderBy<T extends KeyValue<unknown, unknown>, Ve = unknown>(
+  valueExtractor: (
+    value: T extends KeyValue<unknown, infer V> ? V : never
+  ) => Ve,
+  options?: OrderByOptions<Ve>
+) {
+  const limit = options?.limit ?? Infinity
+  const offset = options?.offset ?? 0
+  const comparator =
+    options?.comparator ??
+    ((a, b) => {
+      // Default to JS like ordering
+      if (a === b) return 0
+      if (a < b) return -1
+      return 1
+    })
+
+  return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
+    type KeyType = T extends KeyValue<infer K, unknown> ? K : never
+
+    return stream.pipe(
+      map(
+        ([key, value]) =>
+          [
+            null,
+            [
+              key,
+              valueExtractor(
+                value as T extends KeyValue<unknown, infer V> ? V : never
+              ),
+            ],
+          ] as KeyValue<null, [KeyType, Ve]>
+      ),
+      topK((a, b) => comparator(a[1], b[1]), { limit, offset }),
+      map(([_, [key]]) => [key, null] as KeyValue<KeyType, null>),
+      innerJoin(stream),
+      map(([key, value]) => {
+        return [key, value[1]] as T
+      }),
+      consolidate()
+    )
+  }
+}
+
+/**
+ * Orders the elements and limits the number of results, with optional offset and
+ * annotates the value with the index.
+ * This requires a keyed stream, and uses the `topKWithIndex` operator to order all the elements.
+ *
+ * @param valueExtractor - A function that extracts the value to order by from the element
+ * @param options - An optional object containing comparator, limit and offset properties
+ * @returns A piped operator that orders the elements and limits the number of results
+ */
+export function orderByWithIndex<
+  T extends KeyValue<unknown, unknown>,
+  Ve = unknown,
+>(
+  valueExtractor: (
+    value: T extends KeyValue<unknown, infer V> ? V : never
+  ) => Ve,
+  options?: OrderByOptions<Ve>
+) {
+  const limit = options?.limit ?? Infinity
+  const offset = options?.offset ?? 0
+  const comparator =
+    options?.comparator ??
+    ((a, b) => {
+      // Default to JS like ordering
+      if (a === b) return 0
+      if (a < b) return -1
+      return 1
+    })
+
+  return (
+    stream: IStreamBuilder<T>
+  ): IStreamBuilder<
+    KeyValue<
+      T extends KeyValue<infer K, unknown> ? K : never,
+      [T extends KeyValue<unknown, infer V> ? V : never, number]
+    >
+  > => {
+    type KeyType = T extends KeyValue<infer K, unknown> ? K : never
+    type ValueType = T extends KeyValue<unknown, infer V> ? V : never
+
+    return stream.pipe(
+      map(
+        ([key, value]) =>
+          [
+            null,
+            [
+              key,
+              valueExtractor(
+                value as T extends KeyValue<unknown, infer V> ? V : never
+              ),
+            ],
+          ] as KeyValue<null, [KeyType, Ve]>
+      ),
+      topKWithIndex((a, b) => comparator(a[1], b[1]), { limit, offset }),
+      map(([_, [[key], index]]) => [key, index] as KeyValue<KeyType, number>),
+      innerJoin(stream),
+      map(([key, [index, value]]) => {
+        return [key, [value, index]] as KeyValue<KeyType, [ValueType, number]>
+      }),
+      consolidate()
+    )
+  }
+}
+
+export function orderByWithFractionalIndexBase<
+  T extends KeyValue<unknown, unknown>,
+  Ve = unknown,
+>(
+  topKFunction: typeof topKWithFractionalIndex,
+  valueExtractor: (
+    value: T extends KeyValue<unknown, infer V> ? V : never
+  ) => Ve,
+  options?: OrderByOptions<Ve>
+) {
+  const limit = options?.limit ?? Infinity
+  const offset = options?.offset ?? 0
+  const comparator =
+    options?.comparator ??
+    ((a, b) => {
+      // Default to JS like ordering
+      if (a === b) return 0
+      if (a < b) return -1
+      return 1
+    })
+
+  return (
+    stream: IStreamBuilder<T>
+  ): IStreamBuilder<
+    KeyValue<
+      T extends KeyValue<infer K, unknown> ? K : never,
+      [T extends KeyValue<unknown, infer V> ? V : never, string]
+    >
+  > => {
+    type KeyType = T extends KeyValue<infer K, unknown> ? K : never
+    type ValueType = T extends KeyValue<unknown, infer V> ? V : never
+
+    return stream.pipe(
+      map(
+        ([key, value]) =>
+          [
+            null,
+            [
+              key,
+              valueExtractor(
+                value as T extends KeyValue<unknown, infer V> ? V : never
+              ),
+            ],
+          ] as KeyValue<null, [KeyType, Ve]>
+      ),
+      topKFunction((a, b) => comparator(a[1], b[1]), {
+        limit,
+        offset,
+      }),
+      map(([_, [[key], index]]) => [key, index] as KeyValue<KeyType, string>),
+      innerJoin(stream),
+      map(([key, [index, value]]) => {
+        return [key, [value, index]] as KeyValue<KeyType, [ValueType, string]>
+      }),
+      consolidate()
+    )
+  }
+}
+
+/**
+ * Orders the elements and limits the number of results, with optional offset and
+ * annotates the value with a fractional index.
+ * This requires a keyed stream, and uses the `topKWithFractionalIndex` operator to order all the elements.
+ *
+ * @param valueExtractor - A function that extracts the value to order by from the element
+ * @param options - An optional object containing comparator, limit and offset properties
+ * @returns A piped operator that orders the elements and limits the number of results
+ */
+export function orderByWithFractionalIndex<
+  T extends KeyValue<unknown, unknown>,
+  Ve = unknown,
+>(
+  valueExtractor: (
+    value: T extends KeyValue<unknown, infer V> ? V : never
+  ) => Ve,
+  options?: OrderByOptions<Ve>
+) {
+  return orderByWithFractionalIndexBase(
+    topKWithFractionalIndex,
+    valueExtractor,
+    options
+  )
+}

package/src/operators/orderByBTree.ts

@@ -0,0 +1,20 @@
+import { orderByWithFractionalIndexBase } from "./orderBy.js"
+import { topKWithFractionalIndexBTree } from "./topKWithFractionalIndexBTree.js"
+import type { KeyValue } from "../types.js"
+import type { OrderByOptions } from "./orderBy.js"
+
+export function orderByWithFractionalIndexBTree<
+  T extends KeyValue<unknown, unknown>,
+  Ve = unknown,
+>(
+  valueExtractor: (
+    value: T extends KeyValue<unknown, infer V> ? V : never
+  ) => Ve,
+  options?: OrderByOptions<Ve>
+) {
+  return orderByWithFractionalIndexBase(
+    topKWithFractionalIndexBTree,
+    valueExtractor,
+    options
+  )
+}

package/src/operators/output.ts

@@ -0,0 +1,53 @@
+import { DifferenceStreamWriter, UnaryOperator } 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 outputs the messages in the stream
+ */
+export class OutputOperator<T> extends UnaryOperator<T> {
+  #fn: (data: MultiSet<T>) => void
+
+  constructor(
+    id: number,
+    inputA: DifferenceStreamReader<T>,
+    outputWriter: DifferenceStreamWriter<T>,
+    fn: (data: MultiSet<T>) => void
+  ) {
+    super(id, inputA, outputWriter)
+    this.#fn = fn
+  }
+
+  run(): void {
+    for (const message of this.inputMessages()) {
+      this.#fn(message)
+      this.output.sendData(message)
+    }
+  }
+}
+
+/**
+ * Outputs the messages in the stream
+ * @param fn - The function to call with each message
+ */
+export function output<T>(
+  fn: (data: MultiSet<T>) => void
+): PipedOperator<T, T> {
+  return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
+    const outputStream = new StreamBuilder<T>(
+      stream.graph,
+      new DifferenceStreamWriter<T>()
+    )
+    const operator = new OutputOperator<T>(
+      stream.graph.getNextOperatorId(),
+      stream.connectReader(),
+      outputStream.writer,
+      fn
+    )
+    stream.graph.addOperator(operator)
+    stream.graph.addStream(outputStream.connectReader())
+    return outputStream
+  }
+}

package/src/operators/pipe.ts

@@ -0,0 +1,55 @@
+import type { IStreamBuilder, PipedOperator } from "../types.js"
+
+// Don't judge, this is the only way to type this function.
+// rxjs has very similar code to type its pipe function
+// https://github.com/ReactiveX/rxjs/blob/master/packages/rxjs/src/internal/util/pipe.ts
+// We go to 20 operators deep, because surely that's enough for anyone...
+// A user can always split the pipe into multiple pipes to get around this.
+export function pipe<T, O>(o1: PipedOperator<T, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, T12>, o12: PipedOperator<T12, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, T12>, o12: PipedOperator<T12, T13>, o13: PipedOperator<T13, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, T12>, o12: PipedOperator<T12, T13>, o13: PipedOperator<T13, T14>, o14: PipedOperator<T14, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, T12>, o12: PipedOperator<T12, T13>, o13: PipedOperator<T13, T14>, o14: PipedOperator<T14, T15>, o15: PipedOperator<T15, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, T12>, o12: PipedOperator<T12, T13>, o13: PipedOperator<T13, T14>, o14: PipedOperator<T14, T15>, o15: PipedOperator<T15, T16>, o16: PipedOperator<T16, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, T12>, o12: PipedOperator<T12, T13>, o13: PipedOperator<T13, T14>, o14: PipedOperator<T14, T15>, o15: PipedOperator<T15, T16>, o16: PipedOperator<T16, T17>, o17: PipedOperator<T17, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, T12>, o12: PipedOperator<T12, T13>, o13: PipedOperator<T13, T14>, o14: PipedOperator<T14, T15>, o15: PipedOperator<T15, T16>, o16: PipedOperator<T16, T17>, o17: PipedOperator<T17, T18>, o18: PipedOperator<T18, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, T12>, o12: PipedOperator<T12, T13>, o13: PipedOperator<T13, T14>, o14: PipedOperator<T14, T15>, o15: PipedOperator<T15, T16>, o16: PipedOperator<T16, T17>, o17: PipedOperator<T17, T18>, o18: PipedOperator<T18, T19>, o19: PipedOperator<T19, O>): PipedOperator<T, O>
+// prettier-ignore
+export function pipe<T, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20, O>(o1: PipedOperator<T, T2>, o2: PipedOperator<T2, T3>, o3: PipedOperator<T3, T4>, o4: PipedOperator<T4, T5>, o5: PipedOperator<T5, T6>, o6: PipedOperator<T6, T7>, o7: PipedOperator<T7, T8>, o8: PipedOperator<T8, T9>, o9: PipedOperator<T9, T10>, o10: PipedOperator<T10, T11>, o11: PipedOperator<T11, T12>, o12: PipedOperator<T12, T13>, o13: PipedOperator<T13, T14>, o14: PipedOperator<T14, T15>, o15: PipedOperator<T15, T16>, o16: PipedOperator<T16, T17>, o17: PipedOperator<T17, T18>, o18: PipedOperator<T18, T19>, o19: PipedOperator<T19, T20>, o20: PipedOperator<T20, O>): PipedOperator<T, O>
+
+/**
+ * Creates a new stream by piping the input stream through a series of operators
+ */
+export function pipe<T>(...operators: Array<PipedOperator<any, any>>) {
+  return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
+    return stream.pipe(...operators)
+  }
+}

package/src/operators/reduce.ts

@@ -0,0 +1,124 @@
+import { DifferenceStreamWriter, UnaryOperator } from "../graph.js"
+import { StreamBuilder } from "../d2.js"
+import { MultiSet } from "../multiset.js"
+import { Index } from "../indexes.js"
+import type { DifferenceStreamReader } from "../graph.js"
+import type { IStreamBuilder, KeyValue } from "../types.js"
+
+/**
+ * Base operator for reduction operations (version-free)
+ */
+export class ReduceOperator<K, V1, V2> extends UnaryOperator<[K, V1], [K, V2]> {
+  #index = new Index<K, V1>()
+  #indexOut = new Index<K, V2>()
+  #f: (values: Array<[V1, number]>) => Array<[V2, number]>
+
+  constructor(
+    id: number,
+    inputA: DifferenceStreamReader<[K, V1]>,
+    output: DifferenceStreamWriter<[K, V2]>,
+    f: (values: Array<[V1, number]>) => Array<[V2, number]>
+  ) {
+    super(id, inputA, output)
+    this.#f = f
+  }
+
+  run(): void {
+    // Collect all input messages and update the index
+    const keysTodo = new Set<K>()
+    for (const message of this.inputMessages()) {
+      for (const [item, multiplicity] of message.getInner()) {
+        const [key, value] = item
+        this.#index.addValue(key, [value, multiplicity])
+        keysTodo.add(key)
+      }
+    }
+
+    // For each key, compute the reduction and delta
+    const result: Array<[[K, V2], number]> = []
+    for (const key of keysTodo) {
+      const curr = this.#index.get(key)
+      const currOut = this.#indexOut.get(key)
+      const out = this.#f(curr)
+
+      // Create maps for current and previous outputs using values directly as keys
+      const newOutputMap = new Map<V2, number>()
+      const oldOutputMap = new Map<V2, number>()
+
+      // Process new output
+      for (const [value, multiplicity] of out) {
+        const existing = newOutputMap.get(value) ?? 0
+        newOutputMap.set(value, existing + multiplicity)
+      }
+
+      // Process previous output
+      for (const [value, multiplicity] of currOut) {
+        const existing = oldOutputMap.get(value) ?? 0
+        oldOutputMap.set(value, existing + multiplicity)
+      }
+
+      // First, emit removals for old values that are no longer present
+      for (const [value, multiplicity] of oldOutputMap) {
+        if (!newOutputMap.has(value)) {
+          // Remove the old value entirely
+          result.push([[key, value], -multiplicity])
+          this.#indexOut.addValue(key, [value, -multiplicity])
+        }
+      }
+
+      // Then, emit additions for new values that are not present in old
+      for (const [value, multiplicity] of newOutputMap) {
+        if (!oldOutputMap.has(value)) {
+          // Add the new value only if it has non-zero multiplicity
+          if (multiplicity !== 0) {
+            result.push([[key, value], multiplicity])
+            this.#indexOut.addValue(key, [value, multiplicity])
+          }
+        }
+      }
+
+      // Finally, emit multiplicity changes for values that were present and are still present
+      for (const [value, newMultiplicity] of newOutputMap) {
+        const oldMultiplicity = oldOutputMap.get(value)
+        if (oldMultiplicity !== undefined) {
+          const delta = newMultiplicity - oldMultiplicity
+          // Only emit actual changes, i.e. non-zero deltas
+          if (delta !== 0) {
+            result.push([[key, value], delta])
+            this.#indexOut.addValue(key, [value, delta])
+          }
+        }
+      }
+    }
+
+    if (result.length > 0) {
+      this.output.sendData(new MultiSet(result))
+    }
+  }
+}
+
+/**
+ * Reduces the elements in the stream by key (version-free)
+ */
+export function reduce<
+  KType extends T extends KeyValue<infer K, infer _V> ? K : never,
+  V1Type extends T extends KeyValue<KType, infer V> ? V : never,
+  R,
+  T,
+>(f: (values: Array<[V1Type, number]>) => Array<[R, number]>) {
+  return (stream: IStreamBuilder<T>): IStreamBuilder<KeyValue<KType, R>> => {
+    const output = new StreamBuilder<KeyValue<KType, R>>(
+      stream.graph,
+      new DifferenceStreamWriter<KeyValue<KType, R>>()
+    )
+    const operator = new ReduceOperator<KType, V1Type, R>(
+      stream.graph.getNextOperatorId(),
+      stream.connectReader() as DifferenceStreamReader<KeyValue<KType, V1Type>>,
+      output.writer,
+      f
+    )
+    stream.graph.addOperator(operator)
+    stream.graph.addStream(output.connectReader())
+    return output
+  }
+}

package/src/operators/topK.ts

@@ -0,0 +1,92 @@
+import { MultiSet } from "../multiset.js"
+import { reduce } from "./reduce.js"
+import type { IStreamBuilder, PipedOperator } from "../types"
+import type { KeyValue } from "../types.js"
+
+interface TopKOptions {
+  limit?: number
+  offset?: number
+}
+
+/**
+ * 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 withing 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.
+ *
+ * @param comparator - A function that compares two elements
+ * @param options - An optional object containing limit and offset properties
+ * @returns A piped operator that limits the number of results
+ */
+export function topK<
+  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,
+  options?: TopKOptions
+): PipedOperator<T, T> {
+  const limit = options?.limit ?? Infinity
+  const offset = options?.offset ?? 0
+
+  return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
+    const reduced = stream.pipe(
+      reduce((values) => {
+        // `values` is a list of tuples, first element is the value, second is the multiplicity
+        const consolidated = new MultiSet(values).consolidate()
+        const sortedValues = consolidated
+          .getInner()
+          .sort((a, b) => comparator(a[0] as V1Type, b[0] as V1Type))
+        return sortedValues.slice(offset, offset + limit)
+      })
+    )
+    return reduced as IStreamBuilder<T>
+  }
+}
+
+/**
+ * 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 withing 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.
+ * Adds the index of the element to the result as [key, [value, index]]
+ *
+ * @param comparator - A function that compares two elements
+ * @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 topKWithIndex<
+  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,
+  options?: TopKOptions
+): PipedOperator<T, KeyValue<KType, [V1Type, number]>> {
+  const limit = options?.limit ?? Infinity
+  const offset = options?.offset ?? 0
+
+  return (
+    stream: IStreamBuilder<T>
+  ): IStreamBuilder<KeyValue<KType, [V1Type, number]>> => {
+    const reduced = stream.pipe(
+      reduce<KType, V1Type, [V1Type, number], T>((values) => {
+        // `values` is a list of tuples, first element is the value, second is the multiplicity
+        const consolidated = new MultiSet(values).consolidate()
+        let i = offset
+        const sortedValues = consolidated
+          .getInner()
+          .sort((a, b) => comparator(a[0], b[0]))
+          .slice(offset, offset + limit)
+          .map(([value, multiplicity]): [[V1Type, number], number] => [
+            [value, i++],
+            multiplicity,
+          ])
+        return sortedValues
+      })
+    )
+    return reduced
+  }
+}