// 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; } rc = diskann_node_write(p, vec_col_idx, nodeRowid, validity, vs, neighborIds, nis, qvecs, qs); } sqlite3_free(validity); sqlite3_free(neighborIds); sqlite3_free(qvecs); if (rc != SQLITE_OK) return rc; } 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. */ /** * Scan all nodes and clear any neighbor slot referencing deleted_rowid. * This removes stale reverse edges that the forward-edge repair misses, * preventing data leaks (deleted rowid + quantized vector lingering in * other nodes' blobs). */ static int diskann_scrub_deleted_rowid( vec0_vtab *p, int vec_col_idx, i64 deleted_rowid) { struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx]; struct Vec0DiskannConfig *cfg = &col->diskann; int rc; sqlite3_stmt *stmt = NULL; // Lightweight scan: only read validity + neighbor_ids to find matches char *zSql = sqlite3_mprintf( "SELECT rowid, neighbors_validity, neighbor_ids " "FROM " VEC0_SHADOW_DISKANN_NODES_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; // Collect rowids that need updating (avoid modifying while iterating) i64 *dirty = NULL; int nDirty = 0, capDirty = 0; while (sqlite3_step(stmt) == SQLITE_ROW) { const u8 *validity = (const u8 *)sqlite3_column_blob(stmt, 1); const u8 *ids = (const u8 *)sqlite3_column_blob(stmt, 2); int idsBytes = sqlite3_column_bytes(stmt, 2); if (!validity || !ids) continue; int nSlots = idsBytes / (int)sizeof(i64); if (nSlots > cfg->n_neighbors) nSlots = cfg->n_neighbors; for (int i = 0; i < nSlots; i++) { if (!diskann_validity_get(validity, i)) continue; i64 nid = diskann_neighbor_id_get(ids, i); if (nid == deleted_rowid) { i64 nodeRowid = sqlite3_column_int64(stmt, 0); // Add to dirty list if (nDirty >= capDirty) { capDirty = capDirty ? capDirty * 2 : 16; i64 *tmp = sqlite3_realloc64(dirty, capDirty * sizeof(i64)); if (!tmp) { sqlite3_free(dirty); sqlite3_finalize(stmt); return SQLITE_NOMEM; } dirty = tmp; } dirty[nDirty++] = nodeRowid; break; // one match per node is enough } } } sqlite3_finalize(stmt); // Now do full read/clear/write for each dirty node for (int d = 0; d < nDirty; d++) { u8 *val = NULL, *nids = NULL, *qvecs = NULL; int vs, nis, qs; rc = diskann_node_read(p, vec_col_idx, dirty[d], &val, &vs, &nids, &nis, &qvecs, &qs); if (rc != SQLITE_OK) continue; int modified = 0; for (int i = 0; i < cfg->n_neighbors; i++) { if (diskann_validity_get(val, i) && diskann_neighbor_id_get(nids, i) == deleted_rowid) { diskann_node_clear_neighbor(val, nids, qvecs, i, cfg->quantizer_type, col->dimensions); modified = 1; } } if (modified) { rc = diskann_node_write(p, vec_col_idx, dirty[d], val, vs, nids, nis, qvecs, qs); } sqlite3_free(val); sqlite3_free(nids); sqlite3_free(qvecs); if (rc != SQLITE_OK) break; } sqlite3_free(dirty); return rc; } 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); } // 5. Scrub stale reverse edges — removes deleted rowid + quantized vector // from any node that still references it (data leak prevention) if (rc == SQLITE_OK) { rc = diskann_scrub_deleted_rowid(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 */