mirror of
https://github.com/asg017/sqlite-vec.git
synced 2026-04-25 00:36:56 +02:00
b00865429b
DiskANN's delete repair only fixes forward edges (nodes the deleted node pointed to). Stale reverse edges can cause deleted rowids to appear in search results. Fix: track a 'confirmed' flag on each search candidate, set when the full-precision vector is successfully read during re-ranking. Only confirmed candidates are included in output. Zero additional SQL queries — piggybacks on the existing re-rank vector read. Also adds delete hardening tests: - Rescore: interleaved delete+KNN, rowid_in after deletes, full delete+reinsert cycle - DiskANN: delete+reinsert cycles with KNN verification, interleaved delete+KNN Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1793 lines
60 KiB
C
1793 lines
60 KiB
C
// DiskANN algorithm implementation
|
|
// This file is #include'd into sqlite-vec.c — not compiled separately.
|
|
|
|
// ============================================================
|
|
// DiskANN node blob encode/decode functions
|
|
// ============================================================
|
|
|
|
/** Compute size of validity bitmap in bytes. */
|
|
int diskann_validity_byte_size(int n_neighbors) {
|
|
return n_neighbors / CHAR_BIT;
|
|
}
|
|
|
|
/** Compute size of neighbor IDs blob in bytes. */
|
|
size_t diskann_neighbor_ids_byte_size(int n_neighbors) {
|
|
return (size_t)n_neighbors * sizeof(i64);
|
|
}
|
|
|
|
/** Compute size of quantized vectors blob in bytes. */
|
|
size_t diskann_neighbor_qvecs_byte_size(
|
|
int n_neighbors, enum Vec0DiskannQuantizerType quantizer_type,
|
|
size_t dimensions) {
|
|
return (size_t)n_neighbors *
|
|
diskann_quantized_vector_byte_size(quantizer_type, dimensions);
|
|
}
|
|
|
|
/**
|
|
* Create empty blobs for a new DiskANN node (all neighbors invalid).
|
|
* Caller must free the returned pointers with sqlite3_free().
|
|
*/
|
|
int diskann_node_init(
|
|
int n_neighbors, enum Vec0DiskannQuantizerType quantizer_type,
|
|
size_t dimensions,
|
|
u8 **outValidity, int *outValiditySize,
|
|
u8 **outNeighborIds, int *outNeighborIdsSize,
|
|
u8 **outNeighborQvecs, int *outNeighborQvecsSize) {
|
|
|
|
int validitySize = diskann_validity_byte_size(n_neighbors);
|
|
size_t idsSize = diskann_neighbor_ids_byte_size(n_neighbors);
|
|
size_t qvecsSize = diskann_neighbor_qvecs_byte_size(
|
|
n_neighbors, quantizer_type, dimensions);
|
|
|
|
u8 *validity = sqlite3_malloc(validitySize);
|
|
u8 *ids = sqlite3_malloc(idsSize);
|
|
u8 *qvecs = sqlite3_malloc(qvecsSize);
|
|
|
|
if (!validity || !ids || !qvecs) {
|
|
sqlite3_free(validity);
|
|
sqlite3_free(ids);
|
|
sqlite3_free(qvecs);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
memset(validity, 0, validitySize);
|
|
memset(ids, 0, idsSize);
|
|
memset(qvecs, 0, qvecsSize);
|
|
|
|
*outValidity = validity; *outValiditySize = validitySize;
|
|
*outNeighborIds = ids; *outNeighborIdsSize = (int)idsSize;
|
|
*outNeighborQvecs = qvecs; *outNeighborQvecsSize = (int)qvecsSize;
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
/** Check if neighbor slot i is valid. */
|
|
int diskann_validity_get(const u8 *validity, int i) {
|
|
return (validity[i / CHAR_BIT] >> (i % CHAR_BIT)) & 1;
|
|
}
|
|
|
|
/** Set neighbor slot i as valid (1) or invalid (0). */
|
|
void diskann_validity_set(u8 *validity, int i, int value) {
|
|
if (value) {
|
|
validity[i / CHAR_BIT] |= (1 << (i % CHAR_BIT));
|
|
} else {
|
|
validity[i / CHAR_BIT] &= ~(1 << (i % CHAR_BIT));
|
|
}
|
|
}
|
|
|
|
/** Count the number of valid neighbors. */
|
|
int diskann_validity_count(const u8 *validity, int n_neighbors) {
|
|
int count = 0;
|
|
for (int i = 0; i < n_neighbors; i++) {
|
|
count += diskann_validity_get(validity, i);
|
|
}
|
|
return count;
|
|
}
|
|
|
|
/** Get the rowid of the neighbor in slot i. */
|
|
i64 diskann_neighbor_id_get(const u8 *neighbor_ids, int i) {
|
|
i64 result;
|
|
memcpy(&result, neighbor_ids + i * sizeof(i64), sizeof(i64));
|
|
return result;
|
|
}
|
|
|
|
/** Set the rowid of the neighbor in slot i. */
|
|
void diskann_neighbor_id_set(u8 *neighbor_ids, int i, i64 rowid) {
|
|
memcpy(neighbor_ids + i * sizeof(i64), &rowid, sizeof(i64));
|
|
}
|
|
|
|
/** Get a pointer to the quantized vector in slot i (read-only). */
|
|
const u8 *diskann_neighbor_qvec_get(
|
|
const u8 *qvecs, int i,
|
|
enum Vec0DiskannQuantizerType quantizer_type, size_t dimensions) {
|
|
size_t qvec_size = diskann_quantized_vector_byte_size(quantizer_type, dimensions);
|
|
return qvecs + (size_t)i * qvec_size;
|
|
}
|
|
|
|
/** Copy a quantized vector into slot i. */
|
|
void diskann_neighbor_qvec_set(
|
|
u8 *qvecs, int i, const u8 *src_qvec,
|
|
enum Vec0DiskannQuantizerType quantizer_type, size_t dimensions) {
|
|
size_t qvec_size = diskann_quantized_vector_byte_size(quantizer_type, dimensions);
|
|
memcpy(qvecs + (size_t)i * qvec_size, src_qvec, qvec_size);
|
|
}
|
|
|
|
/**
|
|
* Set neighbor slot i with a rowid and quantized vector, and mark as valid.
|
|
*/
|
|
void diskann_node_set_neighbor(
|
|
u8 *validity, u8 *neighbor_ids, u8 *qvecs, int i,
|
|
i64 neighbor_rowid, const u8 *neighbor_qvec,
|
|
enum Vec0DiskannQuantizerType quantizer_type, size_t dimensions) {
|
|
diskann_validity_set(validity, i, 1);
|
|
diskann_neighbor_id_set(neighbor_ids, i, neighbor_rowid);
|
|
diskann_neighbor_qvec_set(qvecs, i, neighbor_qvec, quantizer_type, dimensions);
|
|
}
|
|
|
|
/**
|
|
* Clear neighbor slot i (mark invalid, zero out data).
|
|
*/
|
|
void diskann_node_clear_neighbor(
|
|
u8 *validity, u8 *neighbor_ids, u8 *qvecs, int i,
|
|
enum Vec0DiskannQuantizerType quantizer_type, size_t dimensions) {
|
|
diskann_validity_set(validity, i, 0);
|
|
diskann_neighbor_id_set(neighbor_ids, i, 0);
|
|
size_t qvec_size = diskann_quantized_vector_byte_size(quantizer_type, dimensions);
|
|
memset(qvecs + (size_t)i * qvec_size, 0, qvec_size);
|
|
}
|
|
|
|
/**
|
|
* Quantize a full-precision float32 vector into the target quantizer format.
|
|
* Output buffer must be pre-allocated with diskann_quantized_vector_byte_size() bytes.
|
|
*/
|
|
int diskann_quantize_vector(
|
|
const f32 *src, size_t dimensions,
|
|
enum Vec0DiskannQuantizerType quantizer_type,
|
|
u8 *out) {
|
|
|
|
switch (quantizer_type) {
|
|
case VEC0_DISKANN_QUANTIZER_BINARY: {
|
|
memset(out, 0, dimensions / CHAR_BIT);
|
|
for (size_t i = 0; i < dimensions; i++) {
|
|
if (src[i] > 0.0f) {
|
|
out[i / CHAR_BIT] |= (1 << (i % CHAR_BIT));
|
|
}
|
|
}
|
|
return SQLITE_OK;
|
|
}
|
|
case VEC0_DISKANN_QUANTIZER_INT8: {
|
|
f32 step = (1.0f - (-1.0f)) / 255.0f;
|
|
for (size_t i = 0; i < dimensions; i++) {
|
|
((i8 *)out)[i] = (i8)(((src[i] - (-1.0f)) / step) - 128.0f);
|
|
}
|
|
return SQLITE_OK;
|
|
}
|
|
}
|
|
return SQLITE_ERROR;
|
|
}
|
|
|
|
/**
|
|
* Compute approximate distance between a full-precision query vector and a
|
|
* quantized neighbor vector. Used during graph traversal.
|
|
*/
|
|
/**
|
|
* Compute distance between a pre-quantized query and a quantized neighbor.
|
|
* The caller is responsible for quantizing the query vector once and passing
|
|
* the result here for each neighbor comparison.
|
|
*/
|
|
static f32 diskann_distance_quantized_precomputed(
|
|
const u8 *query_quantized, const u8 *quantized_neighbor,
|
|
size_t dimensions,
|
|
enum Vec0DiskannQuantizerType quantizer_type,
|
|
enum Vec0DistanceMetrics distance_metric) {
|
|
|
|
switch (quantizer_type) {
|
|
case VEC0_DISKANN_QUANTIZER_BINARY:
|
|
return distance_hamming(query_quantized, quantized_neighbor, &dimensions);
|
|
case VEC0_DISKANN_QUANTIZER_INT8: {
|
|
switch (distance_metric) {
|
|
case VEC0_DISTANCE_METRIC_L2:
|
|
return distance_l2_sqr_int8(query_quantized, quantized_neighbor, &dimensions);
|
|
case VEC0_DISTANCE_METRIC_COSINE:
|
|
return distance_cosine_int8(query_quantized, quantized_neighbor, &dimensions);
|
|
case VEC0_DISTANCE_METRIC_L1:
|
|
return (f32)distance_l1_int8(query_quantized, quantized_neighbor, &dimensions);
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
return FLT_MAX;
|
|
}
|
|
|
|
/**
|
|
* Quantize a float query vector. Returns allocated buffer (caller must free).
|
|
*/
|
|
static u8 *diskann_quantize_query(
|
|
const f32 *query_vector, size_t dimensions,
|
|
enum Vec0DiskannQuantizerType quantizer_type) {
|
|
size_t qsize = diskann_quantized_vector_byte_size(quantizer_type, dimensions);
|
|
u8 *buf = sqlite3_malloc(qsize);
|
|
if (!buf) return NULL;
|
|
diskann_quantize_vector(query_vector, dimensions, quantizer_type, buf);
|
|
return buf;
|
|
}
|
|
|
|
/**
|
|
* Legacy wrapper: quantizes on-the-fly (used by callers that don't pre-quantize).
|
|
*/
|
|
f32 diskann_distance_quantized(
|
|
const void *query_vector, const u8 *quantized_neighbor,
|
|
size_t dimensions,
|
|
enum Vec0DiskannQuantizerType quantizer_type,
|
|
enum Vec0DistanceMetrics distance_metric) {
|
|
|
|
u8 *query_q = diskann_quantize_query((const f32 *)query_vector, dimensions, quantizer_type);
|
|
if (!query_q) return FLT_MAX;
|
|
f32 dist = diskann_distance_quantized_precomputed(
|
|
query_q, quantized_neighbor, dimensions, quantizer_type, distance_metric);
|
|
sqlite3_free(query_q);
|
|
return dist;
|
|
}
|
|
|
|
// ============================================================
|
|
// DiskANN medoid / entry point management
|
|
// ============================================================
|
|
|
|
/**
|
|
* Get the current medoid rowid for the given vector column's DiskANN index.
|
|
* Returns SQLITE_OK with *outMedoid set to the medoid rowid.
|
|
* If the graph is empty, returns SQLITE_OK with *outIsEmpty = 1.
|
|
*/
|
|
static int diskann_medoid_get(vec0_vtab *p, int vec_col_idx,
|
|
i64 *outMedoid, int *outIsEmpty) {
|
|
int rc;
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *key = sqlite3_mprintf("diskann_medoid_%02d", vec_col_idx);
|
|
char *zSql = sqlite3_mprintf(
|
|
"SELECT value FROM " VEC0_SHADOW_INFO_NAME " WHERE key = ?",
|
|
p->schemaName, p->tableName);
|
|
if (!key || !zSql) {
|
|
sqlite3_free(key);
|
|
sqlite3_free(zSql);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) {
|
|
sqlite3_free(key);
|
|
return rc;
|
|
}
|
|
|
|
sqlite3_bind_text(stmt, 1, key, -1, sqlite3_free);
|
|
rc = sqlite3_step(stmt);
|
|
if (rc == SQLITE_ROW) {
|
|
if (sqlite3_column_type(stmt, 0) == SQLITE_NULL) {
|
|
*outIsEmpty = 1;
|
|
} else {
|
|
*outIsEmpty = 0;
|
|
*outMedoid = sqlite3_column_int64(stmt, 0);
|
|
}
|
|
rc = SQLITE_OK;
|
|
} else {
|
|
rc = SQLITE_ERROR;
|
|
}
|
|
sqlite3_finalize(stmt);
|
|
return rc;
|
|
}
|
|
|
|
/**
|
|
* Set the medoid rowid for the given vector column's DiskANN index.
|
|
* Pass isEmpty = 1 to mark the graph as empty (NULL medoid).
|
|
*/
|
|
static int diskann_medoid_set(vec0_vtab *p, int vec_col_idx,
|
|
i64 medoidRowid, int isEmpty) {
|
|
int rc;
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *key = sqlite3_mprintf("diskann_medoid_%02d", vec_col_idx);
|
|
char *zSql = sqlite3_mprintf(
|
|
"UPDATE " VEC0_SHADOW_INFO_NAME " SET value = ?2 WHERE key = ?1",
|
|
p->schemaName, p->tableName);
|
|
if (!key || !zSql) {
|
|
sqlite3_free(key);
|
|
sqlite3_free(zSql);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) {
|
|
sqlite3_free(key);
|
|
return rc;
|
|
}
|
|
|
|
sqlite3_bind_text(stmt, 1, key, -1, sqlite3_free);
|
|
if (isEmpty) {
|
|
sqlite3_bind_null(stmt, 2);
|
|
} else {
|
|
sqlite3_bind_int64(stmt, 2, medoidRowid);
|
|
}
|
|
rc = sqlite3_step(stmt);
|
|
sqlite3_finalize(stmt);
|
|
return (rc == SQLITE_DONE) ? SQLITE_OK : SQLITE_ERROR;
|
|
}
|
|
|
|
|
|
/**
|
|
* Called when deleting a vector. If the deleted vector was the medoid,
|
|
* pick a new one (the first available vector, or set to empty if none remain).
|
|
*/
|
|
static int diskann_medoid_handle_delete(vec0_vtab *p, int vec_col_idx,
|
|
i64 deletedRowid) {
|
|
i64 currentMedoid;
|
|
int isEmpty;
|
|
int rc = diskann_medoid_get(p, vec_col_idx, ¤tMedoid, &isEmpty);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
if (!isEmpty && currentMedoid == deletedRowid) {
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *zSql = sqlite3_mprintf(
|
|
"SELECT rowid FROM " VEC0_SHADOW_VECTORS_N_NAME " WHERE rowid != ?1 LIMIT 1",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
|
|
rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
sqlite3_bind_int64(stmt, 1, deletedRowid);
|
|
rc = sqlite3_step(stmt);
|
|
if (rc == SQLITE_ROW) {
|
|
i64 newMedoid = sqlite3_column_int64(stmt, 0);
|
|
sqlite3_finalize(stmt);
|
|
return diskann_medoid_set(p, vec_col_idx, newMedoid, 0);
|
|
} else {
|
|
sqlite3_finalize(stmt);
|
|
return diskann_medoid_set(p, vec_col_idx, -1, 1);
|
|
}
|
|
}
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
// ============================================================
|
|
// DiskANN database I/O helpers
|
|
// ============================================================
|
|
|
|
/**
|
|
* Read a node's full data from _diskann_nodes.
|
|
* Returns blobs that must be freed by the caller with sqlite3_free().
|
|
*/
|
|
static int diskann_node_read(vec0_vtab *p, int vec_col_idx, i64 rowid,
|
|
u8 **outValidity, int *outValiditySize,
|
|
u8 **outNeighborIds, int *outNeighborIdsSize,
|
|
u8 **outQvecs, int *outQvecsSize) {
|
|
int rc;
|
|
if (!p->stmtDiskannNodeRead[vec_col_idx]) {
|
|
char *zSql = sqlite3_mprintf(
|
|
"SELECT neighbors_validity, neighbor_ids, neighbor_quantized_vectors "
|
|
"FROM " VEC0_SHADOW_DISKANN_NODES_N_NAME " WHERE rowid = ?",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
rc = sqlite3_prepare_v2(p->db, zSql, -1,
|
|
&p->stmtDiskannNodeRead[vec_col_idx], NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
}
|
|
|
|
sqlite3_stmt *stmt = p->stmtDiskannNodeRead[vec_col_idx];
|
|
sqlite3_reset(stmt);
|
|
sqlite3_bind_int64(stmt, 1, rowid);
|
|
|
|
rc = sqlite3_step(stmt);
|
|
if (rc != SQLITE_ROW) {
|
|
return SQLITE_ERROR;
|
|
}
|
|
|
|
int vs = sqlite3_column_bytes(stmt, 0);
|
|
int is = sqlite3_column_bytes(stmt, 1);
|
|
int qs = sqlite3_column_bytes(stmt, 2);
|
|
|
|
// Validate blob sizes against config expectations to detect truncated /
|
|
// corrupt data before any caller iterates using cfg->n_neighbors.
|
|
{
|
|
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
|
|
struct Vec0DiskannConfig *cfg = &col->diskann;
|
|
int expectedVs = diskann_validity_byte_size(cfg->n_neighbors);
|
|
int expectedIs = (int)diskann_neighbor_ids_byte_size(cfg->n_neighbors);
|
|
int expectedQs = (int)diskann_neighbor_qvecs_byte_size(
|
|
cfg->n_neighbors, cfg->quantizer_type, col->dimensions);
|
|
if (vs != expectedVs || is != expectedIs || qs != expectedQs) {
|
|
return SQLITE_CORRUPT;
|
|
}
|
|
}
|
|
|
|
u8 *v = sqlite3_malloc(vs);
|
|
u8 *ids = sqlite3_malloc(is);
|
|
u8 *qv = sqlite3_malloc(qs);
|
|
if (!v || !ids || !qv) {
|
|
sqlite3_free(v);
|
|
sqlite3_free(ids);
|
|
sqlite3_free(qv);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
const void *blobV = sqlite3_column_blob(stmt, 0);
|
|
const void *blobIds = sqlite3_column_blob(stmt, 1);
|
|
const void *blobQv = sqlite3_column_blob(stmt, 2);
|
|
if (!blobV || !blobIds || !blobQv) {
|
|
sqlite3_free(v);
|
|
sqlite3_free(ids);
|
|
sqlite3_free(qv);
|
|
return SQLITE_ERROR;
|
|
}
|
|
memcpy(v, blobV, vs);
|
|
memcpy(ids, blobIds, is);
|
|
memcpy(qv, blobQv, qs);
|
|
|
|
*outValidity = v; *outValiditySize = vs;
|
|
*outNeighborIds = ids; *outNeighborIdsSize = is;
|
|
*outQvecs = qv; *outQvecsSize = qs;
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
/**
|
|
* Write (INSERT OR REPLACE) a node's data to _diskann_nodes.
|
|
*/
|
|
static int diskann_node_write(vec0_vtab *p, int vec_col_idx, i64 rowid,
|
|
const u8 *validity, int validitySize,
|
|
const u8 *neighborIds, int neighborIdsSize,
|
|
const u8 *qvecs, int qvecsSize) {
|
|
int rc;
|
|
if (!p->stmtDiskannNodeWrite[vec_col_idx]) {
|
|
char *zSql = sqlite3_mprintf(
|
|
"INSERT OR REPLACE INTO " VEC0_SHADOW_DISKANN_NODES_N_NAME
|
|
" (rowid, neighbors_validity, neighbor_ids, neighbor_quantized_vectors) "
|
|
"VALUES (?, ?, ?, ?)",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
rc = sqlite3_prepare_v2(p->db, zSql, -1,
|
|
&p->stmtDiskannNodeWrite[vec_col_idx], NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
}
|
|
|
|
sqlite3_stmt *stmt = p->stmtDiskannNodeWrite[vec_col_idx];
|
|
sqlite3_reset(stmt);
|
|
sqlite3_bind_int64(stmt, 1, rowid);
|
|
sqlite3_bind_blob(stmt, 2, validity, validitySize, SQLITE_TRANSIENT);
|
|
sqlite3_bind_blob(stmt, 3, neighborIds, neighborIdsSize, SQLITE_TRANSIENT);
|
|
sqlite3_bind_blob(stmt, 4, qvecs, qvecsSize, SQLITE_TRANSIENT);
|
|
|
|
rc = sqlite3_step(stmt);
|
|
return (rc == SQLITE_DONE) ? SQLITE_OK : SQLITE_ERROR;
|
|
}
|
|
|
|
/**
|
|
* Read the full-precision vector for a given rowid from _vectors.
|
|
* Caller must free *outVector with sqlite3_free().
|
|
*/
|
|
static int diskann_vector_read(vec0_vtab *p, int vec_col_idx, i64 rowid,
|
|
void **outVector, int *outVectorSize) {
|
|
int rc;
|
|
if (!p->stmtVectorsRead[vec_col_idx]) {
|
|
char *zSql = sqlite3_mprintf(
|
|
"SELECT vector FROM " VEC0_SHADOW_VECTORS_N_NAME " WHERE rowid = ?",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
rc = sqlite3_prepare_v2(p->db, zSql, -1,
|
|
&p->stmtVectorsRead[vec_col_idx], NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
}
|
|
|
|
sqlite3_stmt *stmt = p->stmtVectorsRead[vec_col_idx];
|
|
sqlite3_reset(stmt);
|
|
sqlite3_bind_int64(stmt, 1, rowid);
|
|
|
|
rc = sqlite3_step(stmt);
|
|
if (rc != SQLITE_ROW) {
|
|
return SQLITE_ERROR;
|
|
}
|
|
|
|
int sz = sqlite3_column_bytes(stmt, 0);
|
|
const void *blob = sqlite3_column_blob(stmt, 0);
|
|
if (!blob || sz == 0) return SQLITE_ERROR;
|
|
void *vec = sqlite3_malloc(sz);
|
|
if (!vec) return SQLITE_NOMEM;
|
|
memcpy(vec, blob, sz);
|
|
|
|
*outVector = vec;
|
|
*outVectorSize = sz;
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
/**
|
|
* Write a full-precision vector to _vectors.
|
|
*/
|
|
static int diskann_vector_write(vec0_vtab *p, int vec_col_idx, i64 rowid,
|
|
const void *vector, int vectorSize) {
|
|
int rc;
|
|
if (!p->stmtVectorsInsert[vec_col_idx]) {
|
|
char *zSql = sqlite3_mprintf(
|
|
"INSERT OR REPLACE INTO " VEC0_SHADOW_VECTORS_N_NAME
|
|
" (rowid, vector) VALUES (?, ?)",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
rc = sqlite3_prepare_v2(p->db, zSql, -1,
|
|
&p->stmtVectorsInsert[vec_col_idx], NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
}
|
|
|
|
sqlite3_stmt *stmt = p->stmtVectorsInsert[vec_col_idx];
|
|
sqlite3_reset(stmt);
|
|
sqlite3_bind_int64(stmt, 1, rowid);
|
|
sqlite3_bind_blob(stmt, 2, vector, vectorSize, SQLITE_TRANSIENT);
|
|
|
|
rc = sqlite3_step(stmt);
|
|
return (rc == SQLITE_DONE) ? SQLITE_OK : SQLITE_ERROR;
|
|
}
|
|
|
|
// ============================================================
|
|
// DiskANN search data structures
|
|
// ============================================================
|
|
|
|
/**
|
|
* A sorted candidate list for greedy beam search.
|
|
*/
|
|
struct DiskannCandidateList {
|
|
struct Vec0DiskannCandidate *items;
|
|
int count;
|
|
int capacity;
|
|
};
|
|
|
|
static int diskann_candidate_list_init(struct DiskannCandidateList *list, int capacity) {
|
|
list->items = sqlite3_malloc(capacity * sizeof(struct Vec0DiskannCandidate));
|
|
if (!list->items) return SQLITE_NOMEM;
|
|
list->count = 0;
|
|
list->capacity = capacity;
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
static void diskann_candidate_list_free(struct DiskannCandidateList *list) {
|
|
sqlite3_free(list->items);
|
|
list->items = NULL;
|
|
list->count = 0;
|
|
list->capacity = 0;
|
|
}
|
|
|
|
/**
|
|
* Insert a candidate into the sorted list, maintaining sort order by distance.
|
|
* Deduplicates by rowid. If at capacity and new candidate is worse, discards it.
|
|
* Returns 1 if inserted, 0 if discarded.
|
|
*/
|
|
static int diskann_candidate_list_insert(
|
|
struct DiskannCandidateList *list, i64 rowid, f32 distance) {
|
|
|
|
// Check for duplicate
|
|
for (int i = 0; i < list->count; i++) {
|
|
if (list->items[i].rowid == rowid) {
|
|
// Update distance if better
|
|
if (distance < list->items[i].distance) {
|
|
list->items[i].distance = distance;
|
|
// Re-sort this item into position
|
|
struct Vec0DiskannCandidate tmp = list->items[i];
|
|
int j = i - 1;
|
|
while (j >= 0 && list->items[j].distance > tmp.distance) {
|
|
list->items[j + 1] = list->items[j];
|
|
j--;
|
|
}
|
|
list->items[j + 1] = tmp;
|
|
}
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
// If at capacity, check if new candidate is better than worst
|
|
if (list->count >= list->capacity) {
|
|
if (distance >= list->items[list->count - 1].distance) {
|
|
return 0; // Discard
|
|
}
|
|
list->count--; // Make room by dropping the worst
|
|
}
|
|
|
|
// Binary search for insertion point
|
|
int lo = 0, hi = list->count;
|
|
while (lo < hi) {
|
|
int mid = (lo + hi) / 2;
|
|
if (list->items[mid].distance < distance) {
|
|
lo = mid + 1;
|
|
} else {
|
|
hi = mid;
|
|
}
|
|
}
|
|
|
|
// Shift elements to make room
|
|
memmove(&list->items[lo + 1], &list->items[lo],
|
|
(list->count - lo) * sizeof(struct Vec0DiskannCandidate));
|
|
|
|
list->items[lo].rowid = rowid;
|
|
list->items[lo].distance = distance;
|
|
list->items[lo].visited = 0;
|
|
list->items[lo].confirmed = 0;
|
|
list->count++;
|
|
return 1;
|
|
}
|
|
|
|
/**
|
|
* Find the closest unvisited candidate. Returns its index, or -1 if none.
|
|
*/
|
|
static int diskann_candidate_list_next_unvisited(
|
|
const struct DiskannCandidateList *list) {
|
|
for (int i = 0; i < list->count; i++) {
|
|
if (!list->items[i].visited) return i;
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
|
|
|
|
/**
|
|
* Simple hash set for tracking visited rowids during search.
|
|
* Uses open addressing with linear probing.
|
|
*/
|
|
struct DiskannVisitedSet {
|
|
i64 *slots;
|
|
int capacity;
|
|
int count;
|
|
};
|
|
|
|
static int diskann_visited_set_init(struct DiskannVisitedSet *set, int capacity) {
|
|
// Round up to power of 2
|
|
int cap = 16;
|
|
while (cap < capacity) cap *= 2;
|
|
set->slots = sqlite3_malloc(cap * sizeof(i64));
|
|
if (!set->slots) return SQLITE_NOMEM;
|
|
memset(set->slots, 0, cap * sizeof(i64));
|
|
set->capacity = cap;
|
|
set->count = 0;
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
static void diskann_visited_set_free(struct DiskannVisitedSet *set) {
|
|
sqlite3_free(set->slots);
|
|
set->slots = NULL;
|
|
set->capacity = 0;
|
|
set->count = 0;
|
|
}
|
|
|
|
static int diskann_visited_set_contains(const struct DiskannVisitedSet *set, i64 rowid) {
|
|
if (rowid == 0) return 0; // 0 is our sentinel for empty
|
|
int mask = set->capacity - 1;
|
|
int idx = (int)(((u64)rowid * 0x9E3779B97F4A7C15ULL) >> 32) & mask;
|
|
for (int i = 0; i < set->capacity; i++) {
|
|
int slot = (idx + i) & mask;
|
|
if (set->slots[slot] == 0) return 0;
|
|
if (set->slots[slot] == rowid) return 1;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
static int diskann_visited_set_insert(struct DiskannVisitedSet *set, i64 rowid) {
|
|
if (rowid == 0) return 0;
|
|
int mask = set->capacity - 1;
|
|
int idx = (int)(((u64)rowid * 0x9E3779B97F4A7C15ULL) >> 32) & mask;
|
|
for (int i = 0; i < set->capacity; i++) {
|
|
int slot = (idx + i) & mask;
|
|
if (set->slots[slot] == 0) {
|
|
set->slots[slot] = rowid;
|
|
set->count++;
|
|
return 1;
|
|
}
|
|
if (set->slots[slot] == rowid) return 0; // Already there
|
|
}
|
|
return 0; // Full (shouldn't happen with proper sizing)
|
|
}
|
|
|
|
// ============================================================
|
|
// DiskANN greedy beam search (LM-Search)
|
|
// ============================================================
|
|
|
|
/**
|
|
* Perform LM-Search: greedy beam search over the DiskANN graph.
|
|
* Follows Algorithm 1 from the LM-DiskANN paper.
|
|
*/
|
|
static int diskann_search(
|
|
vec0_vtab *p, int vec_col_idx,
|
|
const void *queryVector, size_t dimensions,
|
|
enum VectorElementType elementType,
|
|
int k, int searchListSize,
|
|
i64 *outRowids, f32 *outDistances, int *outCount) {
|
|
|
|
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
|
|
struct Vec0DiskannConfig *cfg = &col->diskann;
|
|
int rc;
|
|
|
|
if (searchListSize <= 0) {
|
|
searchListSize = cfg->search_list_size_search > 0 ? cfg->search_list_size_search : cfg->search_list_size;
|
|
}
|
|
if (searchListSize < k) {
|
|
searchListSize = k;
|
|
}
|
|
|
|
// 1. Get the medoid (entry point)
|
|
i64 medoid;
|
|
int isEmpty;
|
|
rc = diskann_medoid_get(p, vec_col_idx, &medoid, &isEmpty);
|
|
if (rc != SQLITE_OK) return rc;
|
|
if (isEmpty) {
|
|
*outCount = 0;
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
// 2. Compute distance from query to medoid using full-precision vector
|
|
void *medoidVector = NULL;
|
|
int medoidVectorSize;
|
|
rc = diskann_vector_read(p, vec_col_idx, medoid, &medoidVector, &medoidVectorSize);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
f32 medoidDist = vec0_distance_full(queryVector, medoidVector,
|
|
dimensions, elementType,
|
|
col->distance_metric);
|
|
sqlite3_free(medoidVector);
|
|
|
|
// 3. Initialize candidate list and visited set
|
|
struct DiskannCandidateList candidates;
|
|
rc = diskann_candidate_list_init(&candidates, searchListSize);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
struct DiskannVisitedSet visited;
|
|
rc = diskann_visited_set_init(&visited, searchListSize * 4);
|
|
if (rc != SQLITE_OK) {
|
|
diskann_candidate_list_free(&candidates);
|
|
return rc;
|
|
}
|
|
|
|
// Seed with medoid (confirmed — we already read its vector above)
|
|
diskann_candidate_list_insert(&candidates, medoid, medoidDist);
|
|
candidates.items[0].confirmed = 1;
|
|
|
|
// Pre-quantize query vector once for all quantized distance comparisons
|
|
u8 *queryQuantized = NULL;
|
|
if (elementType == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
|
|
queryQuantized = diskann_quantize_query(
|
|
(const f32 *)queryVector, dimensions, cfg->quantizer_type);
|
|
}
|
|
|
|
// 4. Greedy beam search loop (Algorithm 1 from LM-DiskANN paper)
|
|
while (1) {
|
|
int nextIdx = diskann_candidate_list_next_unvisited(&candidates);
|
|
if (nextIdx < 0) break;
|
|
|
|
struct Vec0DiskannCandidate *current = &candidates.items[nextIdx];
|
|
current->visited = 1;
|
|
i64 currentRowid = current->rowid;
|
|
|
|
// Read the node's neighbor data
|
|
u8 *validity = NULL, *neighborIds = NULL, *qvecs = NULL;
|
|
int validitySize, neighborIdsSize, qvecsSize;
|
|
rc = diskann_node_read(p, vec_col_idx, currentRowid,
|
|
&validity, &validitySize,
|
|
&neighborIds, &neighborIdsSize,
|
|
&qvecs, &qvecsSize);
|
|
if (rc != SQLITE_OK) {
|
|
continue; // Skip if node doesn't exist
|
|
}
|
|
|
|
// Insert all valid neighbors with approximate (quantized) distances
|
|
for (int i = 0; i < cfg->n_neighbors; i++) {
|
|
if (!diskann_validity_get(validity, i)) continue;
|
|
|
|
i64 neighborRowid = diskann_neighbor_id_get(neighborIds, i);
|
|
|
|
if (diskann_visited_set_contains(&visited, neighborRowid)) continue;
|
|
|
|
const u8 *neighborQvec = diskann_neighbor_qvec_get(
|
|
qvecs, i, cfg->quantizer_type, dimensions);
|
|
|
|
f32 approxDist;
|
|
if (queryQuantized) {
|
|
approxDist = diskann_distance_quantized_precomputed(
|
|
queryQuantized, neighborQvec, dimensions,
|
|
cfg->quantizer_type, col->distance_metric);
|
|
} else {
|
|
approxDist = diskann_distance_quantized(
|
|
queryVector, neighborQvec, dimensions,
|
|
cfg->quantizer_type, col->distance_metric);
|
|
}
|
|
|
|
diskann_candidate_list_insert(&candidates, neighborRowid, approxDist);
|
|
}
|
|
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
|
|
// Add to visited set
|
|
diskann_visited_set_insert(&visited, currentRowid);
|
|
|
|
// Paper line 13: Re-rank p* using full-precision distance
|
|
// We already have exact distance for medoid; for others, update now
|
|
void *fullVec = NULL;
|
|
int fullVecSize;
|
|
rc = diskann_vector_read(p, vec_col_idx, currentRowid, &fullVec, &fullVecSize);
|
|
if (rc == SQLITE_OK) {
|
|
f32 exactDist = vec0_distance_full(queryVector, fullVec,
|
|
dimensions, elementType,
|
|
col->distance_metric);
|
|
sqlite3_free(fullVec);
|
|
// Update distance in candidate list and re-sort
|
|
diskann_candidate_list_insert(&candidates, currentRowid, exactDist);
|
|
// Mark as confirmed (vector exists, distance is exact)
|
|
for (int ci = 0; ci < candidates.count; ci++) {
|
|
if (candidates.items[ci].rowid == currentRowid) {
|
|
candidates.items[ci].confirmed = 1;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
// If vector read failed, candidate stays unconfirmed (stale edge to deleted node)
|
|
}
|
|
|
|
// 5. Output results — only include confirmed candidates (whose vectors exist)
|
|
int resultCount = 0;
|
|
for (int i = 0; i < candidates.count && resultCount < k; i++) {
|
|
if (candidates.items[i].confirmed) {
|
|
outRowids[resultCount] = candidates.items[i].rowid;
|
|
outDistances[resultCount] = candidates.items[i].distance;
|
|
resultCount++;
|
|
}
|
|
}
|
|
*outCount = resultCount;
|
|
|
|
sqlite3_free(queryQuantized);
|
|
diskann_candidate_list_free(&candidates);
|
|
diskann_visited_set_free(&visited);
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
// ============================================================
|
|
// DiskANN RobustPrune (Algorithm 4 from LM-DiskANN paper)
|
|
// ============================================================
|
|
|
|
/**
|
|
* RobustPrune: Select up to max_neighbors neighbors for node p from a
|
|
* candidate set, using alpha-pruning for diversity.
|
|
*
|
|
* Following Algorithm 4 (LM-Prune):
|
|
* C = C union N_out(p) \ {p}
|
|
* N_out(p) = empty
|
|
* while C not empty:
|
|
* p* = argmin d(p, c) for c in C
|
|
* N_out(p).insert(p*)
|
|
* if |N_out(p)| >= R: break
|
|
* for each p' in C:
|
|
* if alpha * d(p*, p') <= d(p, p'): remove p' from C
|
|
*/
|
|
/**
|
|
* Pure function: given pre-sorted candidates and a distance matrix, select
|
|
* up to max_neighbors using alpha-pruning. inter_distances is a flattened
|
|
* num_candidates x num_candidates matrix where inter_distances[i*num_candidates+j]
|
|
* = d(candidate_i, candidate_j). p_distances[i] = d(p, candidate_i), already sorted.
|
|
* outSelected[i] = 1 if selected. Returns count of selected.
|
|
*/
|
|
int diskann_prune_select(
|
|
const f32 *inter_distances, const f32 *p_distances,
|
|
int num_candidates, f32 alpha, int max_neighbors,
|
|
int *outSelected, int *outCount) {
|
|
|
|
if (num_candidates == 0) {
|
|
*outCount = 0;
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
u8 *active = sqlite3_malloc(num_candidates);
|
|
if (!active) return SQLITE_NOMEM;
|
|
memset(active, 1, num_candidates);
|
|
memset(outSelected, 0, num_candidates * sizeof(int));
|
|
|
|
int selectedCount = 0;
|
|
|
|
for (int round = 0; round < num_candidates && selectedCount < max_neighbors; round++) {
|
|
int bestIdx = -1;
|
|
for (int i = 0; i < num_candidates; i++) {
|
|
if (active[i]) { bestIdx = i; break; }
|
|
}
|
|
if (bestIdx < 0) break;
|
|
|
|
outSelected[bestIdx] = 1;
|
|
selectedCount++;
|
|
active[bestIdx] = 0;
|
|
|
|
for (int i = 0; i < num_candidates; i++) {
|
|
if (!active[i]) continue;
|
|
f32 dist_best_to_i = inter_distances[bestIdx * num_candidates + i];
|
|
if (alpha * dist_best_to_i <= p_distances[i]) {
|
|
active[i] = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
*outCount = selectedCount;
|
|
sqlite3_free(active);
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
static int diskann_robust_prune(
|
|
vec0_vtab *p, int vec_col_idx,
|
|
i64 p_rowid, const void *p_vector,
|
|
i64 *candidates, f32 *candidate_distances, int num_candidates,
|
|
f32 alpha, int max_neighbors,
|
|
i64 *outNeighborRowids, int *outNeighborCount) {
|
|
|
|
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
|
|
int rc;
|
|
|
|
// Remove p itself from candidates
|
|
for (int i = 0; i < num_candidates; i++) {
|
|
if (candidates[i] == p_rowid) {
|
|
candidates[i] = candidates[num_candidates - 1];
|
|
candidate_distances[i] = candidate_distances[num_candidates - 1];
|
|
num_candidates--;
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (num_candidates == 0) {
|
|
*outNeighborCount = 0;
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
// Sort candidates by distance to p (ascending) - insertion sort
|
|
for (int i = 1; i < num_candidates; i++) {
|
|
f32 tmpDist = candidate_distances[i];
|
|
i64 tmpRowid = candidates[i];
|
|
int j = i - 1;
|
|
while (j >= 0 && candidate_distances[j] > tmpDist) {
|
|
candidate_distances[j + 1] = candidate_distances[j];
|
|
candidates[j + 1] = candidates[j];
|
|
j--;
|
|
}
|
|
candidate_distances[j + 1] = tmpDist;
|
|
candidates[j + 1] = tmpRowid;
|
|
}
|
|
|
|
// Active flags
|
|
u8 *active = sqlite3_malloc(num_candidates);
|
|
if (!active) return SQLITE_NOMEM;
|
|
memset(active, 1, num_candidates);
|
|
|
|
// Cache full-precision vectors for inter-candidate distance
|
|
void **candidateVectors = sqlite3_malloc(num_candidates * sizeof(void *));
|
|
if (!candidateVectors) {
|
|
sqlite3_free(active);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
memset(candidateVectors, 0, num_candidates * sizeof(void *));
|
|
|
|
int selectedCount = 0;
|
|
|
|
for (int round = 0; round < num_candidates && selectedCount < max_neighbors; round++) {
|
|
// Find closest active candidate
|
|
int bestIdx = -1;
|
|
for (int i = 0; i < num_candidates; i++) {
|
|
if (active[i]) { bestIdx = i; break; }
|
|
}
|
|
if (bestIdx < 0) break;
|
|
|
|
// Select this candidate
|
|
outNeighborRowids[selectedCount] = candidates[bestIdx];
|
|
selectedCount++;
|
|
active[bestIdx] = 0;
|
|
|
|
// Load selected candidate's vector
|
|
if (!candidateVectors[bestIdx]) {
|
|
int vecSize;
|
|
rc = diskann_vector_read(p, vec_col_idx, candidates[bestIdx],
|
|
&candidateVectors[bestIdx], &vecSize);
|
|
if (rc != SQLITE_OK) continue;
|
|
}
|
|
|
|
// Alpha-prune: remove candidates covered by the selected neighbor
|
|
for (int i = 0; i < num_candidates; i++) {
|
|
if (!active[i]) continue;
|
|
|
|
if (!candidateVectors[i]) {
|
|
int vecSize;
|
|
rc = diskann_vector_read(p, vec_col_idx, candidates[i],
|
|
&candidateVectors[i], &vecSize);
|
|
if (rc != SQLITE_OK) continue;
|
|
}
|
|
|
|
f32 dist_selected_to_i = vec0_distance_full(
|
|
candidateVectors[bestIdx], candidateVectors[i],
|
|
col->dimensions, col->element_type, col->distance_metric);
|
|
|
|
if (alpha * dist_selected_to_i <= candidate_distances[i]) {
|
|
active[i] = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
*outNeighborCount = selectedCount;
|
|
|
|
for (int i = 0; i < num_candidates; i++) {
|
|
sqlite3_free(candidateVectors[i]);
|
|
}
|
|
sqlite3_free(candidateVectors);
|
|
sqlite3_free(active);
|
|
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
/**
|
|
* After RobustPrune selects neighbors, build the node blobs and write to DB.
|
|
* Quantizes each neighbor's vector and packs into the node format.
|
|
*/
|
|
static int diskann_write_pruned_neighbors(
|
|
vec0_vtab *p, int vec_col_idx, i64 nodeRowid,
|
|
const i64 *neighborRowids, int neighborCount) {
|
|
|
|
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
|
|
struct Vec0DiskannConfig *cfg = &col->diskann;
|
|
int rc;
|
|
|
|
u8 *validity, *neighborIds, *qvecs;
|
|
int validitySize, neighborIdsSize, qvecsSize;
|
|
rc = diskann_node_init(cfg->n_neighbors, cfg->quantizer_type,
|
|
col->dimensions,
|
|
&validity, &validitySize,
|
|
&neighborIds, &neighborIdsSize,
|
|
&qvecs, &qvecsSize);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
size_t qvecSize = diskann_quantized_vector_byte_size(
|
|
cfg->quantizer_type, col->dimensions);
|
|
u8 *qvec = sqlite3_malloc(qvecSize);
|
|
if (!qvec) {
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
for (int i = 0; i < neighborCount && i < cfg->n_neighbors; i++) {
|
|
void *neighborVec = NULL;
|
|
int neighborVecSize;
|
|
rc = diskann_vector_read(p, vec_col_idx, neighborRowids[i],
|
|
&neighborVec, &neighborVecSize);
|
|
if (rc != SQLITE_OK) continue;
|
|
|
|
if (col->element_type == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
|
|
diskann_quantize_vector((const f32 *)neighborVec, col->dimensions,
|
|
cfg->quantizer_type, qvec);
|
|
} else {
|
|
memcpy(qvec, neighborVec,
|
|
qvecSize < (size_t)neighborVecSize ? qvecSize : (size_t)neighborVecSize);
|
|
}
|
|
|
|
diskann_node_set_neighbor(validity, neighborIds, qvecs, i,
|
|
neighborRowids[i], qvec,
|
|
cfg->quantizer_type, col->dimensions);
|
|
|
|
sqlite3_free(neighborVec);
|
|
}
|
|
sqlite3_free(qvec);
|
|
|
|
rc = diskann_node_write(p, vec_col_idx, nodeRowid,
|
|
validity, validitySize,
|
|
neighborIds, neighborIdsSize,
|
|
qvecs, qvecsSize);
|
|
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
return rc;
|
|
}
|
|
|
|
// ============================================================
|
|
// DiskANN insert (Algorithm 2 from LM-DiskANN paper)
|
|
// ============================================================
|
|
|
|
/**
|
|
* Add a reverse edge: make target_rowid a neighbor of node_rowid.
|
|
* If node is full, run RobustPrune.
|
|
*/
|
|
static int diskann_add_reverse_edge(
|
|
vec0_vtab *p, int vec_col_idx,
|
|
i64 node_rowid, i64 target_rowid, const void *target_vector) {
|
|
|
|
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
|
|
struct Vec0DiskannConfig *cfg = &col->diskann;
|
|
int rc;
|
|
|
|
u8 *validity = NULL, *neighborIds = NULL, *qvecs = NULL;
|
|
int validitySize, neighborIdsSize, qvecsSize;
|
|
rc = diskann_node_read(p, vec_col_idx, node_rowid,
|
|
&validity, &validitySize,
|
|
&neighborIds, &neighborIdsSize,
|
|
&qvecs, &qvecsSize);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
int currentCount = diskann_validity_count(validity, cfg->n_neighbors);
|
|
|
|
// Check if target is already a neighbor
|
|
for (int i = 0; i < cfg->n_neighbors; i++) {
|
|
if (diskann_validity_get(validity, i) &&
|
|
diskann_neighbor_id_get(neighborIds, i) == target_rowid) {
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
return SQLITE_OK;
|
|
}
|
|
}
|
|
|
|
if (currentCount < cfg->n_neighbors) {
|
|
// Room available: find first empty slot
|
|
for (int i = 0; i < cfg->n_neighbors; i++) {
|
|
if (!diskann_validity_get(validity, i)) {
|
|
size_t qvecSize = diskann_quantized_vector_byte_size(
|
|
cfg->quantizer_type, col->dimensions);
|
|
u8 *qvec = sqlite3_malloc(qvecSize);
|
|
if (!qvec) {
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
if (col->element_type == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
|
|
diskann_quantize_vector((const f32 *)target_vector, col->dimensions,
|
|
cfg->quantizer_type, qvec);
|
|
} else {
|
|
size_t vbs = vector_column_byte_size(*col);
|
|
memcpy(qvec, target_vector, qvecSize < vbs ? qvecSize : vbs);
|
|
}
|
|
|
|
diskann_node_set_neighbor(validity, neighborIds, qvecs, i,
|
|
target_rowid, qvec,
|
|
cfg->quantizer_type, col->dimensions);
|
|
sqlite3_free(qvec);
|
|
break;
|
|
}
|
|
}
|
|
|
|
rc = diskann_node_write(p, vec_col_idx, node_rowid,
|
|
validity, validitySize,
|
|
neighborIds, neighborIdsSize,
|
|
qvecs, qvecsSize);
|
|
} else {
|
|
// Full: lazy replacement — use quantized distances to find the worst
|
|
// existing neighbor and replace it if target is closer. This avoids
|
|
// reading all neighbors' float vectors (the expensive RobustPrune path).
|
|
|
|
// Quantize the node's vector and the target vector for comparison
|
|
void *nodeVector = NULL;
|
|
int nodeVecSize;
|
|
rc = diskann_vector_read(p, vec_col_idx, node_rowid,
|
|
&nodeVector, &nodeVecSize);
|
|
if (rc != SQLITE_OK) {
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
return rc;
|
|
}
|
|
|
|
// Quantize target for node-level comparison
|
|
size_t qvecSize = diskann_quantized_vector_byte_size(
|
|
cfg->quantizer_type, col->dimensions);
|
|
u8 *targetQ = sqlite3_malloc(qvecSize);
|
|
u8 *nodeQ = sqlite3_malloc(qvecSize);
|
|
if (!targetQ || !nodeQ) {
|
|
sqlite3_free(targetQ);
|
|
sqlite3_free(nodeQ);
|
|
sqlite3_free(nodeVector);
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
if (col->element_type == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
|
|
diskann_quantize_vector((const f32 *)target_vector, col->dimensions,
|
|
cfg->quantizer_type, targetQ);
|
|
diskann_quantize_vector((const f32 *)nodeVector, col->dimensions,
|
|
cfg->quantizer_type, nodeQ);
|
|
} else {
|
|
memcpy(targetQ, target_vector, qvecSize);
|
|
memcpy(nodeQ, nodeVector, qvecSize);
|
|
}
|
|
|
|
// Compute quantized distance from node to target
|
|
f32 targetDist = diskann_distance_quantized_precomputed(
|
|
nodeQ, targetQ, col->dimensions,
|
|
cfg->quantizer_type, col->distance_metric);
|
|
|
|
// Find the worst (farthest) existing neighbor using quantized distances
|
|
int worstIdx = -1;
|
|
f32 worstDist = -1.0f;
|
|
for (int i = 0; i < cfg->n_neighbors; i++) {
|
|
if (!diskann_validity_get(validity, i)) continue;
|
|
const u8 *nqvec = diskann_neighbor_qvec_get(
|
|
qvecs, i, cfg->quantizer_type, col->dimensions);
|
|
f32 d = diskann_distance_quantized_precomputed(
|
|
nodeQ, nqvec, col->dimensions,
|
|
cfg->quantizer_type, col->distance_metric);
|
|
if (d > worstDist) {
|
|
worstDist = d;
|
|
worstIdx = i;
|
|
}
|
|
}
|
|
|
|
// Replace worst neighbor if target is closer
|
|
if (worstIdx >= 0 && targetDist < worstDist) {
|
|
diskann_node_set_neighbor(validity, neighborIds, qvecs, worstIdx,
|
|
target_rowid, targetQ,
|
|
cfg->quantizer_type, col->dimensions);
|
|
rc = diskann_node_write(p, vec_col_idx, node_rowid,
|
|
validity, validitySize,
|
|
neighborIds, neighborIdsSize,
|
|
qvecs, qvecsSize);
|
|
} else {
|
|
rc = SQLITE_OK; // target is farther than all existing neighbors, skip
|
|
}
|
|
|
|
sqlite3_free(targetQ);
|
|
sqlite3_free(nodeQ);
|
|
sqlite3_free(nodeVector);
|
|
}
|
|
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
return rc;
|
|
}
|
|
|
|
// ============================================================
|
|
// DiskANN buffer helpers (for batched inserts)
|
|
// ============================================================
|
|
|
|
/**
|
|
* Insert a vector into the _diskann_buffer table.
|
|
*/
|
|
static int diskann_buffer_write(vec0_vtab *p, int vec_col_idx,
|
|
i64 rowid, const void *vector, int vectorSize) {
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *zSql = sqlite3_mprintf(
|
|
"INSERT INTO " VEC0_SHADOW_DISKANN_BUFFER_N_NAME
|
|
" (rowid, vector) VALUES (?, ?)",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
int rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
sqlite3_bind_int64(stmt, 1, rowid);
|
|
sqlite3_bind_blob(stmt, 2, vector, vectorSize, SQLITE_STATIC);
|
|
rc = sqlite3_step(stmt);
|
|
sqlite3_finalize(stmt);
|
|
return (rc == SQLITE_DONE) ? SQLITE_OK : SQLITE_ERROR;
|
|
}
|
|
|
|
/**
|
|
* Delete a vector from the _diskann_buffer table.
|
|
*/
|
|
static int diskann_buffer_delete(vec0_vtab *p, int vec_col_idx, i64 rowid) {
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *zSql = sqlite3_mprintf(
|
|
"DELETE FROM " VEC0_SHADOW_DISKANN_BUFFER_N_NAME " WHERE rowid = ?",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
int rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
sqlite3_bind_int64(stmt, 1, rowid);
|
|
rc = sqlite3_step(stmt);
|
|
sqlite3_finalize(stmt);
|
|
return (rc == SQLITE_DONE) ? SQLITE_OK : SQLITE_ERROR;
|
|
}
|
|
|
|
/**
|
|
* Check if a rowid exists in the _diskann_buffer table.
|
|
* Returns SQLITE_OK and sets *exists to 1 if found, 0 if not.
|
|
*/
|
|
static int diskann_buffer_exists(vec0_vtab *p, int vec_col_idx,
|
|
i64 rowid, int *exists) {
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *zSql = sqlite3_mprintf(
|
|
"SELECT 1 FROM " VEC0_SHADOW_DISKANN_BUFFER_N_NAME " WHERE rowid = ?",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
int rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
sqlite3_bind_int64(stmt, 1, rowid);
|
|
rc = sqlite3_step(stmt);
|
|
*exists = (rc == SQLITE_ROW) ? 1 : 0;
|
|
sqlite3_finalize(stmt);
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
/**
|
|
* Get the count of rows in the _diskann_buffer table.
|
|
*/
|
|
static int diskann_buffer_count(vec0_vtab *p, int vec_col_idx, i64 *count) {
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *zSql = sqlite3_mprintf(
|
|
"SELECT count(*) FROM " VEC0_SHADOW_DISKANN_BUFFER_N_NAME,
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
int rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
rc = sqlite3_step(stmt);
|
|
if (rc == SQLITE_ROW) {
|
|
*count = sqlite3_column_int64(stmt, 0);
|
|
sqlite3_finalize(stmt);
|
|
return SQLITE_OK;
|
|
}
|
|
sqlite3_finalize(stmt);
|
|
return SQLITE_ERROR;
|
|
}
|
|
|
|
// Forward declaration: diskann_insert_graph does the actual graph insertion
|
|
static int diskann_insert_graph(vec0_vtab *p, int vec_col_idx,
|
|
i64 rowid, const void *vector);
|
|
|
|
/**
|
|
* Flush all buffered vectors into the DiskANN graph.
|
|
* Iterates over _diskann_buffer rows and calls diskann_insert_graph for each.
|
|
*/
|
|
static int diskann_flush_buffer(vec0_vtab *p, int vec_col_idx) {
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *zSql = sqlite3_mprintf(
|
|
"SELECT rowid, vector FROM " VEC0_SHADOW_DISKANN_BUFFER_N_NAME,
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
int rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
while ((rc = sqlite3_step(stmt)) == SQLITE_ROW) {
|
|
i64 rowid = sqlite3_column_int64(stmt, 0);
|
|
const void *vector = sqlite3_column_blob(stmt, 1);
|
|
if (!vector) continue;
|
|
// Note: vector is already written to _vectors table, so
|
|
// diskann_insert_graph will skip re-writing it (vector already exists).
|
|
// We call the graph-only insert path.
|
|
int insertRc = diskann_insert_graph(p, vec_col_idx, rowid, vector);
|
|
if (insertRc != SQLITE_OK) {
|
|
sqlite3_finalize(stmt);
|
|
return insertRc;
|
|
}
|
|
}
|
|
sqlite3_finalize(stmt);
|
|
|
|
// Clear the buffer
|
|
zSql = sqlite3_mprintf(
|
|
"DELETE FROM " VEC0_SHADOW_DISKANN_BUFFER_N_NAME,
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
rc = sqlite3_step(stmt);
|
|
sqlite3_finalize(stmt);
|
|
return (rc == SQLITE_DONE) ? SQLITE_OK : SQLITE_ERROR;
|
|
}
|
|
|
|
/**
|
|
* Insert a new vector into the DiskANN graph (graph-only path).
|
|
* The vector must already be written to _vectors table.
|
|
* This is the core graph insertion logic (Algorithm 2: LM-Insert).
|
|
*/
|
|
static int diskann_insert_graph(vec0_vtab *p, int vec_col_idx,
|
|
i64 rowid, const void *vector) {
|
|
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
|
|
struct Vec0DiskannConfig *cfg = &col->diskann;
|
|
int rc;
|
|
|
|
// Handle first insert (empty graph)
|
|
i64 medoid;
|
|
int isEmpty;
|
|
rc = diskann_medoid_get(p, vec_col_idx, &medoid, &isEmpty);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
if (isEmpty) {
|
|
u8 *validity, *neighborIds, *qvecs;
|
|
int validitySize, neighborIdsSize, qvecsSize;
|
|
rc = diskann_node_init(cfg->n_neighbors, cfg->quantizer_type,
|
|
col->dimensions,
|
|
&validity, &validitySize,
|
|
&neighborIds, &neighborIdsSize,
|
|
&qvecs, &qvecsSize);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
rc = diskann_node_write(p, vec_col_idx, rowid,
|
|
validity, validitySize,
|
|
neighborIds, neighborIdsSize,
|
|
qvecs, qvecsSize);
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
return diskann_medoid_set(p, vec_col_idx, rowid, 0);
|
|
}
|
|
|
|
// Search for nearest neighbors
|
|
int L = cfg->search_list_size_insert > 0 ? cfg->search_list_size_insert : cfg->search_list_size;
|
|
i64 *searchRowids = sqlite3_malloc(L * sizeof(i64));
|
|
f32 *searchDistances = sqlite3_malloc(L * sizeof(f32));
|
|
if (!searchRowids || !searchDistances) {
|
|
sqlite3_free(searchRowids);
|
|
sqlite3_free(searchDistances);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
int searchCount;
|
|
rc = diskann_search(p, vec_col_idx, vector, col->dimensions,
|
|
col->element_type, L, L,
|
|
searchRowids, searchDistances, &searchCount);
|
|
if (rc != SQLITE_OK) {
|
|
sqlite3_free(searchRowids);
|
|
sqlite3_free(searchDistances);
|
|
return rc;
|
|
}
|
|
|
|
// RobustPrune to select neighbors for x
|
|
i64 *selectedNeighbors = sqlite3_malloc(cfg->n_neighbors * sizeof(i64));
|
|
int selectedCount = 0;
|
|
if (!selectedNeighbors) {
|
|
sqlite3_free(searchRowids);
|
|
sqlite3_free(searchDistances);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
rc = diskann_robust_prune(p, vec_col_idx, rowid, vector,
|
|
searchRowids, searchDistances, searchCount,
|
|
cfg->alpha, cfg->n_neighbors,
|
|
selectedNeighbors, &selectedCount);
|
|
sqlite3_free(searchRowids);
|
|
sqlite3_free(searchDistances);
|
|
if (rc != SQLITE_OK) {
|
|
sqlite3_free(selectedNeighbors);
|
|
return rc;
|
|
}
|
|
|
|
// Write x's node with selected neighbors
|
|
rc = diskann_write_pruned_neighbors(p, vec_col_idx, rowid,
|
|
selectedNeighbors, selectedCount);
|
|
if (rc != SQLITE_OK) {
|
|
sqlite3_free(selectedNeighbors);
|
|
return rc;
|
|
}
|
|
|
|
// Add bidirectional edges
|
|
for (int i = 0; i < selectedCount; i++) {
|
|
diskann_add_reverse_edge(p, vec_col_idx,
|
|
selectedNeighbors[i], rowid, vector);
|
|
}
|
|
|
|
sqlite3_free(selectedNeighbors);
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
/**
|
|
* Insert a new vector into the DiskANN index (Algorithm 2: LM-Insert).
|
|
* When buffer_threshold > 0, vectors are buffered and flushed in batch.
|
|
*/
|
|
static int diskann_insert(vec0_vtab *p, int vec_col_idx,
|
|
i64 rowid, const void *vector) {
|
|
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
|
|
struct Vec0DiskannConfig *cfg = &col->diskann;
|
|
int rc;
|
|
size_t vectorSize = vector_column_byte_size(*col);
|
|
|
|
// 1. Write full-precision vector to _vectors table (always needed for queries)
|
|
rc = diskann_vector_write(p, vec_col_idx, rowid, vector, (int)vectorSize);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
// 2. If buffering is enabled, write to buffer instead of graph
|
|
if (cfg->buffer_threshold > 0) {
|
|
rc = diskann_buffer_write(p, vec_col_idx, rowid, vector, (int)vectorSize);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
i64 count;
|
|
rc = diskann_buffer_count(p, vec_col_idx, &count);
|
|
if (rc != SQLITE_OK) return rc;
|
|
|
|
if (count >= cfg->buffer_threshold) {
|
|
return diskann_flush_buffer(p, vec_col_idx);
|
|
}
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
// 3. Legacy per-row insert directly into graph
|
|
return diskann_insert_graph(p, vec_col_idx, rowid, vector);
|
|
}
|
|
|
|
/**
|
|
* Returns 1 if ALL vector columns in this table are DiskANN-indexed.
|
|
*/
|
|
// ============================================================
|
|
// DiskANN delete (Algorithm 3 from LM-DiskANN paper)
|
|
// ============================================================
|
|
|
|
static int diskann_node_delete(vec0_vtab *p, int vec_col_idx, i64 rowid) {
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *zSql = sqlite3_mprintf(
|
|
"DELETE FROM " VEC0_SHADOW_DISKANN_NODES_N_NAME " WHERE rowid = ?",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
int rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
sqlite3_bind_int64(stmt, 1, rowid);
|
|
rc = sqlite3_step(stmt);
|
|
sqlite3_finalize(stmt);
|
|
return (rc == SQLITE_DONE) ? SQLITE_OK : SQLITE_ERROR;
|
|
}
|
|
|
|
static int diskann_vector_delete(vec0_vtab *p, int vec_col_idx, i64 rowid) {
|
|
sqlite3_stmt *stmt = NULL;
|
|
char *zSql = sqlite3_mprintf(
|
|
"DELETE FROM " VEC0_SHADOW_VECTORS_N_NAME " WHERE rowid = ?",
|
|
p->schemaName, p->tableName, vec_col_idx);
|
|
if (!zSql) return SQLITE_NOMEM;
|
|
int rc = sqlite3_prepare_v2(p->db, zSql, -1, &stmt, NULL);
|
|
sqlite3_free(zSql);
|
|
if (rc != SQLITE_OK) return rc;
|
|
sqlite3_bind_int64(stmt, 1, rowid);
|
|
rc = sqlite3_step(stmt);
|
|
sqlite3_finalize(stmt);
|
|
return (rc == SQLITE_DONE) ? SQLITE_OK : SQLITE_ERROR;
|
|
}
|
|
|
|
/**
|
|
* Repair graph after deleting a node. Following Algorithm 3 (LM-Delete):
|
|
* For each neighbor n of the deleted node, add deleted node's other neighbors
|
|
* to n's candidate set, then remove the deleted node from n's neighbor list.
|
|
* Uses simple slot replacement rather than full RobustPrune for performance.
|
|
*/
|
|
static int diskann_repair_reverse_edges(
|
|
vec0_vtab *p, int vec_col_idx, i64 deleted_rowid,
|
|
const i64 *deleted_neighbors, int deleted_neighbor_count) {
|
|
|
|
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
|
|
struct Vec0DiskannConfig *cfg = &col->diskann;
|
|
int rc;
|
|
|
|
// For each neighbor of the deleted node, fix their neighbor list
|
|
for (int dn = 0; dn < deleted_neighbor_count; dn++) {
|
|
i64 nodeRowid = deleted_neighbors[dn];
|
|
|
|
u8 *validity = NULL, *neighborIds = NULL, *qvecs = NULL;
|
|
int vs, nis, qs;
|
|
rc = diskann_node_read(p, vec_col_idx, nodeRowid,
|
|
&validity, &vs, &neighborIds, &nis, &qvecs, &qs);
|
|
if (rc != SQLITE_OK) continue;
|
|
|
|
// Find and clear the deleted node's slot
|
|
int clearedSlot = -1;
|
|
for (int i = 0; i < cfg->n_neighbors; i++) {
|
|
if (diskann_validity_get(validity, i) &&
|
|
diskann_neighbor_id_get(neighborIds, i) == deleted_rowid) {
|
|
diskann_node_clear_neighbor(validity, neighborIds, qvecs, i,
|
|
cfg->quantizer_type, col->dimensions);
|
|
clearedSlot = i;
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (clearedSlot >= 0) {
|
|
// Try to fill the cleared slot with one of the deleted node's other neighbors
|
|
for (int di = 0; di < deleted_neighbor_count; di++) {
|
|
i64 candidate = deleted_neighbors[di];
|
|
if (candidate == nodeRowid || candidate == deleted_rowid) continue;
|
|
|
|
// Check not already a neighbor
|
|
int alreadyNeighbor = 0;
|
|
for (int ni = 0; ni < cfg->n_neighbors; ni++) {
|
|
if (diskann_validity_get(validity, ni) &&
|
|
diskann_neighbor_id_get(neighborIds, ni) == candidate) {
|
|
alreadyNeighbor = 1;
|
|
break;
|
|
}
|
|
}
|
|
if (alreadyNeighbor) continue;
|
|
|
|
// Load, quantize, and set
|
|
void *candidateVec = NULL;
|
|
int cvs;
|
|
rc = diskann_vector_read(p, vec_col_idx, candidate, &candidateVec, &cvs);
|
|
if (rc != SQLITE_OK) continue;
|
|
|
|
size_t qvecSize = diskann_quantized_vector_byte_size(
|
|
cfg->quantizer_type, col->dimensions);
|
|
u8 *qvec = sqlite3_malloc(qvecSize);
|
|
if (qvec) {
|
|
if (col->element_type == SQLITE_VEC_ELEMENT_TYPE_FLOAT32) {
|
|
diskann_quantize_vector((const f32 *)candidateVec, col->dimensions,
|
|
cfg->quantizer_type, qvec);
|
|
} else {
|
|
memcpy(qvec, candidateVec,
|
|
qvecSize < (size_t)cvs ? qvecSize : (size_t)cvs);
|
|
}
|
|
diskann_node_set_neighbor(validity, neighborIds, qvecs, clearedSlot,
|
|
candidate, qvec,
|
|
cfg->quantizer_type, col->dimensions);
|
|
sqlite3_free(qvec);
|
|
}
|
|
sqlite3_free(candidateVec);
|
|
break;
|
|
}
|
|
|
|
diskann_node_write(p, vec_col_idx, nodeRowid,
|
|
validity, vs, neighborIds, nis, qvecs, qs);
|
|
}
|
|
|
|
sqlite3_free(validity);
|
|
sqlite3_free(neighborIds);
|
|
sqlite3_free(qvecs);
|
|
}
|
|
|
|
return SQLITE_OK;
|
|
}
|
|
|
|
/**
|
|
* Delete a vector from the DiskANN graph (Algorithm 3: LM-Delete).
|
|
* If the vector is in the buffer (not yet flushed), just remove from buffer.
|
|
*/
|
|
static int diskann_delete(vec0_vtab *p, int vec_col_idx, i64 rowid) {
|
|
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
|
|
struct Vec0DiskannConfig *cfg = &col->diskann;
|
|
int rc;
|
|
|
|
// Check if this rowid is in the buffer (not yet in graph)
|
|
if (cfg->buffer_threshold > 0) {
|
|
int inBuffer = 0;
|
|
rc = diskann_buffer_exists(p, vec_col_idx, rowid, &inBuffer);
|
|
if (rc != SQLITE_OK) return rc;
|
|
if (inBuffer) {
|
|
// Just remove from buffer and _vectors, no graph repair needed
|
|
rc = diskann_buffer_delete(p, vec_col_idx, rowid);
|
|
if (rc == SQLITE_OK) {
|
|
rc = diskann_vector_delete(p, vec_col_idx, rowid);
|
|
}
|
|
return rc;
|
|
}
|
|
}
|
|
|
|
// 1. Read the node to get its neighbor list
|
|
u8 *delValidity = NULL, *delNeighborIds = NULL, *delQvecs = NULL;
|
|
int dvs, dnis, dqs;
|
|
rc = diskann_node_read(p, vec_col_idx, rowid,
|
|
&delValidity, &dvs, &delNeighborIds, &dnis,
|
|
&delQvecs, &dqs);
|
|
if (rc != SQLITE_OK) {
|
|
return SQLITE_OK; // Node doesn't exist, nothing to do
|
|
}
|
|
|
|
i64 *deletedNeighbors = sqlite3_malloc(cfg->n_neighbors * sizeof(i64));
|
|
int deletedNeighborCount = 0;
|
|
if (!deletedNeighbors) {
|
|
sqlite3_free(delValidity);
|
|
sqlite3_free(delNeighborIds);
|
|
sqlite3_free(delQvecs);
|
|
return SQLITE_NOMEM;
|
|
}
|
|
|
|
for (int i = 0; i < cfg->n_neighbors; i++) {
|
|
if (diskann_validity_get(delValidity, i)) {
|
|
deletedNeighbors[deletedNeighborCount++] =
|
|
diskann_neighbor_id_get(delNeighborIds, i);
|
|
}
|
|
}
|
|
|
|
sqlite3_free(delValidity);
|
|
sqlite3_free(delNeighborIds);
|
|
sqlite3_free(delQvecs);
|
|
|
|
// 2. Repair reverse edges
|
|
rc = diskann_repair_reverse_edges(p, vec_col_idx, rowid,
|
|
deletedNeighbors, deletedNeighborCount);
|
|
sqlite3_free(deletedNeighbors);
|
|
|
|
// 3. Delete node and vector
|
|
if (rc == SQLITE_OK) {
|
|
rc = diskann_node_delete(p, vec_col_idx, rowid);
|
|
}
|
|
if (rc == SQLITE_OK) {
|
|
rc = diskann_vector_delete(p, vec_col_idx, rowid);
|
|
}
|
|
|
|
// 4. Handle medoid deletion
|
|
if (rc == SQLITE_OK) {
|
|
rc = diskann_medoid_handle_delete(p, vec_col_idx, rowid);
|
|
}
|
|
|
|
return rc;
|
|
}
|
|
|
|
static int vec0_all_columns_diskann(vec0_vtab *p) {
|
|
for (int i = 0; i < p->numVectorColumns; i++) {
|
|
if (p->vector_columns[i].index_type != VEC0_INDEX_TYPE_DISKANN) return 0;
|
|
}
|
|
return p->numVectorColumns > 0;
|
|
}
|
|
|
|
// ============================================================================
|
|
// Command dispatch
|
|
// ============================================================================
|
|
|
|
static int diskann_handle_command(vec0_vtab *p, const char *command) {
|
|
int col_idx = -1;
|
|
for (int i = 0; i < p->numVectorColumns; i++) {
|
|
if (p->vector_columns[i].index_type == VEC0_INDEX_TYPE_DISKANN) { col_idx = i; break; }
|
|
}
|
|
if (col_idx < 0) return SQLITE_EMPTY;
|
|
|
|
struct Vec0DiskannConfig *cfg = &p->vector_columns[col_idx].diskann;
|
|
|
|
if (strncmp(command, "search_list_size_search=", 24) == 0) {
|
|
int val = atoi(command + 24);
|
|
if (val < 1) { vtab_set_error(&p->base, "search_list_size_search must be >= 1"); return SQLITE_ERROR; }
|
|
cfg->search_list_size_search = val;
|
|
return SQLITE_OK;
|
|
}
|
|
if (strncmp(command, "search_list_size_insert=", 24) == 0) {
|
|
int val = atoi(command + 24);
|
|
if (val < 1) { vtab_set_error(&p->base, "search_list_size_insert must be >= 1"); return SQLITE_ERROR; }
|
|
cfg->search_list_size_insert = val;
|
|
return SQLITE_OK;
|
|
}
|
|
if (strncmp(command, "search_list_size=", 17) == 0) {
|
|
int val = atoi(command + 17);
|
|
if (val < 1) { vtab_set_error(&p->base, "search_list_size must be >= 1"); return SQLITE_ERROR; }
|
|
cfg->search_list_size = val;
|
|
return SQLITE_OK;
|
|
}
|
|
return SQLITE_EMPTY;
|
|
}
|
|
|
|
#ifdef SQLITE_VEC_TEST
|
|
// Expose internal DiskANN data structures and functions for unit testing.
|
|
|
|
int _test_diskann_candidate_list_init(struct DiskannCandidateList *list, int capacity) {
|
|
return diskann_candidate_list_init(list, capacity);
|
|
}
|
|
void _test_diskann_candidate_list_free(struct DiskannCandidateList *list) {
|
|
diskann_candidate_list_free(list);
|
|
}
|
|
int _test_diskann_candidate_list_insert(struct DiskannCandidateList *list, long long rowid, float distance) {
|
|
return diskann_candidate_list_insert(list, (i64)rowid, (f32)distance);
|
|
}
|
|
int _test_diskann_candidate_list_next_unvisited(const struct DiskannCandidateList *list) {
|
|
return diskann_candidate_list_next_unvisited(list);
|
|
}
|
|
int _test_diskann_candidate_list_count(const struct DiskannCandidateList *list) {
|
|
return list->count;
|
|
}
|
|
long long _test_diskann_candidate_list_rowid(const struct DiskannCandidateList *list, int i) {
|
|
return (long long)list->items[i].rowid;
|
|
}
|
|
float _test_diskann_candidate_list_distance(const struct DiskannCandidateList *list, int i) {
|
|
return (float)list->items[i].distance;
|
|
}
|
|
void _test_diskann_candidate_list_set_visited(struct DiskannCandidateList *list, int i) {
|
|
list->items[i].visited = 1;
|
|
}
|
|
|
|
int _test_diskann_visited_set_init(struct DiskannVisitedSet *set, int capacity) {
|
|
return diskann_visited_set_init(set, capacity);
|
|
}
|
|
void _test_diskann_visited_set_free(struct DiskannVisitedSet *set) {
|
|
diskann_visited_set_free(set);
|
|
}
|
|
int _test_diskann_visited_set_contains(const struct DiskannVisitedSet *set, long long rowid) {
|
|
return diskann_visited_set_contains(set, (i64)rowid);
|
|
}
|
|
int _test_diskann_visited_set_insert(struct DiskannVisitedSet *set, long long rowid) {
|
|
return diskann_visited_set_insert(set, (i64)rowid);
|
|
}
|
|
#endif /* SQLITE_VEC_TEST */
|
|
|