apunkt/sqlite-vec
1
0
Fork 0
mirror of https://github.com/asg017/sqlite-vec.git synced 2026-04-26 17:26:25 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Add IVF index for vec0 virtual table

Browse source 
Add inverted file (IVF) index type: partitions vectors into clusters via
k-means, quantizes to int8, and scans only the nearest nprobe partitions at
query time. Includes shadow table management, insert/delete, KNN integration,
compile flag (SQLITE_VEC_ENABLE_IVF), fuzz targets, and tests. Removes
superseded ivf-benchmarks/ directory.
This commit is contained in:
Alex Garcia 2026-03-29 19:46:23 -07:00
parent 43982c144b
commit 3358e127f6
22 changed files with 5237 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

29
tests/fuzz/Makefile
View file

@ -93,13 +93,40 @@ $(TARGET_DIR)/rescore_quantize_edge: rescore-quantize-edge.c $(FUZZ_SRCS) | $(TA
$(TARGET_DIR)/rescore_interleave: rescore-interleave.c $(FUZZ_SRCS) | $(TARGET_DIR)
$(FUZZ_CC) $(FUZZ_CFLAGS) -DSQLITE_VEC_ENABLE_RESCORE $(FUZZ_SRCS) $< -o $@
$(TARGET_DIR)/ivf_create: ivf-create.c $(FUZZ_SRCS) | $(TARGET_DIR)
$(FUZZ_CC) $(FUZZ_CFLAGS) $(FUZZ_SRCS) $< -o $@
$(TARGET_DIR)/ivf_operations: ivf-operations.c $(FUZZ_SRCS) | $(TARGET_DIR)
$(FUZZ_CC) $(FUZZ_CFLAGS) $(FUZZ_SRCS) $< -o $@
$(TARGET_DIR)/ivf_quantize: ivf-quantize.c $(FUZZ_SRCS) | $(TARGET_DIR)
$(FUZZ_CC) $(FUZZ_CFLAGS) $(FUZZ_SRCS) $< -o $@
$(TARGET_DIR)/ivf_kmeans: ivf-kmeans.c $(FUZZ_SRCS) | $(TARGET_DIR)
$(FUZZ_CC) $(FUZZ_CFLAGS) $(FUZZ_SRCS) $< -o $@
$(TARGET_DIR)/ivf_shadow_corrupt: ivf-shadow-corrupt.c $(FUZZ_SRCS) | $(TARGET_DIR)
$(FUZZ_CC) $(FUZZ_CFLAGS) $(FUZZ_SRCS) $< -o $@
$(TARGET_DIR)/ivf_knn_deep: ivf-knn-deep.c $(FUZZ_SRCS) | $(TARGET_DIR)
$(FUZZ_CC) $(FUZZ_CFLAGS) $(FUZZ_SRCS) $< -o $@
$(TARGET_DIR)/ivf_cell_overflow: ivf-cell-overflow.c $(FUZZ_SRCS) | $(TARGET_DIR)
$(FUZZ_CC) $(FUZZ_CFLAGS) $(FUZZ_SRCS) $< -o $@
$(TARGET_DIR)/ivf_rescore: ivf-rescore.c $(FUZZ_SRCS) | $(TARGET_DIR)
$(FUZZ_CC) $(FUZZ_CFLAGS) $(FUZZ_SRCS) $< -o $@
FUZZ_TARGETS = vec0_create exec json numpy \
shadow_corrupt vec0_operations scalar_functions \
vec0_create_full metadata_columns vec_each vec_mismatch \
vec0_delete_completeness \
rescore_operations rescore_create rescore_quantize \
rescore_shadow_corrupt rescore_knn_deep \
rescore_quantize_edge rescore_interleave
rescore_quantize_edge rescore_interleave \
ivf_create ivf_operations \
ivf_quantize ivf_kmeans ivf_shadow_corrupt \
ivf_knn_deep ivf_cell_overflow ivf_rescore
all: $(addprefix $(TARGET_DIR)/,$(FUZZ_TARGETS))

192
tests/fuzz/ivf-cell-overflow.c Normal file
View file

@ -0,0 +1,192 @@
/**
* Fuzz target: IVF cell overflow and boundary conditions.
*
* Pushes cells past VEC0_IVF_CELL_MAX_VECTORS (64) to trigger cell
* splitting, then exercises blob I/O at slot boundaries.
*
* Targets:
* - Cell splitting when n_vectors reaches cap (64)
* - Blob offset arithmetic: slot * vecSize, slot / 8, slot % 8
* - Validity bitmap at byte boundaries (slot 7->8, 15->16, etc.)
* - Insert into full cell -> create new cell path
* - Delete from various slot positions (first, last, middle)
* - Multiple cells per centroid
* - assign-vectors command with multi-cell centroids
*/
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "sqlite-vec.h"
#include "sqlite3.h"
#include <assert.h>
int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
if (size < 8) return 0;
int rc;
sqlite3 *db;
rc = sqlite3_open(":memory:", &db);
assert(rc == SQLITE_OK);
rc = sqlite3_vec_init(db, NULL, NULL);
assert(rc == SQLITE_OK);
// Use small dimensions for speed but enough vectors to overflow cells
int dim = (data[0] % 8) + 2; // 2..9
int nlist = (data[1] % 4) + 1; // 1..4
// We need >64 vectors to overflow a cell
int num_vecs = (data[2] % 64) + 65; // 65..128
int delete_pattern = data[3]; // Controls which vectors to delete
const uint8_t *payload = data + 4;
size_t payload_size = size - 4;
char sql[256];
snprintf(sql, sizeof(sql),
"CREATE VIRTUAL TABLE v USING vec0("
"emb float[%d] indexed by ivf(nlist=%d, nprobe=%d))",
dim, nlist, nlist);
rc = sqlite3_exec(db, sql, NULL, NULL, NULL);
if (rc != SQLITE_OK) { sqlite3_close(db); return 0; }
// Insert enough vectors to overflow at least one cell
sqlite3_stmt *stmtInsert = NULL;
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &stmtInsert, NULL);
if (!stmtInsert) { sqlite3_close(db); return 0; }
size_t offset = 0;
for (int i = 0; i < num_vecs; i++) {
float *vec = sqlite3_malloc(dim * sizeof(float));
if (!vec) break;
for (int d = 0; d < dim; d++) {
if (offset < payload_size) {
vec[d] = ((float)(int8_t)payload[offset++]) / 50.0f;
} else {
// Cluster vectors near specific centroids to ensure some cells overflow
int cluster = i % nlist;
vec[d] = (float)cluster + (float)(i % 10) * 0.01f + d * 0.001f;
}
}
sqlite3_reset(stmtInsert);
sqlite3_bind_int64(stmtInsert, 1, (int64_t)(i + 1));
sqlite3_bind_blob(stmtInsert, 2, vec, dim * sizeof(float), SQLITE_TRANSIENT);
sqlite3_step(stmtInsert);
sqlite3_free(vec);
}
sqlite3_finalize(stmtInsert);
// Train to assign vectors to centroids (triggers cell building)
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('compute-centroids')",
NULL, NULL, NULL);
// Delete vectors at boundary positions based on fuzz data
// This tests validity bitmap manipulation at different slot positions
for (int i = 0; i < num_vecs; i++) {
int byte_idx = i / 8;
if (byte_idx < (int)payload_size && (payload[byte_idx] & (1 << (i % 8)))) {
// Use delete_pattern to thin deletions
if ((delete_pattern + i) % 3 == 0) {
char delsql[64];
snprintf(delsql, sizeof(delsql), "DELETE FROM v WHERE rowid = %d", i + 1);
sqlite3_exec(db, delsql, NULL, NULL, NULL);
}
}
}
// Insert more vectors after deletions (into cells with holes)
{
sqlite3_stmt *si = NULL;
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &si, NULL);
if (si) {
for (int i = 0; i < 10; i++) {
float *vec = sqlite3_malloc(dim * sizeof(float));
if (!vec) break;
for (int d = 0; d < dim; d++)
vec[d] = (float)(i + 200) * 0.01f;
sqlite3_reset(si);
sqlite3_bind_int64(si, 1, (int64_t)(num_vecs + i + 1));
sqlite3_bind_blob(si, 2, vec, dim * sizeof(float), SQLITE_TRANSIENT);
sqlite3_step(si);
sqlite3_free(vec);
}
sqlite3_finalize(si);
}
}
// KNN query that must scan multiple cells per centroid
{
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (qvec) {
for (int d = 0; d < dim; d++) qvec[d] = 0.0f;
sqlite3_stmt *sk = NULL;
snprintf(sql, sizeof(sql),
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT 20");
sqlite3_prepare_v2(db, sql, -1, &sk, NULL);
if (sk) {
sqlite3_bind_blob(sk, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(sk) == SQLITE_ROW) {}
sqlite3_finalize(sk);
}
sqlite3_free(qvec);
}
}
// Test assign-vectors with multi-cell state
// First clear centroids
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('clear-centroids')",
NULL, NULL, NULL);
// Set centroids manually, then assign
for (int c = 0; c < nlist; c++) {
float *cvec = sqlite3_malloc(dim * sizeof(float));
if (!cvec) break;
for (int d = 0; d < dim; d++) cvec[d] = (float)c + d * 0.1f;
char cmd[128];
snprintf(cmd, sizeof(cmd),
"INSERT INTO v(rowid, emb) VALUES ('set-centroid:%d', ?)", c);
sqlite3_stmt *sc = NULL;
sqlite3_prepare_v2(db, cmd, -1, &sc, NULL);
if (sc) {
sqlite3_bind_blob(sc, 1, cvec, dim * sizeof(float), SQLITE_TRANSIENT);
sqlite3_step(sc);
sqlite3_finalize(sc);
}
sqlite3_free(cvec);
}
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('assign-vectors')",
NULL, NULL, NULL);
// Final query after assign-vectors
{
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (qvec) {
for (int d = 0; d < dim; d++) qvec[d] = 1.0f;
sqlite3_stmt *sk = NULL;
sqlite3_prepare_v2(db,
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT 5",
-1, &sk, NULL);
if (sk) {
sqlite3_bind_blob(sk, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(sk) == SQLITE_ROW) {}
sqlite3_finalize(sk);
}
sqlite3_free(qvec);
}
}
// Full scan
sqlite3_exec(db, "SELECT * FROM v", NULL, NULL, NULL);
sqlite3_close(db);
return 0;
}

36
tests/fuzz/ivf-create.c Normal file
View file

@ -0,0 +1,36 @@
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "sqlite-vec.h"
#include "sqlite3.h"
#include <assert.h>
int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
int rc;
sqlite3 *db;
sqlite3_stmt *stmt;
rc = sqlite3_open(":memory:", &db);
assert(rc == SQLITE_OK);
rc = sqlite3_vec_init(db, NULL, NULL);
assert(rc == SQLITE_OK);
sqlite3_str *s = sqlite3_str_new(NULL);
assert(s);
sqlite3_str_appendall(s, "CREATE VIRTUAL TABLE v USING vec0(emb float[4] indexed by ivf(");
sqlite3_str_appendf(s, "%.*s", (int)size, data);
sqlite3_str_appendall(s, "))");
const char *zSql = sqlite3_str_finish(s);
assert(zSql);
rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, NULL);
sqlite3_free((void *)zSql);
if (rc == SQLITE_OK) {
sqlite3_step(stmt);
}
sqlite3_finalize(stmt);
sqlite3_close(db);
return 0;
}

16
tests/fuzz/ivf-create.dict Normal file
View file

