fixes and formating

Author: samwillis

packages/d2mini/src/indexes.ts

@@ -10,9 +10,7 @@ export class Index<K, V> {
   #inner: DefaultMap<K, Map<V, number>>
 
   constructor() {
-    this.#inner = new DefaultMap<K, Map<V, number>>(
-      () => new Map<V, number>(),
-    )
+    this.#inner = new DefaultMap<K, Map<V, number>>(() => new Map<V, number>())
     // #inner is a map of:
     // {
     //   [key]: Map<V, number>  // Direct value-to-multiplicity mapping
@@ -58,7 +56,7 @@ export class Index<K, V> {
     const valueMap = this.#inner.get(key)
     const existingMultiplicity = valueMap.get(val) ?? 0
     const newMultiplicity = existingMultiplicity + multiplicity
-    
+
     if (multiplicity !== 0) {
       if (newMultiplicity === 0) {
         valueMap.delete(val)

packages/d2mini/src/multiset.ts

@@ -1,4 +1,9 @@
-import { DefaultMap, chunkedArrayPush, hash } from './utils.js'
+import {
+  DefaultMap,
+  chunkedArrayPush,
+  hash,
+  globalObjectIdGenerator,
+} from './utils.js'
 
 export type MultiSetArray<T> = [T, number][]
 export type KeyedData<T> = [key: string, value: T]
@@ -66,6 +71,101 @@ export class MultiSet<T> {
    * (record, multiplicity) pair.
    */
   consolidate(): MultiSet<T> {
+    // Check if this looks like a keyed multiset (first item is a tuple of length 2)
+    if (this.#inner.length > 0) {
+      const firstItem = this.#inner[0][0]
+      if (Array.isArray(firstItem) && firstItem.length === 2) {
+        return this.#consolidateKeyed()
+      }
+    }
+
+    // Fall back to original method for unkeyed data
+    return this.#consolidateUnkeyed()
+  }
+
+  /**
+   * Private method for consolidating keyed multisets where keys are strings/numbers
+   * and values are compared by reference equality.
+   *
+   * This method provides significant performance improvements over the hash-based approach
+   * by using WeakMap for object reference tracking and avoiding expensive serialization.
+   *
+   * Special handling for join operations: When values are tuples of length 2 (common in joins),
+   * we unpack them and compare each element individually to maintain proper equality semantics.
+   */
+  #consolidateKeyed(): MultiSet<T> {
+    const consolidated = new Map<string, number>()
+    const values = new Map<string, T>()
+
+    // Use global object ID generator for consistent reference equality
+
+    /**
+     * Special handler for tuples (arrays of length 2) commonly produced by join operations.
+     * Unpacks the tuple and generates an ID based on both elements to ensure proper
+     * consolidation of join results like ['A', null] and [null, 'X'].
+     */
+    const getTupleId = (tuple: any[]): string => {
+      if (tuple.length !== 2) {
+        throw new Error('Expected tuple of length 2')
+      }
+      const [first, second] = tuple
+      return `${globalObjectIdGenerator.getStringId(first)}|${globalObjectIdGenerator.getStringId(second)}`
+    }
+
+    // Process each item in the multiset
+    for (const [data, multiplicity] of this.#inner) {
+      // Verify this is still a keyed item (should be [key, value] pair)
+      if (!Array.isArray(data) || data.length !== 2) {
+        // Found non-keyed item, fall back to unkeyed consolidation
+        return this.#consolidateUnkeyed()
+      }
+
+      const [key, value] = data
+
+      // Verify key is string or number as expected for keyed multisets
+      if (typeof key !== 'string' && typeof key !== 'number') {
+        // Found non-string/number key, fall back to unkeyed consolidation
+        return this.#consolidateUnkeyed()
+      }
+
+      // Generate value ID with special handling for join tuples
+      let valueId: string
+      if (Array.isArray(value) && value.length === 2) {
+        // Special case: value is a tuple from join operations
+        valueId = getTupleId(value)
+      } else {
+        // Regular case: use reference/value equality
+        valueId = globalObjectIdGenerator.getStringId(value)
+      }
+
+      // Create composite key and consolidate
+      const compositeKey = key + '|' + valueId
+      consolidated.set(
+        compositeKey,
+        (consolidated.get(compositeKey) || 0) + multiplicity,
+      )
+
+      // Store the original data for the first occurrence
+      if (!values.has(compositeKey)) {
+        values.set(compositeKey, data as T)
+      }
+    }
+
+    // Build result array, filtering out zero multiplicities
+    const result: MultiSetArray<T> = []
+    for (const [compositeKey, multiplicity] of consolidated) {
+      if (multiplicity !== 0) {
+        result.push([values.get(compositeKey)!, multiplicity])
+      }
+    }
+
+    return new MultiSet(result)
+  }
+
+  /**
+   * Private method for consolidating unkeyed multisets using the original approach.
+   */
+  #consolidateUnkeyed(): MultiSet<T> {
     const consolidated = new DefaultMap<string | number, number>(() => 0)
     const values = new Map<string, any>()
 

packages/d2mini/src/operators/topKWithFractionalIndex.ts

@@ -9,6 +9,7 @@ import { MultiSet } from '../multiset.js'
 import { Index } from '../indexes.js'
 import { generateKeyBetween } from 'fractional-indexing'
 import { binarySearch } from '../utils.js'
+import { globalObjectIdGenerator } from '../utils.js'
 
 export interface TopKWithFractionalIndexOptions {
   limit?: number
@@ -203,7 +204,10 @@ export class TopKWithFractionalIndexOperator<K, V1> extends UnaryOperator<
   protected createTopK(
     offset: number,
     limit: number,
-    comparator: (a: TieBreakerTaggedValue<V1>, b: TieBreakerTaggedValue<V1>) => number,
+    comparator: (
+      a: TieBreakerTaggedValue<V1>,
+      b: TieBreakerTaggedValue<V1>,
+    ) => number,
   ): TopK<TieBreakerTaggedValue<V1>> {
     return new TopKArray(offset, limit, comparator)
   }
@@ -232,7 +236,10 @@ export class TopKWithFractionalIndexOperator<K, V1> extends UnaryOperator<
     this.#index.addValue(key, [value, multiplicity])
     const newMultiplicity = this.#index.getMultiplicity(key, value)
 
-    let res: TopKChanges<TieBreakerTaggedValue<V1>> = { moveIn: null, moveOut: null }
+    let res: TopKChanges<TieBreakerTaggedValue<V1>> = {
+      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
@@ -334,43 +341,19 @@ function mapValue<V, W>(
   return [f(getValue(value)), getIndex(value)]
 }
 
-  // Abstraction for values tagged with a tie breaker
-// Object identity-based tie-breaking using WeakMap
-const objectIds = new WeakMap<object, number>()
-let nextObjectId = 0
-
-function getObjectId(value: any): number {
-  // For primitives, use a simple hash of their string representation
-  if (typeof value !== 'object' || value === null) {
-    // Simple string-based hash for primitives to ensure consistency
-    const str = String(value)
-    let hash = 0
-    for (let i = 0; i < str.length; i++) {
-      const char = str.charCodeAt(i)
-      hash = ((hash << 5) - hash) + char
-      hash = hash & hash // Convert to 32-bit integer
-    }
-    return hash
-  }
-  
-  // For objects, use WeakMap to assign unique IDs
-  if (!objectIds.has(value)) {
-    objectIds.set(value, nextObjectId++)
-  }
-  return objectIds.get(value)!
-}
-
 export type TieBreaker = number
 export type TieBreakerTaggedValue<V> = [V, TieBreaker]
 
 function tagValue<V>(value: V): TieBreakerTaggedValue<V> {
-  return [value, getObjectId(value)]
+  return [value, globalObjectIdGenerator.getId(value)]
 }
 
 function untagValue<V>(tieBreakerTaggedValue: TieBreakerTaggedValue<V>): V {
   return tieBreakerTaggedValue[0]
 }
 
-function getTieBreaker<V>(tieBreakerTaggedValue: TieBreakerTaggedValue<V>): TieBreaker {
+function getTieBreaker<V>(
+  tieBreakerTaggedValue: TieBreakerTaggedValue<V>,
+): TieBreaker {
   return tieBreakerTaggedValue[1]
 }

packages/d2mini/src/operators/topKWithFractionalIndexBTree.ts

@@ -240,7 +240,10 @@ export class TopKWithFractionalIndexBTreeOperator<
   protected override createTopK(
     offset: number,
     limit: number,
-    comparator: (a: TieBreakerTaggedValue<V1>, b: TieBreakerTaggedValue<V1>) => number,
+    comparator: (
+      a: TieBreakerTaggedValue<V1>,
+      b: TieBreakerTaggedValue<V1>,
+    ) => number,
   ): TopK<TieBreakerTaggedValue<V1>> {
     if (!BTree) {
       throw new Error(

packages/d2mini/src/utils.ts

@@ -111,3 +111,53 @@ export function binarySearch<T>(
   }
   return low
 }
+
+/**
+ * Utility for generating unique IDs for objects and values.
+ * Uses WeakMap for object reference tracking and consistent hashing for primitives.
+ */
+export class ObjectIdGenerator {
+  private objectIds = new WeakMap<object, number>()
+  private nextId = 0
+
+  /**
+   * Get a unique identifier for any value.
+   * - Objects: Uses WeakMap for reference-based identity
+   * - Primitives: Uses consistent string-based hashing
+   */
+  getId(value: any): number {
+    // For primitives, use a simple hash of their string representation
+    if (typeof value !== 'object' || value === null) {
+      const str = String(value)
+      let hash = 0
+      for (let i = 0; i < str.length; i++) {
+        const char = str.charCodeAt(i)
+        hash = (hash << 5) - hash + char
+        hash = hash & hash // Convert to 32-bit integer
+      }
+      return hash
+    }
+
+    // For objects, use WeakMap to assign unique IDs
+    if (!this.objectIds.has(value)) {
+      this.objectIds.set(value, this.nextId++)
+    }
+    return this.objectIds.get(value)!
+  }
+
+  /**
+   * Get a string representation of the ID for use in composite keys.
+   */
+  getStringId(value: any): string {
+    if (value === null) return 'null'
+    if (value === undefined) return 'undefined'
+    if (typeof value !== 'object') return String(value)
+
+    return `obj_${this.getId(value)}`
+  }
+}
+
+/**
+ * Global instance for cases where a shared object ID space is needed.
+ */
+export const globalObjectIdGenerator = new ObjectIdGenerator()

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

@@ -3,7 +3,11 @@ import { D2 } from '../../src/d2.js'
 import { MultiSet } from '../../src/multiset.js'
 import { count } from '../../src/operators/count.js'
 import { output } from '../../src/operators/output.js'
-import { KeyedMessageTracker, assertKeyedResults, assertOnlyKeysAffected } from '../test-utils.js'
+import {
+  KeyedMessageTracker,
+  assertKeyedResults,
+  assertOnlyKeysAffected,
+} from '../test-utils.js'
 
 describe('Operators', () => {
   describe('Count operation', () => {
@@ -40,10 +44,10 @@ function testCount() {
     graph.run()
 
     const result = tracker.getResult()
-    
+
     // Assert only keys that have values are affected
     assertOnlyKeysAffected('basic count operation', result.messages, [1, 2, 3])
-    
+
     // Assert the final materialized results are correct
     assertKeyedResults(
       'basic count operation',
@@ -53,7 +57,7 @@ function testCount() {
         [2, 3], // 3 values for key 2
         [3, 1], // 1 value for key 3 (1 + (-1) + 1 = 1)
       ],
-      6 // Expected message count
+      6, // Expected message count
     )
   })
 
@@ -80,18 +84,22 @@ function testCount() {
     graph.run()
 
     const result = tracker.getResult()
-    
+
     // Assert only key 1 is affected
-    assertOnlyKeysAffected('count with all negative multiplicities', result.messages, [1])
-    
+    assertOnlyKeysAffected(
+      'count with all negative multiplicities',
+      result.messages,
+      [1],
+    )
+
     // Assert the final materialized results are correct
     assertKeyedResults(
       'count with all negative multiplicities',
       result,
       [
         [1, -3], // -1 + (-2) = -3
       ],
-      2 // Expected message count
+      2, // Expected message count
     )
   })
 
@@ -126,10 +134,13 @@ function testCount() {
     graph.run()
 
     const result = tracker.getResult()
-    
+
     // Assert only keys 'one' and 'two' are affected
-    assertOnlyKeysAffected('count with multiple batches', result.messages, ['one', 'two'])
-    
+    assertOnlyKeysAffected('count with multiple batches', result.messages, [
+      'one',
+      'two',
+    ])
+
     // Assert the final materialized results are correct
     assertKeyedResults(
       'count with multiple batches',
@@ -138,7 +149,7 @@ function testCount() {
         ['one', 3], // 2 + 1 = 3
         ['two', 1], // 1
       ],
-      5 // Expected message count
+      5, // Expected message count
     )
   })
 
@@ -182,10 +193,13 @@ function testCount() {
     graph.run()
 
     const result = tracker.getResult()
-    
+
     // Assert only keys 'a' and 'c' are affected (NOT 'b')
-    assertOnlyKeysAffected('count incremental updates', result.messages, ['a', 'c'])
-    
+    assertOnlyKeysAffected('count incremental updates', result.messages, [
+      'a',
+      'c',
+    ])
+
     // Assert the final materialized results are correct
     assertKeyedResults(
       'count incremental updates',
@@ -194,7 +208,7 @@ function testCount() {
         ['a', 3], // Count increased from 2 to 3
         ['c', 0], // Count decreased from 1 to 0
       ],
-      4 // Expected message count: remove old 'a', add new 'a', remove old 'c', add new 'c'
+      4, // Expected message count: remove old 'a', add new 'a', remove old 'c', add new 'c'
     )
   })
 }