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 # ============================================================================