@ -0,0 +1,16 @@
"nlist"
"nprobe"
"quantizer"
"oversample"
"binary"
"int8"
"none"
"="
","
"("
")"
"0"
"1"
"128"
"65536"
"65537"

180
tests/fuzz/ivf-kmeans.c Normal file
View file

@ -0,0 +1,180 @@
/**
* Fuzz target: IVF k-means clustering.
*
* Builds a table, inserts fuzz-controlled vectors, then runs
* compute-centroids with fuzz-controlled parameters (nlist, max_iter, seed).
* Targets:
* - kmeans with N < k (clamping), N == 1, k == 1
* - kmeans with duplicate/identical vectors (all distances zero)
* - kmeans with NaN/Inf vectors
* - Empty cluster reassignment path (farthest-point heuristic)
* - Large nlist relative to N
* - The compute-centroids:{json} command parsing
* - clear-centroids followed by compute-centroids (round-trip)
*/
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "sqlite-vec.h"
#include "sqlite3.h"
#include <assert.h>
int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
if (size < 10) return 0;
int rc;
sqlite3 *db;
rc = sqlite3_open(":memory:", &db);
assert(rc == SQLITE_OK);
rc = sqlite3_vec_init(db, NULL, NULL);
assert(rc == SQLITE_OK);
// Parse fuzz header
// Byte 0-1: dimension (1..128)
// Byte 2: nlist for CREATE (1..64)
// Byte 3: nlist override for compute-centroids (0 = use default)
// Byte 4: max_iter (1..50)
// Byte 5-8: seed
// Byte 9: num_vectors (1..64)
// Remaining: vector float data
int dim = (data[0] | (data[1] << 8)) % 128 + 1;
int nlist_create = (data[2] % 64) + 1;
int nlist_override = data[3] % 65; // 0 means use table default
int max_iter = (data[4] % 50) + 1;
uint32_t seed = (uint32_t)data[5] | ((uint32_t)data[6] << 8) |
((uint32_t)data[7] << 16) | ((uint32_t)data[8] << 24);
int num_vecs = (data[9] % 64) + 1;
const uint8_t *payload = data + 10;
size_t payload_size = size - 10;
char sql[256];
snprintf(sql, sizeof(sql),
"CREATE VIRTUAL TABLE v USING vec0("
"emb float[%d] indexed by ivf(nlist=%d, nprobe=%d))",
dim, nlist_create, nlist_create);
rc = sqlite3_exec(db, sql, NULL, NULL, NULL);
if (rc != SQLITE_OK) { sqlite3_close(db); return 0; }
// Insert vectors
sqlite3_stmt *stmtInsert = NULL;
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &stmtInsert, NULL);
if (!stmtInsert) { sqlite3_close(db); return 0; }
size_t offset = 0;
for (int i = 0; i < num_vecs; i++) {
float *vec = sqlite3_malloc(dim * sizeof(float));
if (!vec) break;
for (int d = 0; d < dim; d++) {
if (offset + 4 <= payload_size) {
memcpy(&vec[d], payload + offset, sizeof(float));
offset += 4;
} else if (offset < payload_size) {
// Scale to interesting range including values > 1, < -1
vec[d] = ((float)(int8_t)payload[offset++]) / 5.0f;
} else {
// Reuse earlier bytes to fill remaining dimensions
vec[d] = (float)(i * dim + d) * 0.01f;
}
}
sqlite3_reset(stmtInsert);
sqlite3_bind_int64(stmtInsert, 1, (int64_t)(i + 1));
sqlite3_bind_blob(stmtInsert, 2, vec, dim * sizeof(float), SQLITE_TRANSIENT);
sqlite3_step(stmtInsert);
sqlite3_free(vec);
}
sqlite3_finalize(stmtInsert);
// Exercise compute-centroids with JSON options
{
char cmd[256];
snprintf(cmd, sizeof(cmd),
"INSERT INTO v(rowid) VALUES "
"('compute-centroids:{\"nlist\":%d,\"max_iterations\":%d,\"seed\":%u}')",
nlist_override, max_iter, seed);
sqlite3_exec(db, cmd, NULL, NULL, NULL);
}
// KNN query after training
{
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (qvec) {
for (int d = 0; d < dim; d++) {
qvec[d] = (d < 3) ? 1.0f : 0.0f;
}
sqlite3_stmt *stmtKnn = NULL;
sqlite3_prepare_v2(db,
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT 5",
-1, &stmtKnn, NULL);
if (stmtKnn) {
sqlite3_bind_blob(stmtKnn, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(stmtKnn) == SQLITE_ROW) {}
sqlite3_finalize(stmtKnn);
}
sqlite3_free(qvec);
}
}
// Clear centroids and re-compute to test round-trip
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('clear-centroids')",
NULL, NULL, NULL);
// Insert a few more vectors in untrained state
{
sqlite3_stmt *si = NULL;
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &si, NULL);
if (si) {
for (int i = 0; i < 3; i++) {
float *vec = sqlite3_malloc(dim * sizeof(float));
if (!vec) break;
for (int d = 0; d < dim; d++) vec[d] = (float)(i + 100) * 0.1f;
sqlite3_reset(si);
sqlite3_bind_int64(si, 1, (int64_t)(num_vecs + i + 1));
sqlite3_bind_blob(si, 2, vec, dim * sizeof(float), SQLITE_TRANSIENT);
sqlite3_step(si);
sqlite3_free(vec);
}
sqlite3_finalize(si);
}
}
// Re-train
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('compute-centroids')",
NULL, NULL, NULL);
// Delete some rows after training, then query
sqlite3_exec(db, "DELETE FROM v WHERE rowid = 1", NULL, NULL, NULL);
sqlite3_exec(db, "DELETE FROM v WHERE rowid = 2", NULL, NULL, NULL);
// Query after deletes
{
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (qvec) {
for (int d = 0; d < dim; d++) qvec[d] = 0.5f;
sqlite3_stmt *stmtKnn = NULL;
sqlite3_prepare_v2(db,
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT 10",
-1, &stmtKnn, NULL);
if (stmtKnn) {
sqlite3_bind_blob(stmtKnn, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(stmtKnn) == SQLITE_ROW) {}
sqlite3_finalize(stmtKnn);
}
sqlite3_free(qvec);
}
}
sqlite3_close(db);
return 0;
}

199
tests/fuzz/ivf-knn-deep.c Normal file
View file

