Efficient distinct operator (#80)

* Replace distinct operator atop reduce by a dedicated more efficient distinct operator

* Add additional unit tests for distinct operator

* Fix bug where distinct wasn't properly summing the multiplicites if an element occurs multiple times in the input stream

* Extend distinct with optional argument to determine what to deduplicate by.

* Formatting

* Small change to test

* changest

---------

Co-authored-by: Sam Willis <[email protected]>

Author: kevin-dp

.changeset/blue-papayas-try.md

@@ -0,0 +1,5 @@
+---
+'@electric-sql/d2mini': patch
+---
+
+rebuild the distinct to be more efficient

packages/d2mini/src/operators/distinct.ts

@@ -1,52 +1,98 @@
-import { IStreamBuilder, KeyValue } from '../types.js'
-import { DifferenceStreamReader, DifferenceStreamWriter } from '../graph.js'
+import { IStreamBuilder } from '../types.js'
+import {
+  DifferenceStreamReader,
+  DifferenceStreamWriter,
+  UnaryOperator,
+} from '../graph.js'
 import { StreamBuilder } from '../d2.js'
-import { ReduceOperator } from './reduce.js'
 import { hash } from '../utils.js'
+import { MultiSet } from '../multiset.js'
+
+type HashedValue = string
+type Multiplicity = number
 
 /**
- * Operator that removes duplicates by key (version-free)
+ * Operator that removes duplicates
  */
-export class DistinctOperator<K, V> extends ReduceOperator<K, V, V> {
+export class DistinctOperator<T> extends UnaryOperator<T> {
+  #by: (value: T) => any
+  #values: Map<HashedValue, Multiplicity> // keeps track of the number of times each value has been seen
+
   constructor(
     id: number,
-    inputA: DifferenceStreamReader<[K, V]>,
-    output: DifferenceStreamWriter<[K, V]>,
+    input: DifferenceStreamReader<T>,
+    output: DifferenceStreamWriter<T>,
+    by: (value: T) => any = (value: T) => value,
   ) {
-    const distinctInner = (vals: [V, number][]): [V, number][] => {
-      const consolidated = new Map<string, number>()
-      const values = new Map<string, V>()
-      for (const [val, diff] of vals) {
-        const key = hash(val)
-        consolidated.set(key, (consolidated.get(key) || 0) + diff)
-        values.set(key, val)
+    super(id, input, output)
+    this.#by = by
+    this.#values = new Map()
+  }
+
+  run(): void {
+    const updatedValues = new Map<HashedValue, [Multiplicity, T]>()
+
+    // Compute the new multiplicity for each value
+    for (const message of this.inputMessages()) {
+      for (const [value, diff] of message.getInner()) {
+        const hashedValue = hash(this.#by(value))
+
+        const oldMultiplicity =
+          updatedValues.get(hashedValue)?.[0] ??
+          this.#values.get(hashedValue) ??
+          0
+        const newMultiplicity = oldMultiplicity + diff
+
+        updatedValues.set(hashedValue, [newMultiplicity, value])
       }
-      return Array.from(consolidated.entries())
-        .filter(([_, count]) => count > 0)
-        .map(([key, _]) => [values.get(key) as V, 1])
     }
 
-    super(id, inputA, output, distinctInner)
+    const result: Array<[T, number]> = []
+
+    // Check which values became visible or disappeared
+    for (const [
+      hashedValue,
+      [newMultiplicity, value],
+    ] of updatedValues.entries()) {
+      const oldMultiplicity = this.#values.get(hashedValue) ?? 0
+
+      if (newMultiplicity === 0) {
+        this.#values.delete(hashedValue)
+      } else {
+        this.#values.set(hashedValue, newMultiplicity)
+      }
+
+      if (oldMultiplicity <= 0 && newMultiplicity > 0) {
+        // The value wasn't present in the stream
+        // but with this change it is now present in the stream
+        result.push([value, 1])
+      } else if (oldMultiplicity > 0 && newMultiplicity <= 0) {
+        // The value was present in the stream
+        // but with this change it is no longer present in the stream
+        result.push([value, -1])
+      }
+    }
+
+    if (result.length > 0) {
+      this.output.sendData(new MultiSet(result))
+    }
   }
 }
 
 /**
- * Removes duplicates by key (version-free)
+ * Removes duplicate values
  */
-export function distinct<
-  K extends T extends KeyValue<infer K, infer _V> ? K : never,
-  V extends T extends KeyValue<K, infer V> ? V : never,
-  T,
->() {
-  return (stream: IStreamBuilder<T>): IStreamBuilder<KeyValue<K, V>> => {
-    const output = new StreamBuilder<KeyValue<K, V>>(
+export function distinct<T>(by: (value: T) => any = (value: T) => value) {
+  return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
+    const output = new StreamBuilder<T>(
       stream.graph,
-      new DifferenceStreamWriter<KeyValue<K, V>>(),
+      new DifferenceStreamWriter<T>(),
     )
-    const operator = new DistinctOperator<K, V>(
+    const operator = new DistinctOperator<T>(
       stream.graph.getNextOperatorId(),
-      stream.connectReader() as DifferenceStreamReader<KeyValue<K, V>>,
+      stream.connectReader() as DifferenceStreamReader<T>,
       output.writer,
+      by,
     )
     stream.graph.addOperator(operator)
     stream.graph.addStream(output.connectReader())

packages/d2mini/tests/operators/distinct.test.ts

@@ -5,7 +5,7 @@ import { distinct } from '../../src/operators/distinct.js'
 import { output } from '../../src/operators/output.js'
 
 describe('Operators', () => {
-  describe('Distinct operation', () => {
+  describe('Efficient distinct operation', () => {
     testDistinct()
   })
 })
@@ -45,6 +45,47 @@ function testDistinct() {
     ])
   })
 
+  test('distinct by certain property', () => {
+    const graph = new D2()
+    const input = graph.newInput<[number, { name: string; country: string }]>()
+    const messages: MultiSet<[number, { name: string; country: string }]>[] = []
+
+    input.pipe(
+      distinct(([_, value]) => value.country),
+      output((message) => {
+        messages.push(message)
+      }),
+    )
+
+    graph.finalize()
+
+    input.sendData(
+      new MultiSet([
+        [[1, { name: 'Valter', country: 'Portugal' }], 1],
+        [[2, { name: 'Sam', country: 'UK' }], 1],
+        [[2, { name: 'Kevin', country: 'Belgium' }], 1],
+        [[3, { name: 'Garry', country: 'UK' }], 1],
+        [[4, { name: 'Kyle', country: 'USA' }], 1],
+      ]),
+    )
+
+    graph.run()
+
+    const data = messages.map((m) => m.getInner())[0]
+    const countries = data
+      .map(([[_, value], multiplicity]) => [value.country, multiplicity])
+      .sort()
+
+    expect(countries).toEqual(
+      [
+        ['Belgium', 1],
+        ['Portugal', 1],
+        ['UK', 1],
+        ['USA', 1],
+      ].sort(),
+    )
+  })
+
   test('distinct with updates', () => {
     const graph = new D2()
     const input = graph.newInput<[number, string]>()
@@ -63,18 +104,23 @@ function testDistinct() {
       new MultiSet([
         [[1, 'a'], 1],
         [[1, 'b'], 1],
+        [[1, 'a'], 1],
       ]),
     )
     graph.run()
 
     input.sendData(
       new MultiSet([
         [[1, 'b'], -1],
-        [[1, 'c'], 1],
+        [[1, 'c'], 2],
+        [[1, 'a'], -1],
       ]),
     )
     graph.run()
 
+    input.sendData(new MultiSet([[[1, 'c'], -2]]))
+    graph.run()
+
     const data = messages.map((m) => m.getInner())
 
     expect(data).toEqual([
@@ -86,6 +132,7 @@ function testDistinct() {
         [[1, 'b'], -1],
         [[1, 'c'], 1],
       ],
+      [[[1, 'c'], -1]],
     ])
   })
 
@@ -122,4 +169,40 @@ function testDistinct() {
       ],
     ])
   })
+
+  test('distinct with multiple batches of same key that cancel out', () => {
+    const graph = new D2()
+    const input = graph.newInput<[string, number]>()
+    const messages: MultiSet<[string, number]>[] = []
+
+    input.pipe(
+      distinct(),
+      output((message) => {
+        messages.push(message)
+      }),
+    )
+
+    graph.finalize()
+
+    input.sendData(
+      new MultiSet([
+        [['key1', 1], 2],
+        [['key1', 2], 2],
+        [['key1', 2], 1],
+        [['key2', 1], 1],
+        [['key1', 2], -3], // cancels out the previous addition of [['key2', 2], 3]
+        [['key2', 1], 1],
+      ]),
+    )
+    graph.run()
+
+    const data = messages.map((m) => m.getInner())
+
+    expect(data).toEqual([
+      [
+        [['key1', 1], 1],
+        [['key2', 1], 1],
+      ],
+    ])
+  })
 }