Improve performance of partitionByKey with a big number of elements (#232)

* Improve performance of partitionByKey with a big number of elements
* Improve performance of partitionByKey by not serialising the keys
* partition by key to emit deltas instead of the resulting keys
* update tests for partitionByKey
* update tests for combineKeys
* getGroupedObservable: close outer subscription once the inner stream is found
* add performance tests to partitionByKey
* make tests more predictable
* adjust performance threshold

Co-authored-by: Josep M Sobrepere <jm.sobrepere@gmail.com>

Author: voliva

packages/utils/package.json

@@ -30,7 +30,7 @@
     "build:cjs:dev": "esbuild src/index.tsx --bundle --outfile=./dist/utils.cjs.development.js --target=es2015 --external:react --external:rxjs --external:@react-rxjs/core --format=cjs --sourcemap",
     "build:cjs:prod": "esbuild src/index.tsx --bundle --outfile=./dist/utils.cjs.production.min.js --target=es2015 --external:react --external:rxjs --external:@react-rxjs/core --format=cjs --minify --sourcemap",
     "build:ts": "tsc -p ./tsconfig-build.json --outDir ./dist --skipLibCheck --emitDeclarationOnly",
-    "test": "jest --coverage",
+    "test": "node --expose-gc ../../node_modules/.bin/jest --coverage --runInBand",
     "lint": "prettier --check README.md \"src/**/*.{js,jsx,ts,tsx,json,md}\"",
     "format": "prettier --write README.md \"src/**/*.{js,jsx,ts,tsx,json,md}\"",
     "prepare": "yarn build"

packages/utils/src/combineKeys.test.ts

@@ -1,7 +1,7 @@
 import { concat, NEVER, Observable, of } from "rxjs"
 import { map, scan } from "rxjs/operators"
 import { TestScheduler } from "rxjs/testing"
-import { combineKeys } from "./"
+import { combineKeys, KeyChanges } from "./"
 
 const scheduler = () =>
   new TestScheduler((actual, expected) => {
@@ -288,4 +288,132 @@ describe("combineKeys", () => {
       })
     })
   })
+
+  describe("with deltas", () => {
+    it("emits a map with the latest value of the stream of each key added", () => {
+      scheduler().run(({ expectObservable, cold }) => {
+        const keys = cold<KeyChanges<string>>("  ab---cd---", {
+          a: {
+            type: "add",
+            keys: ["a"],
+          },
+          b: {
+            type: "add",
+            keys: ["b"],
+          },
+          c: {
+            type: "add",
+            keys: ["c"],
+          },
+          d: {
+            type: "add",
+            keys: ["d"],
+          },
+        })
+        const a = cold("     --1---2---")
+        const b = cold("      ---------")
+        const c = cold("          1----")
+        const d = cold("           9---")
+        const expectedStr = "--e--f(gh)"
+
+        const innerStreams: Record<string, Observable<string>> = { a, b, c, d }
+
+        const result = combineKeys(
+          keys,
+          (v): Observable<string> => innerStreams[v],
+        ).pipe(map((x) => Object.fromEntries(x.entries())))
+
+        expectObservable(result).toBe(expectedStr, {
+          e: { a: "1" },
+          f: { a: "1", c: "1" },
+          g: { a: "2", c: "1" },
+          h: { a: "2", c: "1", d: "9" },
+        })
+      })
+    })
+
+    it("removes the entry of a key when that key is removed", () => {
+      scheduler().run(({ expectObservable, cold }) => {
+        const expectedStr = "                  e-f-g-h-"
+        const keys = cold<KeyChanges<string>>("a-b-c-d-", {
+          a: {
+            type: "add",
+            keys: ["a"],
+          },
+          b: {
+            type: "add",
+            keys: ["b"],
+          },
+          c: {
+            type: "remove",
+            keys: ["a"],
+          },
+          d: {
+            type: "remove",
+            keys: ["b"],
+          },
+        })
+
+        const result = combineKeys(keys, (v): Observable<string> => of(v)).pipe(
+          map((x) => Object.fromEntries(x.entries())),
+        )
+
+        expectObservable(result).toBe(expectedStr, {
+          e: { a: "a" },
+          f: { a: "a", b: "b" },
+          g: { b: "b" },
+          h: {},
+        })
+      })
+    })
+
+    it("accepts more than one key change at a time", () => {
+      scheduler().run(({ expectObservable, cold }) => {
+        const expectedStr = "                  e-f-"
+        const keys = cold<KeyChanges<string>>("a-b-", {
+          a: {
+            type: "add",
+            keys: ["a", "b", "c"],
+          },
+          b: {
+            type: "remove",
+            keys: ["c", "b"],
+          },
+        })
+
+        const result = combineKeys(keys, (v): Observable<string> => of(v)).pipe(
+          map((x) => Object.fromEntries(x.entries())),
+        )
+
+        expectObservable(result).toBe(expectedStr, {
+          e: { a: "a", b: "b", c: "c" },
+          f: { a: "a" },
+        })
+      })
+    })
+
+    it("omits removing keys that don't exist", () => {
+      scheduler().run(({ expectObservable, cold }) => {
+        const expectedStr = "                  e---"
+        const keys = cold<KeyChanges<string>>("a-b-", {
+          a: {
+            type: "add",
+            keys: ["a"],
+          },
+          b: {
+            type: "remove",
+            keys: ["b"],
+          },
+        })
+
+        const result = combineKeys(keys, (v): Observable<string> => of(v)).pipe(
+          map((x) => Object.fromEntries(x.entries())),
+        )
+
+        expectObservable(result).toBe(expectedStr, {
+          e: { a: "a" },
+        })
+      })
+    })
+  })
 })

