Make search results deterministic with a rowIndex tie-break

Author: platypii

src/search/heap.js

@@ -32,6 +32,22 @@ export function isBetter(a, b, metric) {
   return a > b
 }
 
+/**
+ * Total order on scored entries: better score wins; ties broken by lower
+ * rowIndex. The tiebreak makes selection independent of the order chunks
+ * arrive (parallel reads complete in nondeterministic order), so identical
+ * queries always return identical results.
+ *
+ * @param {{ rowIndex: number, score: number }} a
+ * @param {{ rowIndex: number, score: number }} b
+ * @param {DistanceMetric} metric
+ * @returns {boolean}
+ */
+function betterEntry(a, b, metric) {
+  if (a.score !== b.score) return isBetter(a.score, b.score, metric)
+  return a.rowIndex < b.rowIndex
+}
+
 /**
  * Bounded heap by score (uses metric to decide "better"). Linear-scan to
  * find the worst entry; fine for the small topK we care about.
@@ -48,17 +64,19 @@ export function pushHeap(heap, candidate, topK, metric) {
   }
   let worstIdx = 0
   for (let i = 1; i < heap.length; i += 1) {
-    if (isBetter(heap[worstIdx].score, heap[i].score, metric)) {
+    if (betterEntry(heap[worstIdx], heap[i], metric)) {
       worstIdx = i
     }
   }
-  if (isBetter(candidate.score, heap[worstIdx].score, metric)) {
+  if (betterEntry(candidate, heap[worstIdx], metric)) {
     heap[worstIdx] = candidate
   }
 }
 
 /**
- * Bounded heap for Hamming candidates (lower hamming is better).
+ * Bounded heap for Hamming candidates (lower hamming is better). Ties on
+ * hamming are broken by lower rowIndex so the kept candidate set does not
+ * depend on the (parallel, nondeterministic) order chunks are scored.
  *
  * @param {{ rowIndex: number, hamming: number }[]} heap
  * @param {{ rowIndex: number, hamming: number }} candidate
@@ -71,21 +89,24 @@ export function pushHammingHeap(heap, candidate, candidatesK) {
   }
   let worstIdx = 0
   for (let i = 1; i < heap.length; i += 1) {
-    if (heap[i].hamming > heap[worstIdx].hamming) worstIdx = i
+    const h = heap[i]; const w = heap[worstIdx]
+    if (h.hamming > w.hamming || (h.hamming === w.hamming && h.rowIndex > w.rowIndex)) worstIdx = i
   }
-  if (candidate.hamming < heap[worstIdx].hamming) {
+  const w = heap[worstIdx]
+  if (candidate.hamming < w.hamming || (candidate.hamming === w.hamming && candidate.rowIndex < w.rowIndex)) {
     heap[worstIdx] = candidate
   }
 }
 
 /**
- * Sort results best-first under the chosen metric.
+ * Sort results best-first under the chosen metric, ties broken by lower
+ * rowIndex for a deterministic order.
  *
  * @param {{ rowIndex: number, score: number }[]} results
  * @param {DistanceMetric} metric
  * @returns {{ rowIndex: number, score: number }[]}
  */
 export function sortHeap(results, metric) {
   const dir = metric === 'euclidean' ? 1 : -1
-  return results.slice().sort((a, b) => dir * (a.score - b.score))
+  return results.slice().sort((a, b) => dir * (a.score - b.score) || a.rowIndex - b.rowIndex)
 }

src/search/rerank.js

@@ -118,7 +118,9 @@ export async function searchRerank({ file, metadata, meta, queryF32, scoringMetr
   }))
 
   const dir = reportedMetric === 'euclidean' ? 1 : -1
-  scored.sort((a, b) => dir * (a.score - b.score))
+  // Tie-break by rowIndex so the winners are independent of the order the
+  // parallel candidate reads completed in (deterministic results).
+  scored.sort((a, b) => dir * (a.score - b.score) || a.rowIndex - b.rowIndex)
   const winners = scored.slice(0, topK)
 
   // Phase 3: fetch ids for just the top-K winners.

test/heap.test.js

@@ -0,0 +1,66 @@
+import { describe, expect, it } from 'vitest'
+import { pushHammingHeap, pushHeap, sortHeap } from '../src/search/heap.js'
+
+/**
+ * Run a sequence of candidates through a bounded heap and return the kept
+ * rowIndices, sorted ascending for comparison.
+ * @param {{ rowIndex: number, hamming: number }[]} candidates
+ * @param {number} k
+ * @returns {number[]}
+ */
+function keepHamming(candidates, k) {
+  /** @type {{ rowIndex: number, hamming: number }[]} */
+  const heap = []
+  for (const c of candidates) pushHammingHeap(heap, c, k)
+  return heap.map(e => e.rowIndex).sort((a, b) => a - b)
+}
+
+describe('heap tie-breaking is deterministic', () => {
+  it('pushHammingHeap keeps the lowest rowIndices when hamming ties, regardless of insertion order', () => {
+    // Five rows all at the same hamming distance; keep 3 -> must keep rows 0,1,2.
+    const rows = [0, 1, 2, 3, 4].map(rowIndex => ({ rowIndex, hamming: 7 }))
+    const forward = keepHamming(rows, 3)
+    const reversed = keepHamming([...rows].reverse(), 3)
+    const shuffled = keepHamming([rows[3], rows[0], rows[4], rows[2], rows[1]], 3)
+    expect(forward).toEqual([0, 1, 2])
+    expect(reversed).toEqual([0, 1, 2])
+    expect(shuffled).toEqual([0, 1, 2])
+  })
+
+  it('pushHammingHeap prefers strictly nearer candidates over ties', () => {
+    const rows = [
+      { rowIndex: 5, hamming: 2 },
+      { rowIndex: 9, hamming: 9 },
+      { rowIndex: 1, hamming: 9 },
+      { rowIndex: 8, hamming: 1 },
+    ]
+    // keep 2 -> the two nearest by hamming (1 then 2): rows 8 and 5.
+    expect(keepHamming(rows, 2)).toEqual([5, 8])
+    expect(keepHamming([...rows].reverse(), 2)).toEqual([5, 8])
+  })
+
+  it('pushHeap keeps best score, breaking ties by lower rowIndex (order-independent)', () => {
+    // cosine: higher score is better. Tied at 0.9 -> keep lower rowIndices.
+    const rows = [0, 1, 2, 3].map(rowIndex => ({ rowIndex, score: 0.9 }))
+    /** @type {{ rowIndex: number, score: number }[]} */
+    const a = []
+    for (const c of rows) pushHeap(a, c, 2, 'cosine')
+    /** @type {{ rowIndex: number, score: number }[]} */
+    const b = []
+    for (const c of [...rows].reverse()) pushHeap(b, c, 2, 'cosine')
+    expect(a.map(e => e.rowIndex).sort((x, y) => x - y)).toEqual([0, 1])
+    expect(b.map(e => e.rowIndex).sort((x, y) => x - y)).toEqual([0, 1])
+  })
+
+  it('sortHeap orders tied scores by ascending rowIndex', () => {
+    const results = [
+      { rowIndex: 4, score: 0.5 },
+      { rowIndex: 1, score: 0.9 },
+      { rowIndex: 7, score: 0.9 },
+      { rowIndex: 2, score: 0.9 },
+    ]
+    expect(sortHeap(results, 'cosine').map(e => e.rowIndex)).toEqual([1, 2, 7, 4])
+    // euclidean: lower score is better, ties still ascend by rowIndex.
+    expect(sortHeap(results, 'euclidean').map(e => e.rowIndex)).toEqual([4, 1, 2, 7])
+  })
+})