From ac363ae4e378c8569210a65d0e3657f8bf2af6bb Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Fri, 23 Oct 2026 08:57:01 -0500 Subject: [PATCH] Read the IVF trained state from _info at each use ivf_is_trained() kept the trained state per connection once it had read it, and compute-centroids, set-centroid and clear-centroids set that cache as they ran. Another connection's command, or a rolled-back one, left it wrong: a connection that had queried before another trained the index scanned only the emptied unassigned cells and found no rows, and one that had trained it, then saw another connection clear it or rolled its own training back, failed every insert and query with "SQL logic error". The state is now read from _info each time, so it is always that of the transaction reading it. --- sqlite-vec-ivf.c | 5 ---- sqlite-vec.c | 2 -- tests/test-ivf.py | 64 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 64 insertions(+), 7 deletions(-) diff --git a/sqlite-vec-ivf.c b/sqlite-vec-ivf.c index 5bc8edbc..76a4dce7 100644 --- a/sqlite-vec-ivf.c +++ b/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; diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..1616bffb 100644 --- a/sqlite-vec.c +++ b/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++) { diff --git a/tests/test-ivf.py b/tests/test-ivf.py index 8b7f566b..9d50f8c2 100644 --- a/tests/test-ivf.py +++ b/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 # ============================================================================