apunkt/sqlite-vec
1
0
Fork 0
mirror of https://github.com/asg017/sqlite-vec.git synced 2026-04-25 16:56:27 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

block WHERE constraints on auxiliary columns in KNN queries

Browse source
This commit is contained in:
Alex Garcia 2024-11-13 15:27:23 -08:00
parent 7d4c023928
commit 743511af55
3 changed files with 81 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

13
sqlite-vec.c
View file

@ -4923,6 +4923,7 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
int iRowidTerm = -1; int iRowidTerm = -1;
int iKTerm = -1; int iKTerm = -1;
int iRowidInTerm = -1; int iRowidInTerm = -1;
int hasAuxConstraint = 0;
#ifdef SQLITE_VEC_DEBUG #ifdef SQLITE_VEC_DEBUG
printf("pIdxInfo->nOrderBy=%d, pIdxInfo->nConstraint=%d\n", pIdxInfo->nOrderBy, pIdxInfo->nConstraint); printf("pIdxInfo->nOrderBy=%d, pIdxInfo->nConstraint=%d\n", pIdxInfo->nOrderBy, pIdxInfo->nConstraint);
@ -4977,6 +4978,11 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
if (op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn == vec0_column_k_idx(p)) { if (op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn == vec0_column_k_idx(p)) {
iKTerm = i; iKTerm = i;
} }
if(
(op != SQLITE_INDEX_CONSTRAINT_LIMIT && op != SQLITE_INDEX_CONSTRAINT_OFFSET)
&& vec0_column_idx_is_auxiliary(p, iColumn)) {
hasAuxConstraint = 1;
}
} }
sqlite3_str *idxStr = sqlite3_str_new(NULL); sqlite3_str *idxStr = sqlite3_str_new(NULL);
@ -5019,6 +5025,13 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
} }
} }
if(hasAuxConstraint) {
// IMP: V25623_09693
vtab_set_error(pVTab, "An illegal WHERE constraint was provided on a vec0 auxiliary column in a KNN query.");
rc = SQLITE_ERROR;
goto done;
}
sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_KNN); sqlite3_str_appendchar(idxStr, 1, VEC0_QUERY_PLAN_KNN);
int argvIndex = 1; int argvIndex = 1;

53
tests/__snapshots__/test-auxiliary.ambr
View file

@ -5,6 +5,59 @@
'message': 'vec0 constructor error: More than 16 auxiliary columns were provided', 'message': 'vec0 constructor error: More than 16 auxiliary columns were provided',
}) })
# --- # ---
# name: test_knn
OrderedDict({
'sql': 'select * from v',
'rows': list([
OrderedDict({
'rowid': 1,
'vector': b'\x00\x00\x80?',
'name': 'alex',
}),
OrderedDict({
'rowid': 2,
'vector': b'\x00\x00\x00@',
'name': 'brian',
}),
OrderedDict({
'rowid': 3,
'vector': b'\x00\x00@@',
'name': 'craig',
}),
]),
})
# ---
# name: test_knn[illegal KNN w/ aux]
dict({
'error': 'OperationalError',
'message': 'An illegal WHERE constraint was provided on a vec0 auxiliary column in a KNN query.',
})
# ---
# name: test_knn[legal KNN w/ aux]
OrderedDict({
'sql': "select *, distance from v where vector match '[5]' and k = 10",
'rows': list([
OrderedDict({
'rowid': 3,
'vector': b'\x00\x00@@',
'name': 'craig',
'distance': 2.0,
}),
OrderedDict({
'rowid': 2,
'vector': b'\x00\x00\x00@',
'name': 'brian',
'distance': 3.0,
}),
OrderedDict({
'rowid': 1,
'vector': b'\x00\x00\x80?',
'name': 'alex',
'distance': 4.0,
}),
]),
})
# ---
# name: test_normal.1 # name: test_normal.1
OrderedDict({ OrderedDict({
'sql': 'select * from v', 'sql': 'select * from v',

16
tests/test-auxiliary.py
View file

@ -81,7 +81,21 @@ def test_deletes(db, snapshot):
def test_knn(db, snapshot): def test_knn(db, snapshot):
pass db.execute("create virtual table v using vec0(vector float[1], +name text)")
db.executemany(
"insert into v(vector, name) values (?, ?)",
[("[1]", "alex"), ("[2]", "brian"), ("[3]", "craig")],
)
assert exec(db, "select * from v") == snapshot()
assert exec(
db, "select *, distance from v where vector match '[5]' and k = 10"
) == snapshot(name="legal KNN w/ aux")
# EVIDENCE-OF: V25623_09693 No aux constraint allowed on KNN queries
assert exec(
db,
"select *, distance from v where vector match '[5]' and k = 10 and name = 'alex'",
) == snapshot(name="illegal KNN w/ aux")
def exec(db, sql, parameters=[]): def exec(db, sql, parameters=[]):