@ -0,0 +1,199 @@
/**
* Fuzz target: IVF KNN search deep paths.
*
* Exercises the full KNN pipeline with fuzz-controlled:
* - nprobe values (including > nlist, =1, =nlist)
* - Query vectors (including adversarial floats)
* - Mix of trained/untrained state
* - Oversample + rescore path (quantizer=int8 with oversample>1)
* - Multiple interleaved KNN queries
* - Candidate array realloc path (many vectors in probed cells)
*
* Targets:
* - ivf_scan_cells_from_stmt: candidate realloc, distance computation
* - ivf_query_knn: centroid sorting, nprobe selection
* - Oversample rescore: re-ranking with full-precision vectors
* - qsort with NaN distances
*/
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "sqlite-vec.h"
#include "sqlite3.h"
#include <assert.h>
static uint16_t read_u16(const uint8_t *p) {
return (uint16_t)(p[0] | (p[1] << 8));
}
int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
if (size < 16) return 0;
int rc;
sqlite3 *db;
rc = sqlite3_open(":memory:", &db);
assert(rc == SQLITE_OK);
rc = sqlite3_vec_init(db, NULL, NULL);
assert(rc == SQLITE_OK);
// Header
int dim = (data[0] % 32) + 2; // 2..33
int nlist = (data[1] % 16) + 1; // 1..16
int nprobe_initial = (data[2] % 20) + 1; // 1..20 (can be > nlist)
int quantizer_type = data[3] % 3; // 0=none, 1=int8, 2=binary
int oversample = (data[4] % 4) + 1; // 1..4
int num_vecs = (data[5] % 80) + 4; // 4..83
int num_queries = (data[6] % 8) + 1; // 1..8
int k_limit = (data[7] % 20) + 1; // 1..20
const uint8_t *payload = data + 8;
size_t payload_size = size - 8;
// For binary quantizer, dimension must be multiple of 8
if (quantizer_type == 2) {
dim = ((dim + 7) / 8) * 8;
if (dim == 0) dim = 8;
}
const char *qname;
switch (quantizer_type) {
case 1: qname = "int8"; break;
case 2: qname = "binary"; break;
default: qname = "none"; break;
}
// Oversample only valid with quantization
if (quantizer_type == 0) oversample = 1;
// Cap nprobe to nlist for CREATE (parser rejects nprobe > nlist)
int nprobe_create = nprobe_initial <= nlist ? nprobe_initial : nlist;
char sql[512];
snprintf(sql, sizeof(sql),
"CREATE VIRTUAL TABLE v USING vec0("
"emb float[%d] indexed by ivf(nlist=%d, nprobe=%d, quantizer=%s%s))",
dim, nlist, nprobe_create, qname,
oversample > 1 ? ", oversample=2" : "");
// If that fails (e.g. oversample with none), try without oversample
rc = sqlite3_exec(db, sql, NULL, NULL, NULL);
if (rc != SQLITE_OK) {
snprintf(sql, sizeof(sql),
"CREATE VIRTUAL TABLE v USING vec0("
"emb float[%d] indexed by ivf(nlist=%d, nprobe=%d, quantizer=%s))",
dim, nlist, nprobe_create, qname);
rc = sqlite3_exec(db, sql, NULL, NULL, NULL);
if (rc != SQLITE_OK) { sqlite3_close(db); return 0; }
}
// Insert vectors
sqlite3_stmt *stmtInsert = NULL;
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &stmtInsert, NULL);
if (!stmtInsert) { sqlite3_close(db); return 0; }
size_t offset = 0;
for (int i = 0; i < num_vecs; i++) {
float *vec = sqlite3_malloc(dim * sizeof(float));
if (!vec) break;
for (int d = 0; d < dim; d++) {
if (offset < payload_size) {
vec[d] = ((float)(int8_t)payload[offset++]) / 20.0f;
} else {
vec[d] = (float)((i * dim + d) % 256 - 128) / 128.0f;
}
}
sqlite3_reset(stmtInsert);
sqlite3_bind_int64(stmtInsert, 1, (int64_t)(i + 1));
sqlite3_bind_blob(stmtInsert, 2, vec, dim * sizeof(float), SQLITE_TRANSIENT);
sqlite3_step(stmtInsert);
sqlite3_free(vec);
}
sqlite3_finalize(stmtInsert);
// Query BEFORE training (flat scan path)
{
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (qvec) {
for (int d = 0; d < dim; d++) qvec[d] = 0.5f;
sqlite3_stmt *sk = NULL;
snprintf(sql, sizeof(sql),
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT %d", k_limit);
sqlite3_prepare_v2(db, sql, -1, &sk, NULL);
if (sk) {
sqlite3_bind_blob(sk, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(sk) == SQLITE_ROW) {}
sqlite3_finalize(sk);
}
sqlite3_free(qvec);
}
}
// Train
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('compute-centroids')",
NULL, NULL, NULL);
// Change nprobe at runtime (can exceed nlist -- tests clamping in query)
{
char cmd[64];
snprintf(cmd, sizeof(cmd),
"INSERT INTO v(rowid) VALUES ('nprobe=%d')", nprobe_initial);
sqlite3_exec(db, cmd, NULL, NULL, NULL);
}
// Multiple KNN queries with different fuzz-derived query vectors
for (int q = 0; q < num_queries; q++) {
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (!qvec) break;
for (int d = 0; d < dim; d++) {
if (offset < payload_size) {
qvec[d] = ((float)(int8_t)payload[offset++]) / 10.0f;
} else {
qvec[d] = (q == 0) ? 1.0f : 0.0f;
}
}
sqlite3_stmt *sk = NULL;
snprintf(sql, sizeof(sql),
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT %d", k_limit);
sqlite3_prepare_v2(db, sql, -1, &sk, NULL);
if (sk) {
sqlite3_bind_blob(sk, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(sk) == SQLITE_ROW) {}
sqlite3_finalize(sk);
}
sqlite3_free(qvec);
}
// Delete half the vectors then query again
for (int i = 1; i <= num_vecs / 2; i++) {
char delsql[64];
snprintf(delsql, sizeof(delsql), "DELETE FROM v WHERE rowid = %d", i);
sqlite3_exec(db, delsql, NULL, NULL, NULL);
}
// Query after mass deletion
{
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (qvec) {
for (int d = 0; d < dim; d++) qvec[d] = -0.5f;
sqlite3_stmt *sk = NULL;
snprintf(sql, sizeof(sql),
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT %d", k_limit);
sqlite3_prepare_v2(db, sql, -1, &sk, NULL);
if (sk) {
sqlite3_bind_blob(sk, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(sk) == SQLITE_ROW) {}
sqlite3_finalize(sk);
}
sqlite3_free(qvec);
}
}
sqlite3_close(db);
return 0;
}

121
tests/fuzz/ivf-operations.c Normal file
View file

@ -0,0 +1,121 @@
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "sqlite-vec.h"
#include "sqlite3.h"
#include <assert.h>
int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
if (size < 6) return 0;
int rc;
sqlite3 *db;
sqlite3_stmt *stmtInsert = NULL;
sqlite3_stmt *stmtDelete = NULL;
sqlite3_stmt *stmtKnn = NULL;
sqlite3_stmt *stmtScan = NULL;
rc = sqlite3_open(":memory:", &db);
assert(rc == SQLITE_OK);
rc = sqlite3_vec_init(db, NULL, NULL);
assert(rc == SQLITE_OK);
rc = sqlite3_exec(db,
"CREATE VIRTUAL TABLE v USING vec0(emb float[4] indexed by ivf(nlist=4, nprobe=4))",
NULL, NULL, NULL);
if (rc != SQLITE_OK) { sqlite3_close(db); return 0; }
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &stmtInsert, NULL);
sqlite3_prepare_v2(db,
"DELETE FROM v WHERE rowid = ?", -1, &stmtDelete, NULL);
sqlite3_prepare_v2(db,
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT 3",
-1, &stmtKnn, NULL);
sqlite3_prepare_v2(db,
"SELECT rowid FROM v", -1, &stmtScan, NULL);
if (!stmtInsert || !stmtDelete || !stmtKnn || !stmtScan) goto cleanup;
size_t i = 0;
while (i + 2 <= size) {
uint8_t op = data[i++] % 7;
uint8_t rowid_byte = data[i++];
int64_t rowid = (int64_t)(rowid_byte % 32) + 1;
switch (op) {
case 0: {
// INSERT: consume 16 bytes for 4 floats, or use what's left
float vec[4] = {0.0f, 0.0f, 0.0f, 0.0f};
for (int j = 0; j < 4 && i < size; j++, i++) {
vec[j] = (float)((int8_t)data[i]) / 10.0f;
}
sqlite3_reset(stmtInsert);
sqlite3_bind_int64(stmtInsert, 1, rowid);
sqlite3_bind_blob(stmtInsert, 2, vec, sizeof(vec), SQLITE_TRANSIENT);
sqlite3_step(stmtInsert);
break;
}
case 1: {
// DELETE
sqlite3_reset(stmtDelete);
sqlite3_bind_int64(stmtDelete, 1, rowid);
sqlite3_step(stmtDelete);
break;
}
case 2: {
// KNN query with a fixed query vector
float qvec[4] = {1.0f, 0.0f, 0.0f, 0.0f};
sqlite3_reset(stmtKnn);
sqlite3_bind_blob(stmtKnn, 1, qvec, sizeof(qvec), SQLITE_STATIC);
while (sqlite3_step(stmtKnn) == SQLITE_ROW) {}
break;
}
case 3: {
// Full scan
sqlite3_reset(stmtScan);
while (sqlite3_step(stmtScan) == SQLITE_ROW) {}
break;
}
case 4: {
// compute-centroids command
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('compute-centroids')",
NULL, NULL, NULL);
break;
}
case 5: {
// clear-centroids command
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('clear-centroids')",
NULL, NULL, NULL);
break;
}
case 6: {
// nprobe=N command
if (i < size) {
uint8_t n = data[i++];
int nprobe = (n % 4) + 1;
char buf[64];
snprintf(buf, sizeof(buf),
"INSERT INTO v(rowid) VALUES ('nprobe=%d')", nprobe);
sqlite3_exec(db, buf, NULL, NULL, NULL);
}
break;
}
}
}
// Final operations — must not crash regardless of prior state
sqlite3_exec(db, "SELECT * FROM v", NULL, NULL, NULL);
cleanup:
sqlite3_finalize(stmtInsert);
sqlite3_finalize(stmtDelete);
sqlite3_finalize(stmtKnn);
sqlite3_finalize(stmtScan);
sqlite3_close(db);
return 0;
}

129
tests/fuzz/ivf-quantize.c Normal file
View file

@ -0,0 +1,129 @@
/**
* Fuzz target: IVF quantization functions.
*
* Directly exercises ivf_quantize_int8 and ivf_quantize_binary with
* fuzz-controlled dimensions and float data. Targets:
* - ivf_quantize_int8: clamping, int8 overflow boundary
* - ivf_quantize_binary: D not divisible by 8, memset(D/8) undercount
* - Round-trip through CREATE TABLE + INSERT with quantized IVF
*/
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "sqlite-vec.h"
#include "sqlite3.h"
#include <assert.h>
int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
if (size < 8) return 0;
int rc;
sqlite3 *db;
rc = sqlite3_open(":memory:", &db);
assert(rc == SQLITE_OK);
rc = sqlite3_vec_init(db, NULL, NULL);
assert(rc == SQLITE_OK);
// Byte 0: quantizer type (0=int8, 1=binary)
// Byte 1: dimension (1..64, but we test edge cases)
// Byte 2: nlist (1..8)
// Byte 3: num_vectors to insert (1..32)
// Remaining: float data
int qtype = data[0] % 2;
int dim = (data[1] % 64) + 1;
int nlist = (data[2] % 8) + 1;
int num_vecs = (data[3] % 32) + 1;
const uint8_t *payload = data + 4;
size_t payload_size = size - 4;
// For binary quantizer, D must be multiple of 8 to avoid the D/8 bug
// in production. But we explicitly want to test non-multiples too to
// find the bug. Use dim as-is.
const char *quantizer = qtype ? "binary" : "int8";
// Binary quantizer needs D multiple of 8 in current code, but let's
// test both valid and invalid dimensions to see what happens.
// For binary with non-multiple-of-8, the code does memset(dst, 0, D/8)
// which underallocates when D%8 != 0.
char sql[256];
snprintf(sql, sizeof(sql),
"CREATE VIRTUAL TABLE v USING vec0("
"emb float[%d] indexed by ivf(nlist=%d, nprobe=%d, quantizer=%s))",
dim, nlist, nlist, quantizer);
rc = sqlite3_exec(db, sql, NULL, NULL, NULL);
if (rc != SQLITE_OK) { sqlite3_close(db); return 0; }
// Insert vectors with fuzz-controlled float values
sqlite3_stmt *stmtInsert = NULL;
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &stmtInsert, NULL);
if (!stmtInsert) { sqlite3_close(db); return 0; }
size_t offset = 0;
for (int i = 0; i < num_vecs && offset < payload_size; i++) {
// Build float vector from fuzz data
float *vec = sqlite3_malloc(dim * sizeof(float));
if (!vec) break;
for (int d = 0; d < dim; d++) {
if (offset + 4 <= payload_size) {
// Use raw bytes as float -- can produce NaN, Inf, denormals
memcpy(&vec[d], payload + offset, sizeof(float));
offset += 4;
} else if (offset < payload_size) {
// Partial: use byte as scaled value
vec[d] = ((float)(int8_t)payload[offset++]) / 50.0f;
} else {
vec[d] = 0.0f;
}
}
sqlite3_reset(stmtInsert);
sqlite3_bind_int64(stmtInsert, 1, (int64_t)(i + 1));
sqlite3_bind_blob(stmtInsert, 2, vec, dim * sizeof(float), SQLITE_TRANSIENT);
sqlite3_step(stmtInsert);
sqlite3_free(vec);
}
sqlite3_finalize(stmtInsert);
// Trigger compute-centroids to exercise kmeans + quantization together
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('compute-centroids')",
NULL, NULL, NULL);
// KNN query with fuzz-derived query vector
{
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (qvec) {
for (int d = 0; d < dim; d++) {
if (offset < payload_size) {
qvec[d] = ((float)(int8_t)payload[offset++]) / 10.0f;
} else {
qvec[d] = 1.0f;
}
}
sqlite3_stmt *stmtKnn = NULL;
sqlite3_prepare_v2(db,
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT 5",
-1, &stmtKnn, NULL);
if (stmtKnn) {
sqlite3_bind_blob(stmtKnn, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(stmtKnn) == SQLITE_ROW) {}
sqlite3_finalize(stmtKnn);
}
sqlite3_free(qvec);
}
}
// Full scan
sqlite3_exec(db, "SELECT * FROM v", NULL, NULL, NULL);
sqlite3_close(db);
return 0;
}

182
tests/fuzz/ivf-rescore.c Normal file
View file

