Reset the cached DiskANN and IVF lookups once their row is read

sqlite-vec-diskann.c

@@ -379,7 +379,8 @@ static int diskann_node_read(vec0_vtab *p, int vec_col_idx, i64 rowid,
 
   rc = sqlite3_step(stmt);
   if (rc != SQLITE_ROW) {
-    return SQLITE_ERROR;
+    rc = SQLITE_ERROR;
+    goto done;
   }
 
   int vs = sqlite3_column_bytes(stmt, 0);
@@ -396,7 +397,8 @@ static int diskann_node_read(vec0_vtab *p, int vec_col_idx, i64 rowid,
     int expectedQs = (int)diskann_neighbor_qvecs_byte_size(
         cfg->n_neighbors, cfg->quantizer_type, col->dimensions);
     if (vs != expectedVs || is != expectedIs || qs != expectedQs) {
-      return SQLITE_CORRUPT;
+      rc = SQLITE_CORRUPT;
+      goto done;
     }
   }
 
@@ -407,7 +409,8 @@ static int diskann_node_read(vec0_vtab *p, int vec_col_idx, i64 rowid,
     sqlite3_free(v);
     sqlite3_free(ids);
     sqlite3_free(qv);
-    return SQLITE_NOMEM;
+    rc = SQLITE_NOMEM;
+    goto done;
   }
 
   const void *blobV = sqlite3_column_blob(stmt, 0);
@@ -417,7 +420,8 @@ static int diskann_node_read(vec0_vtab *p, int vec_col_idx, i64 rowid,
     sqlite3_free(v);
     sqlite3_free(ids);
     sqlite3_free(qv);
-    return SQLITE_ERROR;
+    rc = SQLITE_ERROR;
+    goto done;
   }
   memcpy(v, blobV, vs);
   memcpy(ids, blobIds, is);
@@ -426,7 +430,13 @@ static int diskann_node_read(vec0_vtab *p, int vec_col_idx, i64 rowid,
   *outValidity = v;       *outValiditySize = vs;
   *outNeighborIds = ids;  *outNeighborIdsSize = is;
   *outQvecs = qv;         *outQvecsSize = qs;
-  return SQLITE_OK;
+  rc = SQLITE_OK;
+
+done:
+  // A cached statement left on a row stays in progress until the next
+  // commit, and a connection with a statement in progress refuses VACUUM.
+  sqlite3_reset(stmt);
+  return rc;
 }
 
 /**
@@ -485,19 +495,31 @@ static int diskann_vector_read(vec0_vtab *p, int vec_col_idx, i64 rowid,
 
   rc = sqlite3_step(stmt);
   if (rc != SQLITE_ROW) {
-    return SQLITE_ERROR;
+    rc = SQLITE_ERROR;
+    goto done;
   }
 
   int sz = sqlite3_column_bytes(stmt, 0);
   const void *blob = sqlite3_column_blob(stmt, 0);
-  if (!blob || sz == 0) return SQLITE_ERROR;
+  if (!blob || sz == 0) {
+    rc = SQLITE_ERROR;
+    goto done;
+  }
   void *vec = sqlite3_malloc(sz);
-  if (!vec) return SQLITE_NOMEM;
+  if (!vec) {
+    rc = SQLITE_NOMEM;
+    goto done;
+  }
   memcpy(vec, blob, sz);
 
   *outVector = vec;
   *outVectorSize = sz;
-  return SQLITE_OK;
+  rc = SQLITE_OK;
+
+done:
+  // As in diskann_node_read(): no cached statement stays on a row.
+  sqlite3_reset(stmt);
+  return rc;
 }
 
 /**

sqlite-vec-ivf.c

@@ -609,6 +609,7 @@ static int ivf_delete(vec0_vtab *p, int col_idx, i64 rowid) {
     cell_id = sqlite3_column_int64(s, 0);
     slot = sqlite3_column_int(s, 1);
   }
+  sqlite3_reset(s);
   if (slot < 0) return SQLITE_OK;
 
   // Clear validity bit
@@ -687,6 +688,9 @@ static int ivf_get_vector_data(vec0_vtab *p, i64 rowid, int col_idx,
   if (sqlite3_step(s) != SQLITE_ROW) return SQLITE_EMPTY;
   cell_id = sqlite3_column_int64(s, 0);
   slot = sqlite3_column_int(s, 1);
+  // A cached statement left on a row stays in progress until the next commit,
+  // and a connection with a statement in progress refuses VACUUM.
+  sqlite3_reset(s);
 
   void *buf = sqlite3_malloc(vecSize);
   if (!buf) return SQLITE_NOMEM;

tests/test-cache-finalize.py

@@ -12,6 +12,7 @@
 was nested inside `#if SQLITE_VEC_EXPERIMENTAL_IVF_ENABLE`, so even
 xDisconnect/xDestroy didn't finalize DiskANN stmts in the default build.
 """
+import pytest
 from helpers import _f32
 
 
@@ -32,3 +33,28 @@ def test_vacuum_after_flat_inserts(db):
     db.execute("insert into v(rowid, a) values (1, ?)", (_f32([0.1, 0.2]),))
     db.commit()
     db.execute("VACUUM")
+
+
+@pytest.mark.parametrize(
+    "read",
+    [
+        ("select rowid from v where a match ? and k = 3", (_f32([0.5] * 8),)),
+        ("select rowid, a from v where rowid = 3", ()),
+        ("select rowid, a from v", ()),
+    ],
+    ids=["knn", "point", "fullscan"],
+)
+def test_vacuum_after_diskann_read(db, read):
+    """A read leaves no cached DiskANN statement in progress, so VACUUM runs
+    without a commit in between."""
+    db.execute(
+        "create virtual table v using vec0("
+        "a float[8] indexed by diskann(neighbor_quantizer=binary))"
+    )
+    for i in range(1, 11):
+        db.execute("insert into v(rowid, a) values (?, ?)",
+                   (i, _f32([0.1 * i] * 8)))
+    db.commit()
+    sql, params = read
+    db.execute(sql, params).fetchall()
+    db.execute("VACUUM")

tests/test-ivf.py

@@ -424,3 +424,22 @@ def test_ivf_recall_nprobe_equals_nlist(db):
     assert 50 in rowids
     assert len(rowids) == 10
     assert all(45 <= r <= 55 for r in rowids)
+
+
+@pytest.mark.parametrize(
+    "read",
+    ["SELECT rowid, v FROM t WHERE rowid = 3", "SELECT rowid, v FROM t"],
+    ids=["point", "fullscan"],
+)
+def test_ivf_vacuum_after_read(db, read):
+    """A point query or full scan leaves no cached IVF statement in progress,
+    so VACUUM runs without a commit in between."""
+    db.execute(
+        "CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2))"
+    )
+    for i in range(1, 11):
+        db.execute("INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])])
+    db.execute("INSERT INTO t(t) VALUES ('compute-centroids')")
+    db.commit()
+    db.execute(read).fetchall()
+    db.execute("VACUUM")