Skip to content

Quantize int8 DiskANN columns as int8 vectors #344

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MayCXC wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Quantize int8 DiskANN columns as int8 vectors #344

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 65 additions & 65 deletions sqlite-vec-diskann.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
// ============================================================
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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);
Expand Down
64 changes: 63 additions & 1 deletion tests/test-diskann.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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("<q", ids, slot * 8)[0]
bits = bytearray(2)
for j, x in enumerate(vectors[neighbor]):
if x > 0:
bits[j // 8] |= 1 << (j % 8)
assert qvecs[slot * 2 : slot * 2 + 2] == bytes(bits)
checked += 1
assert checked > 0


# ======================================================================
# Delete edge cases
# ======================================================================
Expand Down