@ -0,0 +1,182 @@
/**
* Fuzz target: IVF oversample + rescore path.
*
* Specifically targets the code path where quantizer != none AND
* oversample > 1, which triggers:
* 1. Quantized KNN scan to collect oversample*k candidates
* 2. Full-precision vector lookup from _ivf_vectors table
* 3. Re-scoring with float32 distances
* 4. Re-sort and truncation
*
* This path has the most complex memory management in the KNN query:
* - Two separate distance computations (quantized + float)
* - Cross-table lookups (cells + vectors KV store)
* - Candidate array resizing
* - qsort over partially re-scored arrays
*
* Also tests the int8 + binary quantization round-trip fidelity
* under adversarial float inputs.
*/
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "sqlite-vec.h"
#include "sqlite3.h"
#include <assert.h>
int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
if (size < 12) return 0;
int rc;
sqlite3 *db;
rc = sqlite3_open(":memory:", &db);
assert(rc == SQLITE_OK);
rc = sqlite3_vec_init(db, NULL, NULL);
assert(rc == SQLITE_OK);
// Header
int quantizer_type = (data[0] % 2) + 1; // 1=int8, 2=binary (never none)
int dim = (data[1] % 32) + 8; // 8..39
int nlist = (data[2] % 8) + 1; // 1..8
int oversample = (data[3] % 4) + 2; // 2..5 (always > 1)
int num_vecs = (data[4] % 60) + 8; // 8..67
int k_limit = (data[5] % 15) + 1; // 1..15
const uint8_t *payload = data + 6;
size_t payload_size = size - 6;
// Binary quantizer needs D multiple of 8
if (quantizer_type == 2) {
dim = ((dim + 7) / 8) * 8;
}
const char *qname = (quantizer_type == 1) ? "int8" : "binary";
char sql[512];
snprintf(sql, sizeof(sql),
"CREATE VIRTUAL TABLE v USING vec0("
"emb float[%d] indexed by ivf(nlist=%d, nprobe=%d, quantizer=%s, oversample=%d))",
dim, nlist, nlist, qname, oversample);
rc = sqlite3_exec(db, sql, NULL, NULL, NULL);
if (rc != SQLITE_OK) { sqlite3_close(db); return 0; }
// Insert vectors with diverse values
sqlite3_stmt *stmtInsert = NULL;
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &stmtInsert, NULL);
if (!stmtInsert) { sqlite3_close(db); return 0; }
size_t offset = 0;
for (int i = 0; i < num_vecs; i++) {
float *vec = sqlite3_malloc(dim * sizeof(float));
if (!vec) break;
for (int d = 0; d < dim; d++) {
if (offset + 4 <= payload_size) {
// Use raw bytes as float for adversarial values
memcpy(&vec[d], payload + offset, sizeof(float));
offset += 4;
// Sanitize: replace NaN/Inf with bounded values to avoid
// poisoning the entire computation. We want edge values,
// not complete nonsense.
if (isnan(vec[d]) || isinf(vec[d])) {
vec[d] = (vec[d] > 0) ? 1e6f : -1e6f;
if (isnan(vec[d])) vec[d] = 0.0f;
}
} else if (offset < payload_size) {
vec[d] = ((float)(int8_t)payload[offset++]) / 30.0f;
} else {
vec[d] = (float)(i * dim + d) * 0.001f;
}
}
sqlite3_reset(stmtInsert);
sqlite3_bind_int64(stmtInsert, 1, (int64_t)(i + 1));
sqlite3_bind_blob(stmtInsert, 2, vec, dim * sizeof(float), SQLITE_TRANSIENT);
sqlite3_step(stmtInsert);
sqlite3_free(vec);
}
sqlite3_finalize(stmtInsert);
// Train
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('compute-centroids')",
NULL, NULL, NULL);
// Multiple KNN queries to exercise rescore path
for (int q = 0; q < 4; q++) {
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (!qvec) break;
for (int d = 0; d < dim; d++) {
if (offset < payload_size) {
qvec[d] = ((float)(int8_t)payload[offset++]) / 10.0f;
} else {
qvec[d] = (q == 0) ? 1.0f : (q == 1) ? -1.0f : 0.0f;
}
}
sqlite3_stmt *sk = NULL;
snprintf(sql, sizeof(sql),
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT %d", k_limit);
sqlite3_prepare_v2(db, sql, -1, &sk, NULL);
if (sk) {
sqlite3_bind_blob(sk, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(sk) == SQLITE_ROW) {}
sqlite3_finalize(sk);
}
sqlite3_free(qvec);
}
// Delete some vectors, then query again (rescore with missing _ivf_vectors rows)
for (int i = 1; i <= num_vecs / 3; i++) {
char delsql[64];
snprintf(delsql, sizeof(delsql), "DELETE FROM v WHERE rowid = %d", i);
sqlite3_exec(db, delsql, NULL, NULL, NULL);
}
{
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (qvec) {
for (int d = 0; d < dim; d++) qvec[d] = 0.5f;
sqlite3_stmt *sk = NULL;
snprintf(sql, sizeof(sql),
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT %d", k_limit);
sqlite3_prepare_v2(db, sql, -1, &sk, NULL);
if (sk) {
sqlite3_bind_blob(sk, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(sk) == SQLITE_ROW) {}
sqlite3_finalize(sk);
}
sqlite3_free(qvec);
}
}
// Retrain after deletions
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('compute-centroids')",
NULL, NULL, NULL);
// Query after retrain
{
float *qvec = sqlite3_malloc(dim * sizeof(float));
if (qvec) {
for (int d = 0; d < dim; d++) qvec[d] = -0.3f;
sqlite3_stmt *sk = NULL;
snprintf(sql, sizeof(sql),
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT %d", k_limit);
sqlite3_prepare_v2(db, sql, -1, &sk, NULL);
if (sk) {
sqlite3_bind_blob(sk, 1, qvec, dim * sizeof(float), SQLITE_TRANSIENT);
while (sqlite3_step(sk) == SQLITE_ROW) {}
sqlite3_finalize(sk);
}
sqlite3_free(qvec);
}
}
sqlite3_close(db);
return 0;
}

228
tests/fuzz/ivf-shadow-corrupt.c Normal file
View file

@ -0,0 +1,228 @@
/**
* Fuzz target: IVF shadow table corruption.
*
* Creates a trained IVF table, then corrupts IVF shadow table blobs
* (centroids, cells validity/rowids/vectors, rowid_map) with fuzz data.
* Then exercises all read/write paths. Must not crash.
*
* Targets:
* - Cell validity bitmap with wrong size
* - Cell rowids blob with wrong size/alignment
* - Cell vectors blob with wrong size
* - Centroid blob with wrong size
* - n_vectors inconsistent with validity bitmap
* - Missing rowid_map entries
* - KNN scan over corrupted cells
* - Insert/delete with corrupted rowid_map
*/
#include <stdint.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "sqlite-vec.h"
#include "sqlite3.h"
#include <assert.h>
int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
if (size < 4) return 0;
int rc;
sqlite3 *db;
rc = sqlite3_open(":memory:", &db);
assert(rc == SQLITE_OK);
rc = sqlite3_vec_init(db, NULL, NULL);
assert(rc == SQLITE_OK);
// Create IVF table and insert enough vectors to train
rc = sqlite3_exec(db,
"CREATE VIRTUAL TABLE v USING vec0("
"emb float[8] indexed by ivf(nlist=2, nprobe=2))",
NULL, NULL, NULL);
if (rc != SQLITE_OK) { sqlite3_close(db); return 0; }
// Insert 10 vectors
{
sqlite3_stmt *si = NULL;
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &si, NULL);
if (!si) { sqlite3_close(db); return 0; }
for (int i = 0; i < 10; i++) {
float vec[8];
for (int d = 0; d < 8; d++) {
vec[d] = (float)(i * 8 + d) * 0.1f;
}
sqlite3_reset(si);
sqlite3_bind_int64(si, 1, i + 1);
sqlite3_bind_blob(si, 2, vec, sizeof(vec), SQLITE_TRANSIENT);
sqlite3_step(si);
}
sqlite3_finalize(si);
}
// Train
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('compute-centroids')",
NULL, NULL, NULL);
// Now corrupt shadow tables based on fuzz input
uint8_t target = data[0] % 10;
const uint8_t *payload = data + 1;
int payload_size = (int)(size - 1);
// Limit payload to avoid huge allocations
if (payload_size > 4096) payload_size = 4096;
sqlite3_stmt *stmt = NULL;
switch (target) {
case 0: {
// Corrupt cell validity blob
rc = sqlite3_prepare_v2(db,
"UPDATE v_ivf_cells00 SET validity = ? WHERE rowid = 1",
-1, &stmt, NULL);
if (rc == SQLITE_OK) {
sqlite3_bind_blob(stmt, 1, payload, payload_size, SQLITE_STATIC);
sqlite3_step(stmt); sqlite3_finalize(stmt);
}
break;
}
case 1: {
// Corrupt cell rowids blob
rc = sqlite3_prepare_v2(db,
"UPDATE v_ivf_cells00 SET rowids = ? WHERE rowid = 1",
-1, &stmt, NULL);
if (rc == SQLITE_OK) {
sqlite3_bind_blob(stmt, 1, payload, payload_size, SQLITE_STATIC);
sqlite3_step(stmt); sqlite3_finalize(stmt);
}
break;
}
case 2: {
// Corrupt cell vectors blob
rc = sqlite3_prepare_v2(db,
"UPDATE v_ivf_cells00 SET vectors = ? WHERE rowid = 1",
-1, &stmt, NULL);
if (rc == SQLITE_OK) {
sqlite3_bind_blob(stmt, 1, payload, payload_size, SQLITE_STATIC);
sqlite3_step(stmt); sqlite3_finalize(stmt);
}
break;
}
case 3: {
// Corrupt centroid blob
rc = sqlite3_prepare_v2(db,
"UPDATE v_ivf_centroids00 SET centroid = ? WHERE centroid_id = 0",
-1, &stmt, NULL);
if (rc == SQLITE_OK) {
sqlite3_bind_blob(stmt, 1, payload, payload_size, SQLITE_STATIC);
sqlite3_step(stmt); sqlite3_finalize(stmt);
}
break;
}
case 4: {
// Set n_vectors to a bogus value (larger than cell capacity)
int bogus_n = 99999;
if (payload_size >= 4) {
memcpy(&bogus_n, payload, 4);
bogus_n = abs(bogus_n) % 100000;
}
char sql[128];
snprintf(sql, sizeof(sql),
"UPDATE v_ivf_cells00 SET n_vectors = %d WHERE rowid = 1", bogus_n);
sqlite3_exec(db, sql, NULL, NULL, NULL);
break;
}
case 5: {
// Delete rowid_map entries (orphan vectors)
sqlite3_exec(db,
"DELETE FROM v_ivf_rowid_map00 WHERE rowid IN (1, 2, 3)",
NULL, NULL, NULL);
break;
}
case 6: {
// Corrupt rowid_map slot values
char sql[128];
int bogus_slot = payload_size > 0 ? (int)payload[0] * 1000 : 99999;
snprintf(sql, sizeof(sql),
"UPDATE v_ivf_rowid_map00 SET slot = %d WHERE rowid = 1", bogus_slot);
sqlite3_exec(db, sql, NULL, NULL, NULL);
break;
}
case 7: {
// Corrupt rowid_map cell_id values
sqlite3_exec(db,
"UPDATE v_ivf_rowid_map00 SET cell_id = 99999 WHERE rowid = 1",
NULL, NULL, NULL);
break;
}
case 8: {
// Delete all centroids (make trained but no centroids)
sqlite3_exec(db,
"DELETE FROM v_ivf_centroids00",
NULL, NULL, NULL);
break;
}
case 9: {
// Set validity to NULL
sqlite3_exec(db,
"UPDATE v_ivf_cells00 SET validity = NULL WHERE rowid = 1",
NULL, NULL, NULL);
break;
}
}
// Exercise all read paths over corrupted state — must not crash
float qvec[8] = {1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
// KNN query
{
sqlite3_stmt *sk = NULL;
sqlite3_prepare_v2(db,
"SELECT rowid, distance FROM v WHERE emb MATCH ? LIMIT 5",
-1, &sk, NULL);
if (sk) {
sqlite3_bind_blob(sk, 1, qvec, sizeof(qvec), SQLITE_STATIC);
while (sqlite3_step(sk) == SQLITE_ROW) {}
sqlite3_finalize(sk);
}
}
// Full scan
sqlite3_exec(db, "SELECT * FROM v", NULL, NULL, NULL);
// Point query
sqlite3_exec(db, "SELECT * FROM v WHERE rowid = 1", NULL, NULL, NULL);
sqlite3_exec(db, "SELECT * FROM v WHERE rowid = 5", NULL, NULL, NULL);
// Delete
sqlite3_exec(db, "DELETE FROM v WHERE rowid = 3", NULL, NULL, NULL);
// Insert after corruption
{
float newvec[8] = {0.5f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f, 0.5f};
sqlite3_stmt *si = NULL;
sqlite3_prepare_v2(db,
"INSERT INTO v(rowid, emb) VALUES (?, ?)", -1, &si, NULL);
if (si) {
sqlite3_bind_int64(si, 1, 100);
sqlite3_bind_blob(si, 2, newvec, sizeof(newvec), SQLITE_STATIC);
sqlite3_step(si);
sqlite3_finalize(si);
}
}
// compute-centroids over corrupted state
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('compute-centroids')",
NULL, NULL, NULL);
// clear-centroids
sqlite3_exec(db,
"INSERT INTO v(rowid) VALUES ('clear-centroids')",
NULL, NULL, NULL);
sqlite3_close(db);
return 0;
}

