Read the IVF trained state from _info at each use

sqlite-vec-ivf.c

@@ -282,7 +282,6 @@ static int ivf_exec(vec0_vtab *p, const char *fmt, int col_idx) {
 }
 
 static int ivf_is_trained(vec0_vtab *p, int col_idx) {
-  if (p->ivfTrainedCache[col_idx] >= 0) return p->ivfTrainedCache[col_idx];
   sqlite3_stmt *stmt = NULL;
   int trained = 0;
   char *zSql = sqlite3_mprintf(
@@ -295,7 +294,6 @@ static int ivf_is_trained(vec0_vtab *p, int col_idx) {
   }
   sqlite3_free(zSql);
   sqlite3_finalize(stmt);
-  p->ivfTrainedCache[col_idx] = trained;
   return trained;
 }
 
@@ -1005,7 +1003,6 @@ static int ivf_cmd_compute_centroids(vec0_vtab *p, int col_idx, int nlist_overri
     if (zSql) { sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); sqlite3_free(zSql);
       sqlite3_step(stmt); sqlite3_finalize(stmt); }
   }
-  p->ivfTrainedCache[col_idx] = 1;
 
   sqlite3_exec(p->db, "RELEASE ivf_train", NULL, NULL, NULL);
   sqlite3_free(vectors); sqlite3_free(rowids); sqlite3_free(centroids); sqlite3_free(assignments);
@@ -1041,7 +1038,6 @@ static int ivf_cmd_set_centroid(vec0_vtab *p, int col_idx, int centroid_id,
       p->schemaName, p->tableName, col_idx);
   if (zSql) { sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); sqlite3_free(zSql);
     sqlite3_step(stmt); sqlite3_finalize(stmt); }
-  p->ivfTrainedCache[col_idx] = 1;
   sqlite3_finalize(p->stmtIvfCentroidsAll[col_idx]); p->stmtIvfCentroidsAll[col_idx] = NULL;
   return SQLITE_OK;
 }
@@ -1157,7 +1153,6 @@ static int ivf_cmd_clear_centroids(vec0_vtab *p, int col_idx) {
       p->schemaName, p->tableName, col_idx);
   if (zSql) { sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL); sqlite3_free(zSql);
     sqlite3_step(stmt); sqlite3_finalize(stmt); }
-  p->ivfTrainedCache[col_idx] = 0;
 
   sqlite3_free(vectors); sqlite3_free(rowids);
   return SQLITE_OK;

sqlite-vec.c

@@ -3574,7 +3574,6 @@ struct vec0_vtab {
 #if SQLITE_VEC_EXPERIMENTAL_IVF_ENABLE
   // IVF cached state per vector column
   char *shadowIvfCellsNames[VEC0_MAX_VECTOR_COLUMNS];   // table name for blob_open
-  int ivfTrainedCache[VEC0_MAX_VECTOR_COLUMNS];          // -1=unknown, 0=no, 1=yes
   sqlite3_stmt *stmtIvfCellMeta[VEC0_MAX_VECTOR_COLUMNS];      // SELECT n_vectors, length(validity)*8 FROM cells WHERE cell_id=?
   sqlite3_stmt *stmtIvfCellUpdateN[VEC0_MAX_VECTOR_COLUMNS];   // UPDATE cells SET n_vectors=n_vectors+? WHERE cell_id=?
   sqlite3_stmt *stmtIvfRowidMapInsert[VEC0_MAX_VECTOR_COLUMNS]; // INSERT INTO rowid_map(rowid,cell_id,slot) VALUES(?,?,?)
@@ -5409,7 +5408,6 @@ static int vec0_init(sqlite3 *db, void *pAux, int argc, const char *const *argv,
     pNew->shadowIvfCellsNames[i] =
         sqlite3_mprintf("%s_ivf_cells%02d", tableName, i);
     if (!pNew->shadowIvfCellsNames[i]) goto error;
-    pNew->ivfTrainedCache[i] = -1; // unknown
   }
 #endif
   for (int i = 0; i < pNew->numMetadataColumns; i++) {

tests/test-ivf.py

@@ -376,6 +376,70 @@ def test_clear_centroids(db):
     assert str(trained) == "0"
 
 
+# ============================================================================
+# Trained state across connections and rollbacks
+# ============================================================================
+
+
+def ivf_connect(path):
+    db = sqlite3.connect(path, isolation_level=None)
+    db.enable_load_extension(True)
+    db.load_extension("dist/vec0")
+    db.enable_load_extension(False)
+    return db
+
+
+def ivf_create_rows(db):
+    db.execute(
+        "CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2, nprobe=2))"
+    )
+    for i in range(20):
+        db.execute("INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])])
+
+
+IVF_KNN = "SELECT rowid FROM t WHERE v MATCH ? AND k = 3"
+
+
+def test_ivf_knn_after_another_connection_trains(tmp_path):
+    """A connection that queried before another ran compute-centroids still
+    finds the rows, which have moved into the trained cells."""
+    path = str(tmp_path / "ivf.db")
+    a = ivf_connect(path)
+    ivf_create_rows(a)
+    assert len(a.execute(IVF_KNN, [_f32([10, 0, 0, 0])]).fetchall()) == 3
+
+    ivf_connect(path).execute("INSERT INTO t(t) VALUES ('compute-centroids')")
+    assert len(a.execute(IVF_KNN, [_f32([10, 0, 0, 0])]).fetchall()) == 3
+
+
+def test_ivf_insert_after_another_connection_clears(tmp_path):
+    """A connection that trained the index can insert after another
+    connection ran clear-centroids."""
+    path = str(tmp_path / "ivf.db")
+    a = ivf_connect(path)
+    ivf_create_rows(a)
+    a.execute("INSERT INTO t(t) VALUES ('compute-centroids')")
+
+    ivf_connect(path).execute("INSERT INTO t(t) VALUES ('clear-centroids')")
+    a.execute("INSERT INTO t(rowid, v) VALUES (20, ?)", [_f32([20, 0, 0, 0])])
+    rows = a.execute(IVF_KNN, [_f32([20, 0, 0, 0])]).fetchall()
+    assert rows[0][0] == 20
+
+
+def test_ivf_trained_state_after_rollback(db):
+    """A rolled-back compute-centroids leaves the index untrained, so inserts
+    and queries work as before it."""
+    db.isolation_level = None
+    ivf_create_rows(db)
+    db.execute("BEGIN")
+    db.execute("INSERT INTO t(t) VALUES ('compute-centroids')")
+    db.execute("ROLLBACK")
+
+    db.execute("INSERT INTO t(rowid, v) VALUES (20, ?)", [_f32([20, 0, 0, 0])])
+    rows = db.execute(IVF_KNN, [_f32([20, 0, 0, 0])]).fetchall()
+    assert rows[0][0] == 20
+
+
 # ============================================================================
 # Delete after training
 # ============================================================================