From 743511af5571061bdc0e3dfb68090f5d69395dbf Mon Sep 17 00:00:00 2001
From: Alex Garcia <alexsebastian.garcia@gmail.com>
Date: Wed, 13 Nov 2024 15:27:23 -0800
Subject: [PATCH] block WHERE constraints on auxiliary columns in KNN queries

---
 sqlite-vec.c                            | 13 ++++++
 tests/__snapshots__/test-auxiliary.ambr | 53 +++++++++++++++++++++++++
 tests/test-auxiliary.py                 | 16 +++++++-
 3 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/sqlite-vec.c b/sqlite-vec.c
index 1acf4d2..1e38bcf 100644
--- a/sqlite-vec.c
+++ b/sqlite-vec.c
@@ -4923,6 +4923,7 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
   int iRowidTerm = -1;
   int iKTerm = -1;
   int iRowidInTerm = -1;
+  int hasAuxConstraint = 0;
 
 #ifdef SQLITE_VEC_DEBUG
   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)) {
       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);
@@ -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);
 
     int argvIndex = 1;
diff --git a/tests/__snapshots__/test-auxiliary.ambr b/tests/__snapshots__/test-auxiliary.ambr
index 8cba12a..33dd836 100644
--- a/tests/__snapshots__/test-auxiliary.ambr
+++ b/tests/__snapshots__/test-auxiliary.ambr
@@ -5,6 +5,59 @@
     '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
   OrderedDict({
     'sql': 'select * from v',
diff --git a/tests/test-auxiliary.py b/tests/test-auxiliary.py
index f09de96..71d8a49 100644
--- a/tests/test-auxiliary.py
+++ b/tests/test-auxiliary.py
@@ -81,7 +81,21 @@ def test_deletes(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=[]):