diff --git a/sqlite-vec-ivf.c b/sqlite-vec-ivf.c index 5bc8edbc..01a40dbf 100644 --- a/sqlite-vec-ivf.c +++ b/sqlite-vec-ivf.c @@ -670,6 +670,41 @@ static int ivf_delete(vec0_vtab *p, int col_idx, i64 rowid) { // Point query // ============================================================================ +/** + * Read a quantized column's full-precision vector from _ivf_vectors. + */ +static int ivf_full_vector_read(vec0_vtab *p, i64 rowid, int col_idx, + void **outVector, int *outVectorSize) { + int vecSize = ivf_full_vec_size(p, col_idx); + int rc = ivf_ensure_stmt(p, &p->stmtIvfVectorsRead[col_idx], + "SELECT vector FROM " VEC0_SHADOW_IVF_VECTORS_NAME " WHERE rowid = ?", + col_idx); + if (rc != SQLITE_OK) return rc; + sqlite3_stmt *s = p->stmtIvfVectorsRead[col_idx]; + sqlite3_reset(s); + sqlite3_bind_int64(s, 1, rowid); + rc = sqlite3_step(s); + if (rc == SQLITE_ROW) { + void *buf = NULL; + if (sqlite3_column_bytes(s, 0) != vecSize) { + rc = SQLITE_CORRUPT; + } else if (!(buf = sqlite3_malloc(vecSize))) { + rc = SQLITE_NOMEM; + } else { + memcpy(buf, sqlite3_column_blob(s, 0), vecSize); + *outVector = buf; + if (outVectorSize) *outVectorSize = vecSize; + rc = SQLITE_OK; + } + } else if (rc == SQLITE_DONE) { + rc = SQLITE_EMPTY; + } + // 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); + return rc; +} + static int ivf_get_vector_data(vec0_vtab *p, i64 rowid, int col_idx, void **outVector, int *outVectorSize) { int rc; @@ -677,6 +712,12 @@ static int ivf_get_vector_data(vec0_vtab *p, i64 rowid, int col_idx, i64 cell_id = 0; int slot = -1; + // A quantized column's cells hold its quantized vectors, which are not the + // vectors it was given. + if (p->vector_columns[col_idx].ivf.quantizer != VEC0_IVF_QUANTIZER_NONE) { + return ivf_full_vector_read(p, rowid, col_idx, outVector, outVectorSize); + } + rc = ivf_ensure_stmt(p, &p->stmtIvfRowidMapLookup[col_idx], "SELECT cell_id, slot FROM " VEC0_SHADOW_IVF_ROWID_MAP_NAME " WHERE rowid = ?", col_idx); @@ -1222,8 +1263,38 @@ static int ivf_scan_cells_from_stmt(vec0_vtab *p, int col_idx, return SQLITE_OK; } +/** + * Rank the rows a `rowid in (...)` constraint names by their exact distances + * to the query. A rowid the table does not have is skipped. + */ +static int ivf_rank_rowids(vec0_vtab *p, int col_idx, const void *queryVector, + struct Array *rowidsIn, + struct IvfCandidate **candidates, + int *nCandidates, int *cap) { + for (size_t i = 0; i < rowidsIn->length; i++) { + i64 rowid = ((const i64 *)rowidsIn->z)[i]; + void *vec = NULL; + int rc = ivf_get_vector_data(p, rowid, col_idx, &vec, NULL); + if (rc == SQLITE_EMPTY) continue; + if (rc != SQLITE_OK) return rc; + if (*nCandidates >= *cap) { + *cap *= 2; + struct IvfCandidate *tmp = sqlite3_realloc64(*candidates, (i64)*cap * sizeof(struct IvfCandidate)); + if (!tmp) { sqlite3_free(vec); return SQLITE_NOMEM; } + *candidates = tmp; + } + (*candidates)[*nCandidates].rowid = rowid; + (*candidates)[*nCandidates].distance = ivf_distance_float(p, col_idx, + (const float *)queryVector, (const float *)vec); + (*nCandidates)++; + sqlite3_free(vec); + } + return SQLITE_OK; +} + static int ivf_query_knn(vec0_vtab *p, int col_idx, const void *queryVector, int queryVectorSize, + struct Array *arrayRowidsIn, i64 k, struct vec0_query_knn_data *knn_data) { UNUSED_PARAMETER(queryVectorSize); int rc; @@ -1246,7 +1317,14 @@ static int ivf_query_knn(vec0_vtab *p, int col_idx, struct IvfCandidate *candidates = sqlite3_malloc64((i64)cap * sizeof(struct IvfCandidate)); if (!candidates) { sqlite3_free(queryQ); return SQLITE_NOMEM; } - if (trained) { + if (arrayRowidsIn) { + // A rowid filter names its candidates, so each is ranked by its exact + // distance; the probed cells hold only the rows near the query, which + // the filter may not name. + rc = ivf_rank_rowids(p, col_idx, queryVector, arrayRowidsIn, + &candidates, &nCandidates, &cap); + if (rc != SQLITE_OK) { sqlite3_free(queryQ); sqlite3_free(candidates); return rc; } + } else if (trained) { // Find top nprobe centroids using quantized distance int nlist = 0; rc = ivf_ensure_stmt(p, &p->stmtIvfCentroidsAll[col_idx], @@ -1333,7 +1411,8 @@ static int ivf_query_knn(vec0_vtab *p, int col_idx, qsort(candidates, nCandidates, sizeof(struct IvfCandidate), ivf_candidate_cmp); // Oversample re-ranking: re-score top (oversample*k) with full-precision vectors - if (oversample > 1 && quantizer != VEC0_IVF_QUANTIZER_NONE && nCandidates > 0) { + if (oversample > 1 && quantizer != VEC0_IVF_QUANTIZER_NONE && nCandidates > 0 && + !arrayRowidsIn) { i64 rescore_n = collect_k < nCandidates ? collect_k : nCandidates; sqlite3_stmt *stmtVec = NULL; char *zSql = sqlite3_mprintf( diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..af0c8bb9 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -3581,6 +3581,7 @@ struct vec0_vtab { sqlite3_stmt *stmtIvfRowidMapLookup[VEC0_MAX_VECTOR_COLUMNS]; // SELECT cell_id,slot FROM rowid_map WHERE rowid=? sqlite3_stmt *stmtIvfRowidMapDelete[VEC0_MAX_VECTOR_COLUMNS]; // DELETE FROM rowid_map WHERE rowid=? sqlite3_stmt *stmtIvfCentroidsAll[VEC0_MAX_VECTOR_COLUMNS]; // SELECT centroid_id,centroid FROM centroids + sqlite3_stmt *stmtIvfVectorsRead[VEC0_MAX_VECTOR_COLUMNS]; // SELECT vector FROM ivf_vectors WHERE rowid=? #endif // select latest chunk from _chunks, getting chunk_id @@ -3695,6 +3696,7 @@ void vec0_free_resources(vec0_vtab *p) { sqlite3_finalize(p->stmtIvfRowidMapLookup[i]); p->stmtIvfRowidMapLookup[i] = NULL; sqlite3_finalize(p->stmtIvfRowidMapDelete[i]); p->stmtIvfRowidMapDelete[i] = NULL; sqlite3_finalize(p->stmtIvfCentroidsAll[i]); p->stmtIvfCentroidsAll[i] = NULL; + sqlite3_finalize(p->stmtIvfVectorsRead[i]); p->stmtIvfVectorsRead[i] = NULL; } #endif #if SQLITE_VEC_ENABLE_DISKANN @@ -8067,7 +8069,8 @@ int vec0Filter_knn(vec0_cursor *pCur, vec0_vtab *p, int idxNum, // IVF dispatch: if vector column has IVF, use IVF query instead of chunk scan if (vector_column->index_type == VEC0_INDEX_TYPE_IVF) { rc = ivf_query_knn(p, vectorColumnIdx, queryVector, - (int)vector_column_byte_size(*vector_column), k, knn_data); + (int)vector_column_byte_size(*vector_column), + arrayRowidsIn, k, knn_data); if (rc != SQLITE_OK) { goto cleanup; } diff --git a/tests/test-ivf-quantization.py b/tests/test-ivf-quantization.py index 1529dadc..722f9e61 100644 --- a/tests/test-ivf-quantization.py +++ b/tests/test-ivf-quantization.py @@ -255,6 +255,32 @@ def test_ivf_quantized_delete(db): assert db.execute("SELECT count(*) FROM t_ivf_vectors00").fetchone()[0] == 9 +@pytest.mark.parametrize("quantizer", ["int8", "binary"]) +def test_ivf_quantized_vectors_read_back(db, quantizer): + """A quantized column returns the vectors it was given from a point + query, a full scan and a KNN query, before and after training.""" + db.execute( + "CREATE VIRTUAL TABLE t USING vec0(" + f"v float[16] indexed by ivf(nlist=2, quantizer={quantizer}))" + ) + vecs = { + i: _f32([((i * 7 + j * 3) % 11) / 10 - 0.5 for j in range(16)]) + for i in range(1, 11) + } + for rowid, v in vecs.items(): + db.execute("INSERT INTO t(rowid, v) VALUES (?, ?)", [rowid, v]) + + for trained in [False, True]: + if trained: + db.execute("INSERT INTO t(t) VALUES ('compute-centroids')") + assert db.execute("SELECT v FROM t WHERE rowid = 3").fetchone()[0] == vecs[3] + assert {r[0]: r[1] for r in db.execute("SELECT rowid, v FROM t")} == vecs + knn = db.execute( + "SELECT rowid, v FROM t WHERE v MATCH ? AND k = 10", [vecs[3]] + ).fetchall() + assert {r[0]: r[1] for r in knn} == vecs + + def test_ivf_binary_rejects_non_multiple_of_8_dims(db): """Binary quantizer requires dimensions divisible by 8.""" with pytest.raises(sqlite3.OperationalError): diff --git a/tests/test-ivf.py b/tests/test-ivf.py index 8b7f566b..f1a36724 100644 --- a/tests/test-ivf.py +++ b/tests/test-ivf.py @@ -2,6 +2,7 @@ import sqlite3 import struct import math +import random from helpers import _f32, exec @@ -319,6 +320,35 @@ def test_ivf_knn_k_larger_than_n(db): assert len(rows) == 5 +@pytest.mark.parametrize("quantizer", ["none", "int8", "binary"]) +def test_ivf_knn_rowid_in(db, quantizer): + """A rowid filter returns the k nearest of the rows it names, as a flat + column does, whether or not their cells are probed.""" + db.execute( + "CREATE VIRTUAL TABLE t USING vec0(" + f"v float[8] indexed by ivf(nlist=4, nprobe=1, quantizer={quantizer}))" + ) + db.execute("CREATE VIRTUAL TABLE flat USING vec0(v float[8])") + rng = random.Random(7) + for i in range(40): + v = _f32([rng.uniform(-1, 1) for _ in range(8)]) + db.execute("INSERT INTO t(rowid, v) VALUES (?, ?)", [i, v]) + db.execute("INSERT INTO flat(rowid, v) VALUES (?, ?)", [i, v]) + query = db.execute("SELECT v FROM flat WHERE rowid = 3").fetchone()[0] + + for trained in [False, True]: + if trained: + db.execute("INSERT INTO t(t) VALUES ('compute-centroids')") + for rowids in ["1, 2", "0, 20, 39", "5, 17, 30, 38, 99"]: + sql = ( + "SELECT rowid, distance FROM {} WHERE v MATCH ? AND k = 3" + f" AND rowid IN ({rowids})" + ) + got = [tuple(r) for r in db.execute(sql.format("t"), [query])] + want = [tuple(r) for r in db.execute(sql.format("flat"), [query])] + assert got == want + + # ============================================================================ # Manual centroid import (set-centroid, assign-vectors) # ============================================================================