Skip to content

Apply rowid filters to DiskANN KNN queries #348

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

Apply rowid filters to DiskANN KNN queries #348

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
89 changes: 86 additions & 3 deletions sqlite-vec.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7645,6 +7645,74 @@ int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks,
#endif

#if SQLITE_VEC_ENABLE_DISKANN
#if COMPILER_SUPPORTS_VTAB_IN
/**
* Rank the rows a `rowid in (...)` constraint names by their exact distances
* to the query, keeping the k nearest in ascending order of distance. A rowid
* or id the table does not have is skipped. The vectors are read from
* _vectorsNN, which holds the buffered rows as well as the graph's.
*/
static int diskann_knn_rowids_in(vec0_vtab *p, int vec_col_idx,
sqlite3_value *rowidsIn,
const void *queryVector, size_t dimensions,
enum VectorElementType elementType, i64 k,
i64 *outRowids, f32 *outDistances,
int *outCount) {
struct VectorColumnDefinition *col = &p->vector_columns[vec_col_idx];
int vectorSize = (int)vector_column_byte_size(*col);
sqlite3_stmt *stmt = NULL;
sqlite3_value *item;
int count = 0;
int rc;
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, &stmt, NULL);
sqlite3_free(zSql);
if (rc != SQLITE_OK) return rc;

for (rc = sqlite3_vtab_in_first(rowidsIn, &item); rc == SQLITE_OK && item;
rc = sqlite3_vtab_in_next(rowidsIn, &item)) {
i64 rowid;
if (p->pkIsText) {
rc = vec0_rowid_from_id(p, item, &rowid);
if (rc == SQLITE_EMPTY) continue;
if (rc != SQLITE_OK) break;
} else {
rowid = sqlite3_value_int64(item);
}

sqlite3_reset(stmt);
sqlite3_bind_int64(stmt, 1, rowid);
rc = sqlite3_step(stmt);
if (rc == SQLITE_DONE) continue;
if (rc != SQLITE_ROW) break;
if (sqlite3_column_bytes(stmt, 0) != vectorSize) {
rc = SQLITE_CORRUPT;
break;
}
f32 distance = vec0_distance_full(queryVector, sqlite3_column_blob(stmt, 0),
dimensions, elementType,
col->distance_metric);

if (count == k && distance >= outDistances[count - 1]) continue;
int i = count < k ? count++ : count - 1;
while (i > 0 && outDistances[i - 1] > distance) {
outRowids[i] = outRowids[i - 1];
outDistances[i] = outDistances[i - 1];
i--;
}
outRowids[i] = rowid;
outDistances[i] = distance;
}
sqlite3_finalize(stmt);
if (rc != SQLITE_DONE) return rc;
*outCount = count;
return SQLITE_OK;
}
#endif

