Replace a deleted DiskANN medoid with a node of the graph

sqlite-vec-diskann.c

@@ -314,7 +314,9 @@ static int diskann_medoid_set(vec0_vtab *p, int vec_col_idx,
 
 /**
  * Called when deleting a vector. If the deleted vector was the medoid,
- * pick a new one (the first available vector, or set to empty if none remain).
+ * pick a new one: the first remaining node of the graph, or set to empty if
+ * none remain. A vector still in the buffer has no node, so it cannot be the
+ * entry point of a search.
  */
 static int diskann_medoid_handle_delete(vec0_vtab *p, int vec_col_idx,
                                           i64 deletedRowid) {
@@ -326,7 +328,7 @@ static int diskann_medoid_handle_delete(vec0_vtab *p, int vec_col_idx,
   if (!isEmpty && currentMedoid == deletedRowid) {
     sqlite3_stmt *stmt = NULL;
     char *zSql = sqlite3_mprintf(
-        "SELECT rowid FROM " VEC0_SHADOW_VECTORS_N_NAME " WHERE rowid != ?1 LIMIT 1",
+        "SELECT rowid FROM " VEC0_SHADOW_DISKANN_NODES_N_NAME " WHERE rowid != ?1 LIMIT 1",
         p->schemaName, p->tableName, vec_col_idx);
     if (!zSql) return SQLITE_NOMEM;
 
@@ -340,10 +342,10 @@ static int diskann_medoid_handle_delete(vec0_vtab *p, int vec_col_idx,
       i64 newMedoid = sqlite3_column_int64(stmt, 0);
       sqlite3_finalize(stmt);
       return diskann_medoid_set(p, vec_col_idx, newMedoid, 0);
-    } else {
-      sqlite3_finalize(stmt);
-      return diskann_medoid_set(p, vec_col_idx, -1, 1);
     }
+    sqlite3_finalize(stmt);
+    if (rc != SQLITE_DONE) return rc;
+    return diskann_medoid_set(p, vec_col_idx, -1, 1);
   }
   return SQLITE_OK;
 }

tests/test-diskann.py

@@ -889,6 +889,41 @@ def test_diskann_delete_preserves_graph_connectivity(db):
         f"Only {len(reachable)}/{len(all_rowids)} nodes reachable"
 
 
+def test_diskann_delete_medoid_with_rows_in_buffer(db):
+    """Deleting the medoid while rows wait in the buffer makes another node
+    of the graph the medoid, so every node is still found by its own vector."""
+    db.execute("""
+        CREATE VIRTUAL TABLE t USING vec0(
+            emb float[8] INDEXED BY diskann(
+                neighbor_quantizer=int8, n_neighbors=8, buffer_threshold=10
+            )
+        )
+    """)
+    import random
+    random.seed(11)
+    vectors = {}
+    for rowid in [*range(100, 110), 1, 2, 3]:
+        vectors[rowid] = _f32([random.uniform(-1, 1) for _ in range(8)])
+        db.execute("INSERT INTO t(rowid, emb) VALUES (?, ?)", [rowid, vectors[rowid]])
+    assert db.execute("SELECT count(*) FROM t_diskann_buffer00").fetchone()[0] == 3
+
+    medoid = db.execute(
+        "SELECT value FROM t_info WHERE key = 'diskann_medoid_00'"
+    ).fetchone()[0]
+    db.execute("DELETE FROM t WHERE rowid = ?", [medoid])
+
+    new_medoid = db.execute(
+        "SELECT value FROM t_info WHERE key = 'diskann_medoid_00'"
+    ).fetchone()[0]
+    nodes = [r[0] for r in db.execute("SELECT rowid FROM t_diskann_nodes00")]
+    assert new_medoid in nodes
+    for rowid in nodes:
+        rows = db.execute(
+            "SELECT rowid FROM t WHERE emb MATCH ? AND k = 1", [vectors[rowid]]
+        ).fetchall()
+        assert [r[0] for r in rows] == [rowid]
+
+
 # ======================================================================
 # Update scenarios
 # ======================================================================