From 9e74f15eb0ba2b2b62a9b09a3d93399e1515eef4 Mon Sep 17 00:00:00 2001
From: Aaron Paterson <apaterson@pm.me>
Date: Fri, 23 Oct 2026 08:17:33 -0500
Subject: [PATCH] Give every DiskANN result its exact distance

A DiskANN search ranks the neighbors it reaches by quantized distance,
then gives each candidate it visits its exact distance through
diskann_candidate_list_insert(), which only ever lowers a distance. A
candidate whose exact distance was larger than its quantized one kept
the quantized one, so the results came back with distances that were
not theirs, in an order that was not the exact one.

diskann_candidate_list_confirm() replaces a visited candidate's distance
with its exact one, larger or smaller, moves it to its place, and
confirms it; a candidate the list dropped after it was visited returns
visited, when it still ranks inside the list.
---
 sqlite-vec-diskann.c  | 42 +++++++++++++++++++++++++++++++++---------
 tests/test-diskann.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/sqlite-vec-diskann.c b/sqlite-vec-diskann.c
index 5bd298bc..f2592125 100644
--- a/sqlite-vec-diskann.c
+++ b/sqlite-vec-diskann.c
@@ -613,6 +613,38 @@ static int diskann_candidate_list_insert(
   return 1;
 }
 
+/**
+ * Give a visited candidate its exact distance and confirm it (its vector
+ * exists). The exact distance replaces the quantized one whether it is larger
+ * or smaller, and the candidate moves to its place in the order. A candidate
+ * the list dropped after it was visited comes back, still visited, when its
+ * exact distance ranks inside the list.
+ */
+static void diskann_candidate_list_confirm(
+    struct DiskannCandidateList *list, i64 rowid, f32 distance) {
+  int i = 0;
+  while (i < list->count && list->items[i].rowid != rowid) i++;
+  if (i == list->count) {
+    if (!diskann_candidate_list_insert(list, rowid, distance)) return;
+    i = 0;
+    while (list->items[i].rowid != rowid) i++;
+    list->items[i].visited = 1;
+  }
+
+  struct Vec0DiskannCandidate item = list->items[i];
+  item.distance = distance;
+  item.confirmed = 1;
+  while (i > 0 && list->items[i - 1].distance > distance) {
+    list->items[i] = list->items[i - 1];
+    i--;
+  }
+  while (i + 1 < list->count && list->items[i + 1].distance < distance) {
+    list->items[i] = list->items[i + 1];
+    i++;
+  }
+  list->items[i] = item;
+}
+
 /**
  * Find the closest unvisited candidate. Returns its index, or -1 if none.
  */
@@ -815,15 +847,7 @@ static int diskann_search(
                                              dimensions, elementType,
                                              col->distance_metric);
       sqlite3_free(fullVec);
-      // Update distance in candidate list and re-sort
-      diskann_candidate_list_insert(&candidates, currentRowid, exactDist);
-      // Mark as confirmed (vector exists, distance is exact)
-      for (int ci = 0; ci < candidates.count; ci++) {
-        if (candidates.items[ci].rowid == currentRowid) {
-          candidates.items[ci].confirmed = 1;
-          break;
-        }
-      }
+      diskann_candidate_list_confirm(&candidates, currentRowid, exactDist);
     }
     // If vector read failed, candidate stays unconfirmed (stale edge to deleted node)
   }
diff --git a/tests/test-diskann.py b/tests/test-diskann.py
index 4e651604..4a239f67 100644
--- a/tests/test-diskann.py
+++ b/tests/test-diskann.py
@@ -1091,6 +1091,34 @@ def test_diskann_knn_wrong_query_dimensions(db):
     assert "error" in result
 
 
+def test_diskann_result_distances_are_exact(db):
+    """Every result carries its exact distance, in order, including the rows
+    whose exact distance is larger than their quantized one."""
+    db.execute("""
+        CREATE VIRTUAL TABLE t USING vec0(
+            emb float[32] distance_metric=cosine INDEXED BY diskann(neighbor_quantizer=int8)
+        )
+    """)
+    import random
+    random.seed(3)
+    for i in range(1, 301):
+        vec = [random.gauss(0, 1) for _ in range(32)]
+        db.execute("INSERT INTO t(rowid, emb) VALUES (?, ?)", [i, _f32(vec)])
+
+    for _ in range(30):
+        query = _f32([random.gauss(0, 1) for _ in range(32)])
+        rows = db.execute(
+            "SELECT rowid, distance, vec_distance_cosine(emb, ?) "
+            "FROM t WHERE emb MATCH ? AND k = 10",
+            [query, query],
+        ).fetchall()
+        assert len(rows) == 10
+        for rowid, distance, exact in rows:
+            assert distance == pytest.approx(exact, abs=1e-6)
+        distances = [r[1] for r in rows]
+        assert distances == sorted(distances)
+
+
 def test_diskann_graph_connectivity_after_many_deletes(db):
     """After many deletes, the graph should still be connected enough for search."""
     db.execute("""