/**
* Handle a KNN query using the DiskANN graph search.
*/
Expand All @@ -7661,16 +7729,20 @@ static int vec0Filter_knn_diskann(
if (!knn_data) return SQLITE_NOMEM;
memset(knn_data, 0, sizeof(*knn_data));

// Parse query_idx and k_idx from idxStr
// Parse query_idx, k_idx and rowid_in_idx from idxStr
int query_idx = -1;
int k_idx = -1;
int rowid_in_idx = -1;
for (int i = 0; i < argc; i++) {
if (idxStr[1 + (i * 4)] == VEC0_IDXSTR_KIND_KNN_MATCH) {
query_idx = i;
}
if (idxStr[1 + (i * 4)] == VEC0_IDXSTR_KIND_KNN_K) {
k_idx = i;
}
if (idxStr[1 + (i * 4)] == VEC0_IDXSTR_KIND_KNN_ROWID_IN) {
rowid_in_idx = i;
}
}
assert(query_idx >= 0);
assert(k_idx >= 0);
Expand Down Expand Up @@ -7719,7 +7791,17 @@ static int vec0Filter_knn_diskann(
return SQLITE_NOMEM;
}

int resultCount;
int resultCount = 0;
#if COMPILER_SUPPORTS_VTAB_IN
// A rowid filter names its candidates, so each is ranked by its exact
// distance; the graph search reaches only the rows near the query, which
// the filter may not name.
if (rowid_in_idx >= 0) {
rc = diskann_knn_rowids_in(p, vectorColumnIdx, argv[rowid_in_idx],
queryVector, dimensions, elementType, k,
resultRowids, resultDistances, &resultCount);
} else
#endif
rc = diskann_search(p, vectorColumnIdx, queryVector, dimensions,
elementType, (int)k, 0,
resultRowids, resultDistances, &resultCount);
Expand All @@ -7734,7 +7816,8 @@ static int vec0Filter_knn_diskann(

// Scan _diskann_buffer for any buffered (unflushed) vectors and merge
// with graph results. This ensures no recall loss for buffered vectors.
{
// A rowid filter has already ranked the buffered rows it names.
if (rowid_in_idx < 0) {
sqlite3_stmt *bufStmt = NULL;
char *zSql = sqlite3_mprintf(
"SELECT rowid, vector FROM " VEC0_SHADOW_DISKANN_BUFFER_N_NAME,
Expand Down
54 changes: 54 additions & 0 deletions tests/test-diskann.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1329,3 +1329,57 @@ def test_diskann_delete_scrubs_all_references(db):
f"Node {node_rowid} slot {j // 8} still references "
f"deleted rowid {target}"
)


# ======================================================================
# KNN with a rowid filter
# ======================================================================


def _rowid_filter_rows(db, table_options):
db.execute(f"CREATE VIRTUAL TABLE t USING vec0(emb float[8]{table_options})")
db.execute("CREATE VIRTUAL TABLE flat USING vec0(emb float[8])")
for i in range(1, 21):
vec = _f32([((i * 7 + j * 3) % 11) / 10 - 0.5 for j in range(8)])
db.execute("INSERT INTO t(rowid, emb) VALUES (?, ?)", [i, vec])
db.execute("INSERT INTO flat(rowid, emb) VALUES (?, ?)", [i, vec])
return db.execute("SELECT emb FROM flat WHERE rowid = 3").fetchone()[0]


@pytest.mark.parametrize(
"table_options",
[
" INDEXED BY diskann(neighbor_quantizer=binary, n_neighbors=8)",
" INDEXED BY diskann(neighbor_quantizer=binary, n_neighbors=8, buffer_threshold=16)",
],
ids=["graph", "buffered"],
)
def test_diskann_knn_rowid_in(db, table_options):
"""A rowid filter returns the k nearest of the rows it names, as a flat
column does, whether those rows are in the graph or the buffer."""
query = _rowid_filter_rows(db, table_options)
for rowids in ["1, 2", "2, 19, 20", "4, 5, 6, 7, 8, 9", "3, 3, 999"]:
sql = f"SELECT rowid, distance FROM {{}} WHERE emb MATCH ? AND k = 3 AND rowid IN ({rowids})"
got = [tuple(r) for r in db.execute(sql.format("t"), [query])]
want = [tuple(r) for r in db.execute(sql.format("flat"), [query])]
assert got == want


def test_diskann_knn_id_in_text_primary_key(db):
"""A filter on a text primary key returns the nearest of the rows it names
and skips an id the table does not have."""
db.execute("""
CREATE VIRTUAL TABLE t USING vec0(
id text primary key,
emb float[8] INDEXED BY diskann(neighbor_quantizer=binary, n_neighbors=8)
)
""")
for i in range(1, 21):
vec = _f32([((i * 7 + j * 3) % 11) / 10 - 0.5 for j in range(8)])
db.execute("INSERT INTO t(id, emb) VALUES (?, ?)", [f"row{i}", vec])
query = db.execute("SELECT emb FROM t WHERE id = 'row3'").fetchone()[0]
rows = db.execute(
"SELECT id FROM t WHERE emb MATCH ? AND k = 3 AND id IN ('row1', 'row2', 'missing')",
[query],
).fetchall()
assert [r[0] for r in rows] == ["row1", "row2"]