From f4b54c6aa46eb0123af3253c66372a4aa2f858b1 Mon Sep 17 00:00:00 2001 From: Aaron Paterson Date: Fri, 23 Oct 2026 08:14:33 -0500 Subject: [PATCH] Quantize int8 DiskANN columns as int8 vectors A DiskANN search quantized a query that was not float32 through diskann_distance_quantized(), which read it as floats, four times its bytes: every insert into an int8 column, whose graph insertion searches with the new vector, read past it. The neighbors of an int8 column were stored by copying the start of their vectors, so under the binary quantizer a neighbor's bits were its first bytes. diskann_quantize_element_vector() quantizes a vector of either element type: a float32 vector as before, and an int8 vector as itself under the int8 quantizer or as the bits of its elements above 0 under the binary quantizer, as a float32 vector is. The query and every neighbor go through it, and the wrapper that quantized on the fly is gone. --- sqlite-vec-diskann.c | 130 +++++++++++++++++++++--------------------- tests/test-diskann.py | 64 ++++++++++++++++++++- 2 files changed, 128 insertions(+), 66 deletions(-) diff --git a/sqlite-vec-diskann.c b/sqlite-vec-diskann.c index 5bd298bc..93f9db2e 100644 --- a/sqlite-vec-diskann.c +++ b/sqlite-vec-diskann.c @@ -165,6 +165,43 @@ int diskann_quantize_vector( return SQLITE_ERROR; } +/** + * Quantize a vector of a DiskANN column's element type into the target + * quantizer format. A float32 vector goes through diskann_quantize_vector(). + * An int8 vector is its own int8 quantized form, and under the binary + * quantizer each element above 0 sets its bit, as for a float32 vector. + * Output buffer must be pre-allocated with diskann_quantized_vector_byte_size() bytes. + */ +static int diskann_quantize_element_vector( + const void *src, enum VectorElementType element_type, size_t dimensions, + enum Vec0DiskannQuantizerType quantizer_type, u8 *out) { + switch (element_type) { + case SQLITE_VEC_ELEMENT_TYPE_FLOAT32: + return diskann_quantize_vector((const f32 *)src, dimensions, + quantizer_type, out); + case SQLITE_VEC_ELEMENT_TYPE_INT8: { + const i8 *v = (const i8 *)src; + switch (quantizer_type) { + case VEC0_DISKANN_QUANTIZER_BINARY: + memset(out, 0, dimensions / CHAR_BIT); + for (size_t i = 0; i < dimensions; i++) { + if (v[i] > 0) { + out[i / CHAR_BIT] |= (1 << (i % CHAR_BIT)); + } + } + return SQLITE_OK; + case VEC0_DISKANN_QUANTIZER_INT8: + memcpy(out, v, dimensions); + return SQLITE_OK; + } + break; + } + case SQLITE_VEC_ELEMENT_TYPE_BIT: + break; + } + return SQLITE_ERROR; +} + /** * Compute approximate distance between a full-precision query vector and a * quantized neighbor vector. Used during graph traversal. @@ -199,35 +236,20 @@ static f32 diskann_distance_quantized_precomputed( } /** - * Quantize a float query vector. Returns allocated buffer (caller must free). + * Quantize a query vector of the column's element type. Returns allocated + * buffer (caller must free), or NULL when out of memory. */ static u8 *diskann_quantize_query( - const f32 *query_vector, size_t dimensions, - enum Vec0DiskannQuantizerType quantizer_type) { + const void *query_vector, enum VectorElementType element_type, + 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); + diskann_quantize_element_vector(query_vector, element_type, 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 // ============================================================ @@ -747,10 +769,12 @@ static int diskann_search( 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); + u8 *queryQuantized = diskann_quantize_query(queryVector, elementType, + dimensions, cfg->quantizer_type); + if (!queryQuantized) { + diskann_candidate_list_free(&candidates); + diskann_visited_set_free(&visited); + return SQLITE_NOMEM; } // 4. Greedy beam search loop (Algorithm 1 from LM-DiskANN paper) @@ -784,16 +808,9 @@ static int diskann_search( 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); - } + f32 approxDist = diskann_distance_quantized_precomputed( + queryQuantized, neighborQvec, dimensions, + cfg->quantizer_type, col->distance_metric); diskann_candidate_list_insert(&candidates, neighborRowid, approxDist); } @@ -1057,13 +1074,8 @@ static int diskann_write_pruned_neighbors( &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_quantize_element_vector(neighborVec, col->element_type, + col->dimensions, cfg->quantizer_type, qvec); diskann_node_set_neighbor(validity, neighborIds, qvecs, i, neighborRowids[i], qvec, @@ -1135,13 +1147,9 @@ static int diskann_add_reverse_edge( 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_quantize_element_vector(target_vector, col->element_type, + col->dimensions, cfg->quantizer_type, + qvec); diskann_node_set_neighbor(validity, neighborIds, qvecs, i, target_rowid, qvec, @@ -1187,15 +1195,11 @@ static int diskann_add_reverse_edge( 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); - } + diskann_quantize_element_vector(target_vector, col->element_type, + col->dimensions, cfg->quantizer_type, + targetQ); + diskann_quantize_element_vector(nodeVector, col->element_type, + col->dimensions, cfg->quantizer_type, nodeQ); // Compute quantized distance from node to target f32 targetDist = diskann_distance_quantized_precomputed( @@ -1605,13 +1609,9 @@ static int diskann_repair_reverse_edges( 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_quantize_element_vector(candidateVec, col->element_type, + col->dimensions, cfg->quantizer_type, + qvec); diskann_node_set_neighbor(validity, neighborIds, qvecs, clearedSlot, candidate, qvec, cfg->quantizer_type, col->dimensions); diff --git a/tests/test-diskann.py b/tests/test-diskann.py index 4e651604..7bcbbdcc 100644 --- a/tests/test-diskann.py +++ b/tests/test-diskann.py @@ -1,7 +1,7 @@ import sqlite3 import struct import pytest -from helpers import _f32, exec +from helpers import _f32, _int8, exec def test_diskann_create_basic(db): @@ -779,6 +779,68 @@ def test_diskann_insert_int8_quantizer_knn(db): assert distances[i] <= distances[i + 1] +def test_diskann_int8_column_knn_matches_exact(db): + """An int8 column under the int8 quantizer: a search whose list holds every + row returns the exact nearest rows.""" + db.execute(""" + CREATE VIRTUAL TABLE t USING vec0( + emb int8[64] INDEXED BY diskann(neighbor_quantizer=int8, search_list_size=128) + ) + """) + import random + random.seed(9) + for i in range(1, 101): + vec = [random.randint(-128, 127) for _ in range(64)] + db.execute("INSERT INTO t(rowid, emb) VALUES (?, vec_int8(?))", [i, _int8(vec)]) + + for _ in range(10): + query = _int8([random.randint(-128, 127) for _ in range(64)]) + rows = db.execute( + "SELECT rowid FROM t WHERE emb MATCH vec_int8(?) AND k = 10", [query] + ).fetchall() + exact = db.execute( + "SELECT rowid FROM t ORDER BY vec_distance_l2(emb, vec_int8(?)) LIMIT 10", + [query], + ).fetchall() + assert [r[0] for r in rows] == [r[0] for r in exact] + + +def test_diskann_int8_column_binary_neighbors_are_sign_bits(db): + """An int8 column under the binary quantizer keeps each neighbor's sign bits: + an element above 0 sets its bit.""" + db.execute(""" + CREATE VIRTUAL TABLE t USING vec0( + emb int8[16] INDEXED BY diskann(neighbor_quantizer=binary, n_neighbors=8) + ) + """) + import random + random.seed(5) + vectors = {} + for i in range(1, 13): + vectors[i] = [random.randint(-128, 127) for _ in range(16)] + db.execute( + "INSERT INTO t(rowid, emb) VALUES (?, vec_int8(?))", + [i, _int8(vectors[i])], + ) + + checked = 0 + for validity, ids, qvecs in db.execute( + "SELECT neighbors_validity, neighbor_ids, neighbor_quantized_vectors " + "FROM t_diskann_nodes00" + ): + for slot in range(8): + if not (validity[slot // 8] >> (slot % 8)) & 1: + continue + neighbor = struct.unpack_from(" 0: + bits[j // 8] |= 1 << (j % 8) + assert qvecs[slot * 2 : slot * 2 + 2] == bytes(bits) + checked += 1 + assert checked > 0 + + # ====================================================================== # Delete edge cases # ======================================================================