Give every DiskANN result its exact distance

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)
   }

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("""