From 74d426a0006572a68a5f85ce6d9f450742fd92f9 Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Fri, 23 Oct 2026 08:44:55 -0500 Subject: [PATCH] Read a quantized IVF column's vectors from _ivf_vectors A quantized IVF column keeps each vector's int8 or binary form in its cell and the vector as it was inserted in _ivf_vectors, but ivf_get_vector_data() returned the cell's bytes for every column. A point query then read the column's float32 elements past the end of the quantized vector, and a full scan or a KNN query returned the quantized bytes as the column's value. A quantized column's vectors are now read from _ivf_vectors. --- sqlite-vec-ivf.c | 41 ++++++++++++++++++++++++++++++++++ sqlite-vec.c | 2 ++ tests/test-ivf-quantization.py | 26 +++++++++++++++++++++ 3 files changed, 69 insertions(+) diff --git a/sqlite-vec-ivf.c b/sqlite-vec-ivf.c index 5bc8edbc..221472f8 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); diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..1e681f94 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 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):