37
tests/sqlite-vec-internal.h
View file

@ -5,6 +5,10 @@
#include <stddef.h>
#include <stdint.h>
#ifndef SQLITE_VEC_ENABLE_IVF
#define SQLITE_VEC_ENABLE_IVF 1
#endif
int min_idx(
const float *distances,
int32_t n,
@ -68,8 +72,36 @@ enum Vec0IndexType {
#ifdef SQLITE_VEC_ENABLE_RESCORE
VEC0_INDEX_TYPE_RESCORE = 2,
#endif
VEC0_INDEX_TYPE_IVF = 3,
};
enum Vec0RescoreQuantizerType {
VEC0_RESCORE_QUANTIZER_BIT = 1,
VEC0_RESCORE_QUANTIZER_INT8 = 2,
};
struct Vec0RescoreConfig {
enum Vec0RescoreQuantizerType quantizer_type;
int oversample;
};
#if SQLITE_VEC_ENABLE_IVF
enum Vec0IvfQuantizer {
VEC0_IVF_QUANTIZER_NONE = 0,
VEC0_IVF_QUANTIZER_INT8 = 1,
VEC0_IVF_QUANTIZER_BINARY = 2,
};
struct Vec0IvfConfig {
int nlist;
int nprobe;
int quantizer;
int oversample;
};
#else
struct Vec0IvfConfig { char _unused; };
#endif
#ifdef SQLITE_VEC_ENABLE_RESCORE
enum Vec0RescoreQuantizerType {
VEC0_RESCORE_QUANTIZER_BIT = 1,
@ -93,6 +125,7 @@ struct VectorColumnDefinition {
#ifdef SQLITE_VEC_ENABLE_RESCORE
struct Vec0RescoreConfig rescore;
#endif
struct Vec0IvfConfig ivf;
};
int vec0_parse_vector_column(const char *source, int source_length,
@ -114,6 +147,10 @@ void _test_rescore_quantize_float_to_int8(const float *src, int8_t *dst, size_t
size_t _test_rescore_quantized_byte_size_bit(size_t dimensions);
size_t _test_rescore_quantized_byte_size_int8(size_t dimensions);
#endif
#if SQLITE_VEC_ENABLE_IVF
void ivf_quantize_int8(const float *src, int8_t *dst, int D);
void ivf_quantize_binary(const float *src, uint8_t *dst, int D);
#endif
#endif
#endif /* SQLITE_VEC_INTERNAL_H */

575
tests/test-ivf-mutations.py Normal file
View file

@ -0,0 +1,575 @@
"""
Thorough IVF mutation tests: insert, delete, update, KNN correctness,
error cases, edge cases, and cell overflow scenarios.
"""
import pytest
import sqlite3
import struct
import math
from helpers import _f32, exec
@pytest.fixture()
def db():
db = sqlite3.connect(":memory:")
db.row_factory = sqlite3.Row
db.enable_load_extension(True)
db.load_extension("dist/vec0")
db.enable_load_extension(False)
return db
def ivf_total_vectors(db, table="t", col=0):
"""Count total vectors across all IVF cells."""
return db.execute(
f"SELECT COALESCE(SUM(n_vectors), 0) FROM {table}_ivf_cells{col:02d}"
).fetchone()[0]
def ivf_unassigned_count(db, table="t", col=0):
return db.execute(
f"SELECT COALESCE(SUM(n_vectors), 0) FROM {table}_ivf_cells{col:02d} WHERE centroid_id = -1"
).fetchone()[0]
def ivf_assigned_count(db, table="t", col=0):
return db.execute(
f"SELECT COALESCE(SUM(n_vectors), 0) FROM {table}_ivf_cells{col:02d} WHERE centroid_id >= 0"
).fetchone()[0]
def knn(db, query, k, table="t", col="v"):
"""Run a KNN query and return list of (rowid, distance) tuples."""
rows = db.execute(
f"SELECT rowid, distance FROM {table} WHERE {col} MATCH ? AND k = ?",
[_f32(query), k],
).fetchall()
return [(r[0], r[1]) for r in rows]
# ============================================================================
# Single row insert + KNN
# ============================================================================
def test_insert_single_row_knn(db):
db.execute("CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf())")
db.execute("INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([1, 0, 0, 0])])
results = knn(db, [1, 0, 0, 0], 5)
assert len(results) == 1
assert results[0][0] == 1
assert results[0][1] < 0.001
# ============================================================================
# Batch insert + KNN recall
# ============================================================================
def test_batch_insert_knn_recall(db):
"""Insert 200 vectors, train, verify KNN recall with nprobe=nlist."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=8, nprobe=8))"
)
for i in range(200):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
assert ivf_total_vectors(db) == 200
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
assert ivf_assigned_count(db) == 200
# Query near 100 -- closest should be rowid 100
results = knn(db, [100.0, 0, 0, 0], 10)
assert len(results) == 10
assert results[0][0] == 100
assert results[0][1] < 0.01
# All results should be near 100
rowids = {r[0] for r in results}
assert all(95 <= r <= 105 for r in rowids)
# ============================================================================
# Delete rows, verify they're gone from KNN
# ============================================================================
def test_delete_rows_gone_from_knn(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([float(i), 0, 0, 0])],
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
# Delete rowid 10
db.execute("DELETE FROM t WHERE rowid = 10")
results = knn(db, [10.0, 0, 0, 0], 20)
rowids = {r[0] for r in results}
assert 10 not in rowids
def test_delete_all_rows_empty_results(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2, nprobe=2))"
)
for i in range(10):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
for i in range(10):
db.execute("DELETE FROM t WHERE rowid = ?", [i])
assert ivf_total_vectors(db) == 0
results = knn(db, [5.0, 0, 0, 0], 10)
assert len(results) == 0
# ============================================================================
# Insert after delete (reuse rowids)
# ============================================================================
def test_insert_after_delete_reuse_rowid(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2, nprobe=2))"
)
for i in range(10):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
# Delete rowid 5
db.execute("DELETE FROM t WHERE rowid = 5")
# Re-insert rowid 5 with a very different vector
db.execute(
"INSERT INTO t(rowid, v) VALUES (5, ?)", [_f32([999.0, 0, 0, 0])]
)
# KNN near 999 should find rowid 5
results = knn(db, [999.0, 0, 0, 0], 1)
assert len(results) >= 1
assert results[0][0] == 5
# ============================================================================
# Update vectors (INSERT OR REPLACE), verify KNN reflects new values
# ============================================================================
def test_update_vector_via_delete_insert(db):
"""vec0 IVF update: delete then re-insert with new vector."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2, nprobe=2))"
)
for i in range(10):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
# "Update" rowid 3: delete and re-insert with new vector
db.execute("DELETE FROM t WHERE rowid = 3")
db.execute(
"INSERT INTO t(rowid, v) VALUES (3, ?)",
[_f32([100.0, 0, 0, 0])],
)
# KNN near 100 should find rowid 3
results = knn(db, [100.0, 0, 0, 0], 1)
assert results[0][0] == 3
# ============================================================================
# Error cases: IVF + auxiliary/metadata/partition key columns
# ============================================================================
def test_error_ivf_with_auxiliary_column(db):
result = exec(
db,
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(), +extra text)",
)
assert "error" in result
assert "auxiliary" in result.get("message", "").lower()
def test_error_ivf_with_metadata_column(db):
result = exec(
db,
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(), genre text)",
)
assert "error" in result
assert "metadata" in result.get("message", "").lower()
def test_error_ivf_with_partition_key(db):
result = exec(
db,
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(), user_id integer partition key)",
)
assert "error" in result
assert "partition" in result.get("message", "").lower()
def test_flat_with_auxiliary_still_works(db):
"""Regression guard: flat-indexed tables with aux columns should still work."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4], +extra text)"
)
db.execute(
"INSERT INTO t(rowid, v, extra) VALUES (1, ?, 'hello')",
[_f32([1, 0, 0, 0])],
)
row = db.execute("SELECT extra FROM t WHERE rowid = 1").fetchone()
assert row[0] == "hello"
def test_flat_with_metadata_still_works(db):
"""Regression guard: flat-indexed tables with metadata columns should still work."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4], genre text)"
)
db.execute(
"INSERT INTO t(rowid, v, genre) VALUES (1, ?, 'rock')",
[_f32([1, 0, 0, 0])],
)
row = db.execute("SELECT genre FROM t WHERE rowid = 1").fetchone()
assert row[0] == "rock"
def test_flat_with_partition_key_still_works(db):
"""Regression guard: flat-indexed tables with partition key should still work."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4], user_id integer partition key)"
)
db.execute(
"INSERT INTO t(rowid, v, user_id) VALUES (1, ?, 42)",
[_f32([1, 0, 0, 0])],
)
row = db.execute("SELECT user_id FROM t WHERE rowid = 1").fetchone()
assert row[0] == 42
# ============================================================================
# Edge cases
# ============================================================================
def test_zero_vectors(db):
"""Insert zero vectors, verify KNN still works."""
db.execute("CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf())")
for i in range(5):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([0, 0, 0, 0])],
)
results = knn(db, [0, 0, 0, 0], 5)
assert len(results) == 5
# All distances should be 0
for _, dist in results:
assert dist < 0.001
def test_large_values(db):
"""Insert vectors with very large and small values."""
db.execute("CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf())")
db.execute(
"INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([1e6, 1e6, 1e6, 1e6])]
)
db.execute(
"INSERT INTO t(rowid, v) VALUES (2, ?)", [_f32([1e-6, 1e-6, 1e-6, 1e-6])]
)
db.execute(
"INSERT INTO t(rowid, v) VALUES (3, ?)", [_f32([-1e6, -1e6, -1e6, -1e6])]
)
results = knn(db, [1e6, 1e6, 1e6, 1e6], 3)
assert results[0][0] == 1
def test_single_row_compute_centroids(db):
"""Single row table, compute-centroids should still work."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=1))"
)
db.execute(
"INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([1, 2, 3, 4])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
assert ivf_assigned_count(db) == 1
results = knn(db, [1, 2, 3, 4], 1)
assert len(results) == 1
assert results[0][0] == 1
# ============================================================================
# Cell overflow (many vectors in one cell)
# ============================================================================
def test_cell_overflow_many_vectors(db):
"""Insert >64 vectors that all go to same centroid. Should create multiple cells."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=0))"
)
# Insert 100 very similar vectors
for i in range(100):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([1.0 + i * 0.001, 0, 0, 0])],
)
# Set a single centroid so all vectors go there
db.execute(
"INSERT INTO t(rowid, v) VALUES ('set-centroid:0', ?)",
[_f32([1.0, 0, 0, 0])],
)
db.execute("INSERT INTO t(rowid) VALUES ('assign-vectors')")
assert ivf_assigned_count(db) == 100
# Should have more than 1 cell (64 max per cell)
cell_count = db.execute(
"SELECT count(*) FROM t_ivf_cells00 WHERE centroid_id = 0"
).fetchone()[0]
assert cell_count >= 2 # 100 / 64 = 2 cells needed
# All vectors should be queryable
results = knn(db, [1.0, 0, 0, 0], 100)
assert len(results) == 100
# ============================================================================
# Large batch with training
# ============================================================================
def test_large_batch_with_training(db):
"""Insert 500, train, insert 500 more, verify total is 1000."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=16, nprobe=16))"
)
for i in range(500):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
for i in range(500, 1000):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
assert ivf_total_vectors(db) == 1000
# KNN should still work
results = knn(db, [750.0, 0, 0, 0], 5)
assert len(results) == 5
assert results[0][0] == 750
# ============================================================================
# KNN after interleaved insert/delete
# ============================================================================
def test_knn_after_interleaved_insert_delete(db):
"""Insert 20, train, delete 10 closest to query, verify remaining."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4, nprobe=4))"
)
for i in range(20):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
# Delete rowids 0-9 (closest to query at 5.0)
for i in range(10):
db.execute("DELETE FROM t WHERE rowid = ?", [i])
results = knn(db, [5.0, 0, 0, 0], 10)
rowids = {r[0] for r in results}
# None of the deleted rowids should appear
assert all(r >= 10 for r in rowids)
assert len(results) == 10
def test_knn_empty_centroids_after_deletes(db):
"""Some centroids may become empty after deletes. Should not crash."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4, nprobe=2))"
)
# Insert vectors clustered near 4 points
for i in range(40):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i % 10) * 10, 0, 0, 0])],
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
# Delete a bunch, potentially emptying some centroids
for i in range(30):
db.execute("DELETE FROM t WHERE rowid = ?", [i])
# Should not crash even with empty centroids
results = knn(db, [50.0, 0, 0, 0], 20)
assert len(results) <= 10 # only 10 left
# ============================================================================
# KNN returns correct distances
# ============================================================================
def test_knn_correct_distances(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2, nprobe=2))"
)
db.execute("INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([0, 0, 0, 0])])
db.execute("INSERT INTO t(rowid, v) VALUES (2, ?)", [_f32([3, 0, 0, 0])])
db.execute("INSERT INTO t(rowid, v) VALUES (3, ?)", [_f32([0, 4, 0, 0])])
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
results = knn(db, [0, 0, 0, 0], 3)
result_map = {r[0]: r[1] for r in results}
# L2 distances (sqrt of sum of squared differences)
assert abs(result_map[1] - 0.0) < 0.01
assert abs(result_map[2] - 3.0) < 0.01 # sqrt(3^2) = 3
assert abs(result_map[3] - 4.0) < 0.01 # sqrt(4^2) = 4
# ============================================================================
# Delete in flat mode leaves no orphan rowid_map entries
# ============================================================================
def test_delete_flat_mode_rowid_map_count(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4))"
)
for i in range(5):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
db.execute("DELETE FROM t WHERE rowid = 0")
db.execute("DELETE FROM t WHERE rowid = 2")
db.execute("DELETE FROM t WHERE rowid = 4")
assert db.execute("SELECT count(*) FROM t_ivf_rowid_map00").fetchone()[0] == 2
assert ivf_unassigned_count(db) == 2
# ============================================================================
# Duplicate rowid insert
# ============================================================================
def test_delete_reinsert_as_update(db):
"""Simulate update via delete + insert on same rowid."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2, nprobe=2))"
)
db.execute("INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([1, 0, 0, 0])])
# Delete then re-insert as "update"
db.execute("DELETE FROM t WHERE rowid = 1")
db.execute("INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([99, 0, 0, 0])])
results = knn(db, [99, 0, 0, 0], 1)
assert len(results) == 1
assert results[0][0] == 1
assert results[0][1] < 0.01
def test_duplicate_rowid_insert_fails(db):
"""Inserting a duplicate rowid should fail with a constraint error."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2, nprobe=2))"
)
db.execute("INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([1, 0, 0, 0])])
result = exec(
db,
"INSERT INTO t(rowid, v) VALUES (1, ?)",
[_f32([99, 0, 0, 0])],
)
assert "error" in result
# ============================================================================
# Interleaved insert/delete with KNN correctness
# ============================================================================
def test_interleaved_ops_correctness(db):
"""Complex sequence of inserts and deletes, verify KNN is always correct."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4, nprobe=4))"
)
# Phase 1: Insert 50 vectors
for i in range(50):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
# Phase 2: Delete even-numbered rowids
for i in range(0, 50, 2):
db.execute("DELETE FROM t WHERE rowid = ?", [i])
# Phase 3: Insert new vectors with higher rowids
for i in range(50, 75):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(i), 0, 0, 0])],
)
# Phase 4: Delete some of the new ones
for i in range(60, 70):
db.execute("DELETE FROM t WHERE rowid = ?", [i])
# KNN query: should only find existing vectors
results = knn(db, [25.0, 0, 0, 0], 50)
rowids = {r[0] for r in results}
# Verify no deleted rowids appear
deleted = set(range(0, 50, 2)) | set(range(60, 70))
assert len(rowids & deleted) == 0
# Verify we get the right count (25 odd + 15 new - 10 deleted new = 30)
expected_alive = set(range(1, 50, 2)) | set(range(50, 60)) | set(range(70, 75))
assert rowids.issubset(expected_alive)

255
tests/test-ivf-quantization.py Normal file
View file

@ -0,0 +1,255 @@
import pytest
import sqlite3
from helpers import _f32, exec
@pytest.fixture()
def db():
db = sqlite3.connect(":memory:")
db.row_factory = sqlite3.Row
db.enable_load_extension(True)
db.load_extension("dist/vec0")
db.enable_load_extension(False)
return db
# ============================================================================
# Parser tests — quantizer and oversample options
# ============================================================================
def test_ivf_quantizer_binary(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[768] indexed by ivf(nlist=64, quantizer=binary, oversample=10))"
)
tables = [
r[0]
for r in db.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY 1"
).fetchall()
]
assert "t_ivf_centroids00" in tables
assert "t_ivf_cells00" in tables
def test_ivf_quantizer_int8(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[768] indexed by ivf(nlist=64, quantizer=int8))"
)
tables = [
r[0]
for r in db.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY 1"
).fetchall()
]
assert "t_ivf_centroids00" in tables
def test_ivf_quantizer_none_explicit(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[768] indexed by ivf(quantizer=none))"
)
# Should work — same as no quantizer
tables = [
r[0]
for r in db.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY 1"
).fetchall()
]
assert "t_ivf_centroids00" in tables
def test_ivf_quantizer_all_params(db):
"""All params together: nlist, nprobe, quantizer, oversample."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[768] distance_metric=cosine "
"indexed by ivf(nlist=128, nprobe=16, quantizer=int8, oversample=4))"
)
tables = [
r[0]
for r in db.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY 1"
).fetchall()
]
assert "t_ivf_centroids00" in tables
def test_ivf_error_oversample_without_quantizer(db):
"""oversample > 1 without quantizer should error."""
result = exec(
db,
"CREATE VIRTUAL TABLE t USING vec0("
"v float[768] indexed by ivf(oversample=10))",
)
assert "error" in result
def test_ivf_error_unknown_quantizer(db):
result = exec(
db,
"CREATE VIRTUAL TABLE t USING vec0("
"v float[768] indexed by ivf(quantizer=pq))",
)
assert "error" in result
def test_ivf_oversample_1_without_quantizer_ok(db):
"""oversample=1 (default) is fine without quantizer."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[768] indexed by ivf(nlist=64))"
)
# Should succeed — oversample defaults to 1
# ============================================================================
# Functional tests — insert, train, query with quantized IVF
# ============================================================================
def test_ivf_int8_insert_and_query(db):
"""int8 quantized IVF: insert, train, query."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[4] indexed by ivf(nlist=2, quantizer=int8, oversample=4))"
)
for i in range(20):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
# Should be able to query
rows = db.execute(
"SELECT rowid, distance FROM t WHERE v MATCH ? AND k = 5",
[_f32([10.0, 0, 0, 0])],
).fetchall()
assert len(rows) == 5
# Top result should be close to 10
assert rows[0][0] in range(8, 13)
# Full vectors should be in _ivf_vectors table
fv_count = db.execute("SELECT count(*) FROM t_ivf_vectors00").fetchone()[0]
assert fv_count == 20
def test_ivf_binary_insert_and_query(db):
"""Binary quantized IVF: insert, train, query."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[8] indexed by ivf(nlist=2, quantizer=binary, oversample=4))"
)
for i in range(20):
# Vectors with varying sign patterns
v = [(i * 0.1 - 1.0) + j * 0.3 for j in range(8)]
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32(v)]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
rows = db.execute(
"SELECT rowid FROM t WHERE v MATCH ? AND k = 5",
[_f32([0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5])],
).fetchall()
assert len(rows) == 5
# Full vectors stored
fv_count = db.execute("SELECT count(*) FROM t_ivf_vectors00").fetchone()[0]
assert fv_count == 20
def test_ivf_int8_cell_sizes_smaller(db):
"""Cell blobs should be smaller with int8 quantization."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[768] indexed by ivf(nlist=2, quantizer=int8, oversample=1))"
)
for i in range(10):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(x) / 768 for x in range(768)])],
)
# Cell vectors blob: 10 vectors at int8 = 10 * 768 = 7680 bytes
# vs float32 = 10 * 768 * 4 = 30720 bytes
# But cells have capacity 64, so blob = 64 * 768 = 49152 (int8) vs 64*768*4=196608 (float32)
blob_size = db.execute(
"SELECT length(vectors) FROM t_ivf_cells00 LIMIT 1"
).fetchone()[0]
# int8: 64 slots * 768 bytes = 49152
assert blob_size == 64 * 768
def test_ivf_binary_cell_sizes_smaller(db):
"""Cell blobs should be much smaller with binary quantization."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[768] indexed by ivf(nlist=2, quantizer=binary, oversample=1))"
)
for i in range(10):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([float(x) / 768 for x in range(768)])],
)
blob_size = db.execute(
"SELECT length(vectors) FROM t_ivf_cells00 LIMIT 1"
).fetchone()[0]
# binary: 64 slots * 768/8 bytes = 6144
assert blob_size == 64 * (768 // 8)
def test_ivf_int8_oversample_improves_recall(db):
"""Oversample re-ranking should improve recall over oversample=1."""
# Create two tables: one with oversample=1, one with oversample=10
db.execute(
"CREATE VIRTUAL TABLE t1 USING vec0("
"v float[4] indexed by ivf(nlist=4, quantizer=int8, oversample=1))"
)
db.execute(
"CREATE VIRTUAL TABLE t2 USING vec0("
"v float[4] indexed by ivf(nlist=4, quantizer=int8, oversample=10))"
)
for i in range(100):
v = _f32([i * 0.1, (i * 0.1) % 3, (i * 0.3) % 5, i * 0.01])
db.execute("INSERT INTO t1(rowid, v) VALUES (?, ?)", [i, v])
db.execute("INSERT INTO t2(rowid, v) VALUES (?, ?)", [i, v])
db.execute("INSERT INTO t1(rowid) VALUES ('compute-centroids')")
db.execute("INSERT INTO t2(rowid) VALUES ('compute-centroids')")
db.execute("INSERT INTO t1(rowid) VALUES ('nprobe=4')")
db.execute("INSERT INTO t2(rowid) VALUES ('nprobe=4')")
query = _f32([5.0, 1.5, 2.5, 0.5])
r1 = db.execute("SELECT rowid FROM t1 WHERE v MATCH ? AND k=10", [query]).fetchall()
r2 = db.execute("SELECT rowid FROM t2 WHERE v MATCH ? AND k=10", [query]).fetchall()
# Both should return 10 results
assert len(r1) == 10
assert len(r2) == 10
# oversample=10 should have at least as good recall (same or better ordering)
def test_ivf_quantized_delete(db):
"""Delete should remove from both cells and _ivf_vectors."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0("
"v float[4] indexed by ivf(nlist=2, quantizer=int8, oversample=1))"
)
for i in range(10):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
assert db.execute("SELECT count(*) FROM t_ivf_vectors00").fetchone()[0] == 10
db.execute("DELETE FROM t WHERE rowid = 5")
# _ivf_vectors should have 9 rows
assert db.execute("SELECT count(*) FROM t_ivf_vectors00").fetchone()[0] == 9

