mirror of
https://github.com/asg017/sqlite-vec.git
synced 2026-04-25 00:36:56 +02:00
Add IVF index for vec0 virtual table
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:
parent
43982c144b
commit
3358e127f6
22 changed files with 5237 additions and 28 deletions
|
|
@ -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
192
tests/fuzz/ivf-cell-overflow.c
Normal 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
36
tests/fuzz/ivf-create.c
Normal 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
16
tests/fuzz/ivf-create.dict
Normal 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
180
tests/fuzz/ivf-kmeans.c
Normal 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
199
tests/fuzz/ivf-knn-deep.c
Normal 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
121
tests/fuzz/ivf-operations.c
Normal 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
129
tests/fuzz/ivf-quantize.c
Normal 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
182
tests/fuzz/ivf-rescore.c
Normal 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
228
tests/fuzz/ivf-shadow-corrupt.c
Normal 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;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue