Incremental topK with fractional index (electric-sql#72)

* WIP incremental topKWithFractionalIndex

* Incremental topKWithFractionalIndex

* Fix tests to not assume particular fractional indices as those are implementation details.

* Introduce a TopK data structure

* B+ tree variant of topKWithFractionalIndex

* Extend unit tests to test all insertion and deletion cases

* Formatting

* Unit test for duplicate values

* Expose useTree option also on sortBy operator

* Split array and B+ tree variants in separate operators

* Add missing imports

* Add a orderBy operator that uses topK with B+ tree variant

* Formatting

* Trigger CI

* Remove useTree option

* Changeset

* Dynamically import B+ tree library

* Also run orderByWithFractionalIndex tests for both the array and B+ tree variant

* Split tree version of orderBy into its own file to enable tree shaking the sorted-btree dependency.

* Improved changeset

Author: kevin-dp

.changeset/chilly-icons-design.md

@@ -0,0 +1,5 @@
+---
+'@electric-sql/d2mini': patch
+---
+
+Introduce topKWithFractionalIndexBTree and orderByWithFractionalIndexBTree operators. These variants use a B+ tree which is more efficient on big collections as its time complexity is logarithmic.

packages/d2mini/package.json

@@ -50,6 +50,7 @@
   },
   "dependencies": {
     "fractional-indexing": "^3.2.0",
-    "murmurhash-js": "^1.0.0"
+    "murmurhash-js": "^1.0.0",
+    "sorted-btree": "^1.8.1"
   }
 }

packages/d2mini/src/indexes.ts

@@ -35,6 +35,13 @@ export class Index<K, V> {
     return [...valueMap.values()]
   }
 
+  getMultiplicity(key: K, value: V): number {
+    const valueMap = this.#inner.get(key)
+    const valueHash = hash(value)
+    const [, multiplicity] = valueMap.get(valueHash)
+    return multiplicity
+  }
+
   entries() {
     return this.#inner.entries()
   }

packages/d2mini/src/operators/orderBy.ts

@@ -6,7 +6,7 @@ import { map } from './map.js'
 import { innerJoin } from './join.js'
 import { consolidate } from './consolidate.js'
 
-interface OrderByOptions<Ve> {
+export interface OrderByOptions<Ve> {
   comparator?: (a: Ve, b: Ve) => number
   limit?: number
   offset?: number
@@ -128,19 +128,11 @@ export function orderByWithIndex<
   }
 }
 
-/**
- * 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<
+export function orderByWithFractionalIndexBase<
   T extends KeyValue<unknown, unknown>,
   Ve = unknown,
 >(
+  topK: typeof topKWithFractionalIndex,
   valueExtractor: (
     value: T extends KeyValue<unknown, infer V> ? V : never,
   ) => Ve,
@@ -181,7 +173,7 @@ export function orderByWithFractionalIndex<
             ],
           ] as KeyValue<null, [K, Ve]>,
       ),
-      topKWithFractionalIndex((a, b) => comparator(a[1], b[1]), {
+      topK((a, b) => comparator(a[1], b[1]), {
         limit,
         offset,
       }),
@@ -194,3 +186,28 @@ export function orderByWithFractionalIndex<
     )
   }
 }
+
+/**
+ * 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,
+  )
+}

packages/d2mini/src/operators/orderByBTree.ts

@@ -0,0 +1,19 @@
+import { KeyValue } from '../types.js'
+import { OrderByOptions, orderByWithFractionalIndexBase } from './orderBy.js'
+import { topKWithFractionalIndexBTree } from './topKWithFractionalIndexBTree.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,
+  )
+}