diff --git a/sqlite-vec-diskann.c b/sqlite-vec-diskann.c index 5bd298bc..7c2af0b1 100644 --- a/sqlite-vec-diskann.c +++ b/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; } /** diff --git a/sqlite-vec-ivf.c b/sqlite-vec-ivf.c index 5bc8edbc..7a65e2e0 100644 --- a/sqlite-vec-ivf.c +++ b/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; diff --git a/tests/test-cache-finalize.py b/tests/test-cache-finalize.py index c81f45a8..bd01fcaf 100644 --- a/tests/test-cache-finalize.py +++ b/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") diff --git a/tests/test-ivf.py b/tests/test-ivf.py index 8b7f566b..eed5fe43 100644 --- a/tests/test-ivf.py +++ b/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")