packages/utils/src/combineKeys.ts

@@ -1,3 +1,4 @@
+import { KeyChanges } from "./partitionByKey"
 import { Observable, Subscription } from "rxjs"
 
 export interface MapWithChanges<K, V> extends Map<K, V> {
@@ -12,7 +13,7 @@ export interface MapWithChanges<K, V> extends Map<K, V> {
  * @returns An stream with a map containing the latest value from the stream of each key.
  */
 export const combineKeys = <K, T>(
-  keys$: Observable<Array<K> | Set<K>>,
+  keys$: Observable<Iterable<K> | KeyChanges<K>>,
   getInner$: (key: K) => Observable<T>,
 ): Observable<MapWithChanges<K, T>> =>
   new Observable((observer) => {
@@ -24,7 +25,7 @@ export const combineKeys = <K, T>(
 
     const next = () => {
       if (!updatingSource) {
-        const result = Object.assign(new Map(currentValue), {
+        const result = Object.assign(currentValue, {
           changes,
         })
         changes = new Set<K>()
@@ -36,20 +37,42 @@ export const combineKeys = <K, T>(
     const subscription = keys$.subscribe(
       (nextKeysArr) => {
         updatingSource = true
-        const nextKeys = new Set(nextKeysArr)
-        innerSubscriptions.forEach((sub, key) => {
-          if (!nextKeys.has(key)) {
-            sub.unsubscribe()
-            innerSubscriptions.delete(key)
-            if (currentValue.has(key)) {
-              changes.add(key)
-              currentValue.delete(key)
-            }
-          } else {
-            nextKeys.delete(key)
+
+        const keys = new Set(
+          inputIsKeyChanges(nextKeysArr) ? nextKeysArr.keys : nextKeysArr,
+        )
+
+        if (inputIsKeyChanges(nextKeysArr)) {
+          if (nextKeysArr.type === "remove") {
+            keys.forEach((key) => {
+              const sub = innerSubscriptions.get(key)
+              if (!sub) return
+              sub.unsubscribe()
+              innerSubscriptions.delete(key)
+              if (currentValue.has(key)) {
+                changes.add(key)
+                currentValue.delete(key)
+              }
+            })
+            // Keys after this block is the list of keys to add. Clear it.
+            keys.clear()
           }
-        })
-        nextKeys.forEach((key) => {
+        } else {
+          innerSubscriptions.forEach((sub, key) => {
+            if (!keys.has(key)) {
+              sub.unsubscribe()
+              innerSubscriptions.delete(key)
+              if (currentValue.has(key)) {
+                changes.add(key)
+                currentValue.delete(key)
+              }
+            } else {
+              keys.delete(key)
+            }
+          })
+        }
+
+        keys.forEach((key) => {
           innerSubscriptions.set(
             key,
             getInner$(key).subscribe(
@@ -66,12 +89,13 @@ export const combineKeys = <K, T>(
             ),
           )
         })
+
         updatingSource = false
         // If there are no changes but the nextKeys are an empty iterator
         // and we have never emitted before, that means that the first
         // value that keys$ has emitted is an empty iterator, therefore
         // we should emit an empy Map
-        if (changes.size || (isPristine && !nextKeys.size)) next()
+        if (changes.size || (isPristine && !keys.size)) next()
       },
       (e) => {
         observer.error(e)
@@ -88,3 +112,9 @@ export const combineKeys = <K, T>(
       })
     }
   })
+
+function inputIsKeyChanges<K>(
+  input: Iterable<K> | KeyChanges<K>,
+): input is KeyChanges<K> {
+  return "type" in input && "keys" in input
+}

packages/utils/src/index.tsx

@@ -2,7 +2,7 @@ export { combineKeys, MapWithChanges } from "./combineKeys"
 export { createSignal } from "./createSignal"
 export { createKeyedSignal } from "./createKeyedSignal"
 export { mergeWithKey } from "./mergeWithKey"
-export { partitionByKey } from "./partitionByKey"
+export { partitionByKey, KeyChanges } from "./partitionByKey"
 export { suspend } from "./suspend"
 export { suspended } from "./suspended"
 export { switchMapSuspended } from "./switchMapSuspended"