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
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;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue