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