426
tests/test-ivf.py Normal file
View file

@ -0,0 +1,426 @@
import pytest
import sqlite3
import struct
import math
from helpers import _f32, exec
@pytest.fixture()
def db():
db = sqlite3.connect(":memory:")
db.row_factory = sqlite3.Row
db.enable_load_extension(True)
db.load_extension("dist/vec0")
db.enable_load_extension(False)
return db
def ivf_total_vectors(db, table="t", col=0):
"""Count total vectors across all IVF cells."""
return db.execute(
f"SELECT COALESCE(SUM(n_vectors), 0) FROM {table}_ivf_cells{col:02d}"
).fetchone()[0]
def ivf_unassigned_count(db, table="t", col=0):
"""Count vectors in unassigned cells (centroid_id=-1)."""
return db.execute(
f"SELECT COALESCE(SUM(n_vectors), 0) FROM {table}_ivf_cells{col:02d} WHERE centroid_id = -1"
).fetchone()[0]
def ivf_assigned_count(db, table="t", col=0):
"""Count vectors in trained cells (centroid_id >= 0)."""
return db.execute(
f"SELECT COALESCE(SUM(n_vectors), 0) FROM {table}_ivf_cells{col:02d} WHERE centroid_id >= 0"
).fetchone()[0]
# ============================================================================
# Parser tests
# ============================================================================
def test_ivf_create_defaults(db):
"""ivf() with no args uses defaults."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf())"
)
tables = [
r[0]
for r in db.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY 1"
).fetchall()
]
assert "t_ivf_centroids00" in tables
assert "t_ivf_cells00" in tables
assert "t_ivf_rowid_map00" in tables
def test_ivf_create_custom_params(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=64, nprobe=8))"
)
tables = [
r[0]
for r in db.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY 1"
).fetchall()
]
assert "t_ivf_centroids00" in tables
assert "t_ivf_cells00" in tables
def test_ivf_create_with_distance_metric(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] distance_metric=cosine indexed by ivf(nlist=16))"
)
tables = [
r[0]
for r in db.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY 1"
).fetchall()
]
assert "t_ivf_centroids00" in tables
def test_ivf_create_error_unknown_key(db):
result = exec(
db,
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(bogus=1))",
)
assert "error" in result
def test_ivf_create_error_nprobe_gt_nlist(db):
result = exec(
db,
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4, nprobe=10))",
)
assert "error" in result
# ============================================================================
# Shadow table tests
# ============================================================================
def test_ivf_shadow_tables_created(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=8))"
)
trained = db.execute(
"SELECT value FROM t_info WHERE key = 'ivf_trained_0'"
).fetchone()[0]
assert str(trained) == "0"
# No cells yet (created lazily on first insert)
count = db.execute(
"SELECT count(*) FROM t_ivf_cells00"
).fetchone()[0]
assert count == 0
def test_ivf_drop_cleans_up(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4))"
)
db.execute("DROP TABLE t")
tables = [
r[0]
for r in db.execute(
"SELECT name FROM sqlite_master WHERE type='table'"
).fetchall()
]
assert not any("ivf" in t for t in tables)
# ============================================================================
# Insert tests (flat mode)
# ============================================================================
def test_ivf_insert_flat_mode(db):
"""Before training, vectors go to unassigned cell."""
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4))"
)
db.execute("INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([1, 2, 3, 4])])
db.execute("INSERT INTO t(rowid, v) VALUES (2, ?)", [_f32([5, 6, 7, 8])])
assert ivf_unassigned_count(db) == 2
assert ivf_assigned_count(db) == 0
# rowid_map should have 2 entries
assert db.execute("SELECT count(*) FROM t_ivf_rowid_map00").fetchone()[0] == 2
def test_ivf_delete_flat_mode(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4))"
)
db.execute("INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([1, 2, 3, 4])])
db.execute("INSERT INTO t(rowid, v) VALUES (2, ?)", [_f32([5, 6, 7, 8])])
db.execute("DELETE FROM t WHERE rowid = 1")
assert ivf_unassigned_count(db) == 1
assert db.execute("SELECT count(*) FROM t_ivf_rowid_map00").fetchone()[0] == 1
# ============================================================================
# KNN flat mode tests
# ============================================================================
def test_ivf_knn_flat_mode(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4))"
)
db.execute("INSERT INTO t(rowid, v) VALUES (1, ?)", [_f32([1, 0, 0, 0])])
db.execute("INSERT INTO t(rowid, v) VALUES (2, ?)", [_f32([2, 0, 0, 0])])
db.execute("INSERT INTO t(rowid, v) VALUES (3, ?)", [_f32([9, 0, 0, 0])])
rows = db.execute(
"SELECT rowid, distance FROM t WHERE v MATCH ? AND k = 2",
[_f32([1.5, 0, 0, 0])],
).fetchall()
assert len(rows) == 2
rowids = {r[0] for r in rows}
assert rowids == {1, 2}
def test_ivf_knn_flat_empty(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4))"
)
rows = db.execute(
"SELECT rowid FROM t WHERE v MATCH ? AND k = 5",
[_f32([1, 0, 0, 0])],
).fetchall()
assert len(rows) == 0
# ============================================================================
# compute-centroids tests
# ============================================================================
def test_compute_centroids(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4))"
)
for i in range(40):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)",
[i, _f32([i % 10, i // 10, 0, 0])],
)
assert ivf_unassigned_count(db) == 40
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
# After training: unassigned cell should be gone (or empty), vectors in trained cells
assert ivf_unassigned_count(db) == 0
assert ivf_assigned_count(db) == 40
assert db.execute("SELECT count(*) FROM t_ivf_centroids00").fetchone()[0] == 4
trained = db.execute(
"SELECT value FROM t_info WHERE key='ivf_trained_0'"
).fetchone()[0]
assert str(trained) == "1"
def test_compute_centroids_recompute(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2))"
)
for i in range(20):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
assert db.execute("SELECT count(*) FROM t_ivf_centroids00").fetchone()[0] == 2
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
assert db.execute("SELECT count(*) FROM t_ivf_centroids00").fetchone()[0] == 2
assert ivf_assigned_count(db) == 20
# ============================================================================
# Insert after training (assigned mode)
# ============================================================================
def test_ivf_insert_after_training(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2))"
)
for i in range(20):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
db.execute(
"INSERT INTO t(rowid, v) VALUES (100, ?)", [_f32([5, 0, 0, 0])]
)
# Should be in a trained cell, not unassigned
row = db.execute(
"SELECT m.cell_id, c.centroid_id FROM t_ivf_rowid_map00 m "
"JOIN t_ivf_cells00 c ON c.rowid = m.cell_id "
"WHERE m.rowid = 100"
).fetchone()
assert row is not None
assert row[1] >= 0 # centroid_id >= 0 means trained cell
# ============================================================================
# KNN after training (IVF probe mode)
# ============================================================================
def test_ivf_knn_after_training(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=4, nprobe=4))"
)
for i in range(100):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
rows = db.execute(
"SELECT rowid, distance FROM t WHERE v MATCH ? AND k = 5",
[_f32([50.0, 0, 0, 0])],
).fetchall()
assert len(rows) == 5
assert rows[0][0] == 50
assert rows[0][1] < 0.01
def test_ivf_knn_k_larger_than_n(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2, nprobe=2))"
)
for i in range(5):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
rows = db.execute(
"SELECT rowid FROM t WHERE v MATCH ? AND k = 100",
[_f32([0, 0, 0, 0])],
).fetchall()
assert len(rows) == 5
# ============================================================================
# Manual centroid import (set-centroid, assign-vectors)
# ============================================================================
def test_set_centroid_and_assign(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=0))"
)
for i in range(20):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute(
"INSERT INTO t(rowid, v) VALUES ('set-centroid:0', ?)",
[_f32([5, 0, 0, 0])],
)
db.execute(
"INSERT INTO t(rowid, v) VALUES ('set-centroid:1', ?)",
[_f32([15, 0, 0, 0])],
)
assert db.execute("SELECT count(*) FROM t_ivf_centroids00").fetchone()[0] == 2
db.execute("INSERT INTO t(rowid) VALUES ('assign-vectors')")
assert ivf_unassigned_count(db) == 0
assert ivf_assigned_count(db) == 20
# ============================================================================
# clear-centroids
# ============================================================================
def test_clear_centroids(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2))"
)
for i in range(20):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
assert db.execute("SELECT count(*) FROM t_ivf_centroids00").fetchone()[0] == 2
db.execute("INSERT INTO t(rowid) VALUES ('clear-centroids')")
assert db.execute("SELECT count(*) FROM t_ivf_centroids00").fetchone()[0] == 0
assert ivf_unassigned_count(db) == 20
trained = db.execute(
"SELECT value FROM t_info WHERE key='ivf_trained_0'"
).fetchone()[0]
assert str(trained) == "0"
# ============================================================================
# Delete after training
# ============================================================================
def test_ivf_delete_after_training(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=2))"
)
for i in range(10):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
assert ivf_assigned_count(db) == 10
db.execute("DELETE FROM t WHERE rowid = 5")
assert ivf_assigned_count(db) == 9
assert db.execute("SELECT count(*) FROM t_ivf_rowid_map00").fetchone()[0] == 9
# ============================================================================
# Recall test
# ============================================================================
def test_ivf_recall_nprobe_equals_nlist(db):
db.execute(
"CREATE VIRTUAL TABLE t USING vec0(v float[4] indexed by ivf(nlist=8, nprobe=8))"
)
for i in range(100):
db.execute(
"INSERT INTO t(rowid, v) VALUES (?, ?)", [i, _f32([i, 0, 0, 0])]
)
db.execute("INSERT INTO t(rowid) VALUES ('compute-centroids')")
rows = db.execute(
"SELECT rowid FROM t WHERE v MATCH ? AND k = 10",
[_f32([50.0, 0, 0, 0])],
).fetchall()
rowids = {r[0] for r in rows}
# 45 and 55 are equidistant from 50, so either may appear in top 10
assert 50 in rowids
assert len(rowids) == 10
assert all(45 <= r <= 55 for r in rowids)

419
tests/test-unit.c
View file

@ -577,6 +577,182 @@ void test_vec0_parse_vector_column() {
assert(rc == SQLITE_ERROR);
}
#if SQLITE_VEC_ENABLE_IVF
// IVF: indexed by ivf() — defaults
{
const char *input = "v float[4] indexed by ivf()";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.index_type == VEC0_INDEX_TYPE_IVF);
assert(col.dimensions == 4);
assert(col.index_type == VEC0_INDEX_TYPE_IVF);
assert(col.ivf.nlist == 128); // default
assert(col.ivf.nprobe == 10); // default
sqlite3_free(col.name);
}
// IVF: indexed by ivf(nlist=8) — nprobe auto-clamped to 8
{
const char *input = "v float[4] indexed by ivf(nlist=8)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.index_type == VEC0_INDEX_TYPE_IVF);
assert(col.index_type == VEC0_INDEX_TYPE_IVF);
assert(col.ivf.nlist == 8);
assert(col.ivf.nprobe == 8); // clamped from default 10
sqlite3_free(col.name);
}
// IVF: indexed by ivf(nlist=64, nprobe=8)
{
const char *input = "v float[4] indexed by ivf(nlist=64, nprobe=8)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.index_type == VEC0_INDEX_TYPE_IVF);
assert(col.ivf.nlist == 64);
assert(col.ivf.nprobe == 8);
sqlite3_free(col.name);
}
// IVF: with distance_metric before indexed by
{
const char *input = "v float[4] distance_metric=cosine indexed by ivf(nlist=16)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.index_type == VEC0_INDEX_TYPE_IVF);
assert(col.distance_metric == VEC0_DISTANCE_METRIC_COSINE);
assert(col.index_type == VEC0_INDEX_TYPE_IVF);
assert(col.ivf.nlist == 16);
sqlite3_free(col.name);
}
// IVF: nlist=0 (deferred)
{
const char *input = "v float[4] indexed by ivf(nlist=0)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.nlist == 0);
sqlite3_free(col.name);
}
// IVF error: nprobe > nlist
{
const char *input = "v float[4] indexed by ivf(nlist=4, nprobe=10)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_ERROR);
}
// IVF error: unknown key
{
const char *input = "v float[4] indexed by ivf(bogus=1)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_ERROR);
}
// IVF error: unknown index type (hnsw not supported)
{
const char *input = "v float[4] indexed by hnsw()";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_ERROR);
}
// Not IVF: no ivf config
{
const char *input = "v float[4]";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.index_type == VEC0_INDEX_TYPE_FLAT);
sqlite3_free(col.name);
}
// IVF: quantizer=binary
{
const char *input = "v float[768] indexed by ivf(nlist=128, quantizer=binary)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.index_type == VEC0_INDEX_TYPE_IVF);
assert(col.ivf.nlist == 128);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_BINARY);
assert(col.ivf.oversample == 1);
sqlite3_free(col.name);
}
// IVF: quantizer=int8
{
const char *input = "v float[768] indexed by ivf(nlist=64, quantizer=int8)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
sqlite3_free(col.name);
}
// IVF: quantizer=none (explicit)
{
const char *input = "v float[768] indexed by ivf(quantizer=none)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_NONE);
sqlite3_free(col.name);
}
// IVF: oversample=10 with quantizer
{
const char *input = "v float[768] indexed by ivf(nlist=128, quantizer=binary, oversample=10)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_BINARY);
assert(col.ivf.oversample == 10);
assert(col.ivf.nlist == 128);
sqlite3_free(col.name);
}
// IVF: all params
{
const char *input = "v float[768] distance_metric=cosine indexed by ivf(nlist=256, nprobe=16, quantizer=int8, oversample=4)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.distance_metric == VEC0_DISTANCE_METRIC_COSINE);
assert(col.ivf.nlist == 256);
assert(col.ivf.nprobe == 16);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
assert(col.ivf.oversample == 4);
sqlite3_free(col.name);
}
// IVF error: oversample > 1 without quantizer
{
const char *input = "v float[768] indexed by ivf(oversample=10)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_ERROR);
}
// IVF error: unknown quantizer value
{
const char *input = "v float[768] indexed by ivf(quantizer=pq)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_ERROR);
}
// IVF: quantizer with defaults (nlist=128 default, nprobe=10 default)
{
const char *input = "v float[768] indexed by ivf(quantizer=binary, oversample=5)";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.nlist == 128);
assert(col.ivf.nprobe == 10);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_BINARY);
assert(col.ivf.oversample == 5);
sqlite3_free(col.name);
}
#else
// When IVF is disabled, parsing "ivf" should fail
{
const char *input = "v float[4] indexed by ivf()";
rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
assert(rc == SQLITE_ERROR);
}
#endif /* SQLITE_VEC_ENABLE_IVF */
printf(" All vec0_parse_vector_column tests passed.\n");
}
@ -821,6 +997,38 @@ void test_rescore_quantize_float_to_int8() {
float src[8] = {5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f};
_test_rescore_quantize_float_to_int8(src, dst, 8);
for (int i = 0; i < 8; i++) {
#if SQLITE_VEC_ENABLE_IVF
void test_ivf_quantize_int8() {
printf("Starting %s...\n", __func__);
// Basic values in [-1, 1] range
{
float src[] = {0.0f, 1.0f, -1.0f, 0.5f};
int8_t dst[4];
ivf_quantize_int8(src, dst, 4);
assert(dst[0] == 0);
assert(dst[1] == 127);
assert(dst[2] == -127);
assert(dst[3] == 63); // 0.5 * 127 = 63.5, truncated to 63
}
// Clamping: values beyond [-1, 1]
{
float src[] = {2.0f, -3.0f, 100.0f, -0.01f};
int8_t dst[4];
ivf_quantize_int8(src, dst, 4);
assert(dst[0] == 127); // clamped to 1.0
assert(dst[1] == -127); // clamped to -1.0
assert(dst[2] == 127); // clamped to 1.0
assert(dst[3] == (int8_t)(-0.01f * 127.0f));
}
// Zero vector
{
float src[] = {0.0f, 0.0f, 0.0f, 0.0f};
int8_t dst[4];
ivf_quantize_int8(src, dst, 4);
for (int i = 0; i < 4; i++) {
assert(dst[i] == 0);
}
}
@ -882,6 +1090,103 @@ void test_rescore_quantized_byte_size() {
}
void test_vec0_parse_vector_column_rescore() {
// Negative zero
{
float src[] = {-0.0f};
int8_t dst[1];
ivf_quantize_int8(src, dst, 1);
assert(dst[0] == 0);
}
// Single element
{
float src[] = {0.75f};
int8_t dst[1];
ivf_quantize_int8(src, dst, 1);
assert(dst[0] == (int8_t)(0.75f * 127.0f));
}
// Boundary: exactly 1.0 and -1.0
{
float src[] = {1.0f, -1.0f};
int8_t dst[2];
ivf_quantize_int8(src, dst, 2);
assert(dst[0] == 127);
assert(dst[1] == -127);
}
printf(" All ivf_quantize_int8 tests passed.\n");
}
void test_ivf_quantize_binary() {
printf("Starting %s...\n", __func__);
// Basic sign-bit quantization: positive -> 1, negative/zero -> 0
{
float src[] = {1.0f, -1.0f, 0.5f, -0.5f, 0.0f, 0.1f, -0.1f, 2.0f};
uint8_t dst[1];
ivf_quantize_binary(src, dst, 8);
// bit 0: 1.0 > 0 -> 1 (LSB)
// bit 1: -1.0 -> 0
// bit 2: 0.5 > 0 -> 1
// bit 3: -0.5 -> 0
// bit 4: 0.0 -> 0 (not > 0)
// bit 5: 0.1 > 0 -> 1
// bit 6: -0.1 -> 0
// bit 7: 2.0 > 0 -> 1
// Expected: bits 0,2,5,7 = 0b10100101 = 0xA5
assert(dst[0] == 0xA5);
}
// All positive
{
float src[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
uint8_t dst[1];
ivf_quantize_binary(src, dst, 8);
assert(dst[0] == 0xFF);
}
// All negative
{
float src[] = {-1.0f, -2.0f, -3.0f, -4.0f, -5.0f, -6.0f, -7.0f, -8.0f};
uint8_t dst[1];
ivf_quantize_binary(src, dst, 8);
assert(dst[0] == 0x00);
}
// All zero (zero is NOT > 0, so all bits should be 0)
{
float src[] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
uint8_t dst[1];
ivf_quantize_binary(src, dst, 8);
assert(dst[0] == 0x00);
}
// Multi-byte: 16 dimensions -> 2 bytes
{
float src[16];
for (int i = 0; i < 16; i++) src[i] = (i % 2 == 0) ? 1.0f : -1.0f;
uint8_t dst[2];
ivf_quantize_binary(src, dst, 16);
// Even indices are positive: bits 0,2,4,6 in each byte
// byte 0: bits 0,2,4,6 = 0b01010101 = 0x55
// byte 1: same pattern = 0x55
assert(dst[0] == 0x55);
assert(dst[1] == 0x55);
}
// Single byte, only first bit set
{
float src[] = {0.1f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f};
uint8_t dst[1];
ivf_quantize_binary(src, dst, 8);
assert(dst[0] == 0x01);
}
printf(" All ivf_quantize_binary tests passed.\n");
}
void test_ivf_config_parsing() {
printf("Starting %s...\n", __func__);
struct VectorColumnDefinition col;
int rc;
@ -955,6 +1260,116 @@ void test_vec0_parse_vector_column_rescore() {
}
#endif /* SQLITE_VEC_ENABLE_RESCORE */
// Default IVF config
{
const char *s = "v float[4] indexed by ivf()";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_OK);
assert(col.index_type == VEC0_INDEX_TYPE_IVF);
assert(col.ivf.nlist == 128); // default
assert(col.ivf.nprobe == 10); // default
assert(col.ivf.quantizer == 0); // VEC0_IVF_QUANTIZER_NONE
sqlite3_free(col.name);
}
// Custom nlist and nprobe
{
const char *s = "v float[4] indexed by ivf(nlist=64, nprobe=8)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.nlist == 64);
assert(col.ivf.nprobe == 8);
sqlite3_free(col.name);
}
// nlist=0 (deferred)
{
const char *s = "v float[4] indexed by ivf(nlist=0)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.nlist == 0);
sqlite3_free(col.name);
}
// Quantizer options
{
const char *s = "v float[8] indexed by ivf(quantizer=int8)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
sqlite3_free(col.name);
}
{
const char *s = "v float[8] indexed by ivf(quantizer=binary)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_BINARY);
sqlite3_free(col.name);
}
// nprobe > nlist (explicit) should fail
{
const char *s = "v float[4] indexed by ivf(nlist=4, nprobe=10)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_ERROR);
}
// Unknown key
{
const char *s = "v float[4] indexed by ivf(bogus=1)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_ERROR);
}
// nlist > max (65536) should fail
{
const char *s = "v float[4] indexed by ivf(nlist=65537)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_ERROR);
}
// nlist at max boundary (65536) should succeed
{
const char *s = "v float[4] indexed by ivf(nlist=65536)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.nlist == 65536);
sqlite3_free(col.name);
}
// oversample > 1 without quantization should fail
{
const char *s = "v float[4] indexed by ivf(oversample=4)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_ERROR);
}
// oversample with quantizer should succeed
{
const char *s = "v float[8] indexed by ivf(quantizer=int8, oversample=4)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.oversample == 4);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
sqlite3_free(col.name);
}
// All options combined
{
const char *s = "v float[8] indexed by ivf(nlist=32, nprobe=4, quantizer=int8, oversample=2)";
rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
assert(rc == SQLITE_OK);
assert(col.ivf.nlist == 32);
assert(col.ivf.nprobe == 4);
assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
assert(col.ivf.oversample == 2);
sqlite3_free(col.name);
}
printf(" All ivf_config_parsing tests passed.\n");
}
#endif /* SQLITE_VEC_ENABLE_IVF */
int main() {
printf("Starting unit tests...\n");
@ -982,6 +1397,10 @@ int main() {
test_rescore_quantize_float_to_int8();
test_rescore_quantized_byte_size();
test_vec0_parse_vector_column_rescore();
#if SQLITE_VEC_ENABLE_IVF
test_ivf_quantize_int8();
test_ivf_quantize_binary();
test_ivf_config_parsing();
#endif
printf("All unit tests passed.\n");
}