Take a KNN query's vector, k and mmr_lambda from a joined table

ARCHITECTURE.md

@@ -135,6 +135,33 @@ is associated with `argv[1]`. The 2nd block at byte offset `5-8` (inclusive) is
 associated with `argv[2]` and so on. Each block describes what kind of value or
 filter the given `argv[i]` value is.
 
+A KNN plan needs its arguments: the MATCH on a vector column, `k`, and the
+hidden parameter columns declared after `k` (`mmr_lambda` among them). When one
+of them reads another table of the query, SQLite also weighs plans that scan
+`vec0` before that table, in which the argument is unusable. `xBestIndex`
+returns `SQLITE_CONSTRAINT` for such a plan (`vec0_constraint_is_knn_argument`),
+so the planner takes one that gives `vec0` every argument.
+
+Filters on rowid, partition key, metadata and distance columns are left to the
+plan. A filter whose value comes from another table is unusable in a plan that
+scans `vec0` first, and SQLite takes such a plan whenever the query allows it,
+in either FROM order, applying the filter to the k rows `vec0` returns. That
+keeps a lookup join at k rows, and `vec0` cannot tell a lookup join from a
+filter join: in
+
+```sql
+select v.rowid, c.name from v join categories c on c.id = v.category
+where v.embedding match ? and k = 10
+```
+
+refusing the plan without `v.category = c.id` would run one KNN per
+`categories` row and return 10 rows per category. A query that reads the other
+table first, because that table supplies a KNN argument, a `CROSS JOIN` names it
+first, or `vec0` is on the right of a `LEFT JOIN` with the filter in `ON`, has
+the filter usable in every plan, so `vec0` runs once per row of that table and
+applies the filter before the k. So does a correlated subquery, whose outer
+columns are constants to `vec0`.
+
 #### `VEC0_IDXSTR_KIND_KNN_MATCH` (`'{'`)
 
 `argv[i]` is the query vector of the KNN query.

sqlite-vec.c

@@ -5744,6 +5744,23 @@ typedef enum {
   VEC0_DISTANCE_CONSTRAINT_LE = 'd',
 } vec0_distance_constraint_operator;
 
+/**
+ * Whether constraint i is an argument of a KNN query: a MATCH on a vector
+ * column, or `k` or a hidden parameter column declared after it (`mmr_lambda`
+ * among them). Filters on rowid, partition key, metadata and distance columns
+ * are not; ARCHITECTURE.md's idxStr section says why they are left to the plan.
+ */
+static int vec0_constraint_is_knn_argument(vec0_vtab *p,
+                                           sqlite3_index_info *pIdxInfo,
+                                           int i) {
+  int iColumn = pIdxInfo->aConstraint[i].iColumn;
+  int op = pIdxInfo->aConstraint[i].op;
+  if (op == SQLITE_INDEX_CONSTRAINT_MATCH) {
+    return vec0_column_idx_is_vector(p, iColumn);
+  }
+  return op == SQLITE_INDEX_CONSTRAINT_EQ && iColumn >= vec0_column_k_idx(p);
+}
+
 static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
   vec0_vtab *p = (vec0_vtab *)pVTab;
   /**
@@ -5771,6 +5788,23 @@ static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) {
   printf("pIdxInfo->nOrderBy=%d, pIdxInfo->nConstraint=%d\n", pIdxInfo->nOrderBy, pIdxInfo->nConstraint);
 #endif
 
+  // When a KNN argument reads another table of the query, SQLite also weighs
+  // plans that scan vec0 before that table, in which the argument is unusable.
+  // Such a plan is refused, the way json_each refuses a plan that cannot supply
+  // its argument, so the planner takes one that gives vec0 every argument.
+  for (int i = 0; i < pIdxInfo->nConstraint; i++) {
+    if (pIdxInfo->aConstraint[i].op == SQLITE_INDEX_CONSTRAINT_MATCH &&
+        vec0_column_idx_is_vector(p, pIdxInfo->aConstraint[i].iColumn)) {
+      for (int j = 0; j < pIdxInfo->nConstraint; j++) {
+        if (!pIdxInfo->aConstraint[j].usable &&
+            vec0_constraint_is_knn_argument(p, pIdxInfo, j)) {
+          return SQLITE_CONSTRAINT;
+        }
+      }
+      break;
+    }
+  }
+
   for (int i = 0; i < pIdxInfo->nConstraint; i++) {
     u8 vtabIn = 0;
 

tests/test-knn-joined-constraints.py

@@ -0,0 +1,92 @@
+import pytest
+
+# Six points on a line, the first three tagged 1 and in partition 1, the rest
+# tagged 2 and in partition 2, so the k = 2 nearest the origin are rowids 1 and
+# 2 overall, 4 and 5 among tag 2, and 1 and 6 with mmr_lambda = 0.3.
+
+
+@pytest.fixture()
+def fill(db):
+    db.execute(
+        "create virtual table v using vec0("
+        "part integer partition key, embedding float[2] distance_metric=l2, tag integer)"
+    )
+    db.executemany(
+        "insert into v(rowid, part, embedding, tag) values (?, ?, ?, ?)",
+        [[i, 1 if i <= 3 else 2, f"[{i},0]", 1 if i <= 3 else 2] for i in range(1, 7)],
+    )
+    db.execute("create table w(q text, n integer, lambda real)")
+    db.execute("insert into w values ('[0,0]', 2, 0.3)")
+    db.execute(
+        "create table categories(id integer primary key, q text, n integer, lambda real)"
+    )
+    db.execute(
+        "insert into categories values (1, '[0,0]', 2, 1.0), (2, '[0,0]', 2, 1.0)"
+    )
+    return db
+
+
+def rowids(db, sql):
+    return [row[0] for row in db.execute(sql)]
+
+
+@pytest.mark.parametrize("tables", ["w, v", "v, w"])
+def test_the_query_vector_can_come_from_another_table(fill, tables):
+    knn = f"select v.rowid from {tables} where v.embedding match w.q and v.k = 2"
+    assert rowids(fill, knn) == [1, 2]
+
+
+@pytest.mark.parametrize("tables", ["w, v", "v, w"])
+def test_k_can_come_from_another_table(fill, tables):
+    knn = f"select v.rowid from {tables} where v.embedding match '[0,0]' and v.k = w.n"
+    assert rowids(fill, knn) == [1, 2]
+
+
+@pytest.mark.parametrize("tables", ["w, v", "v, w"])
+def test_mmr_lambda_can_come_from_another_table(fill, tables):
+    knn = (
+        f"select v.rowid from {tables} where v.embedding match '[0,0]' and v.k = 2 "
+        "and v.mmr_lambda = w.lambda"
+    )
+    assert rowids(fill, knn) == [1, 6]
+
+
+# A table that supplies the query vector, k or mmr_lambda is read before v, so
+# vec0 runs one KNN per row of it, and a filter whose value comes from that
+# table applies before the k: the 2 nearest in each category.
+@pytest.mark.parametrize(
+    "arguments",
+    [
+        "v.embedding match c.q and v.k = 2",
+        "v.embedding match '[0,0]' and v.k = c.n",
+        "v.embedding match '[0,0]' and v.k = 2 and v.mmr_lambda = c.lambda",
+    ],
+)
+@pytest.mark.parametrize("column", ["tag", "part"])
+def test_a_table_that_supplies_an_argument_filters_before_the_k(fill, column, arguments):
+    knn = (
+        f"select c.id, v.rowid from v join categories c on c.id = v.{column} "
+        f"where {arguments}"
+    )
+    assert [tuple(row) for row in fill.execute(knn)] == [(1, 1), (1, 2), (2, 4), (2, 5)]
+
+
+# A filter whose value comes from a table that supplies no KNN argument runs on
+# the k rows the KNN returns, in either FROM order, so a lookup join returns k
+# rows. Running one KNN per categories row instead would return k rows per
+# category; ARCHITECTURE.md's idxStr section says why vec0 leaves the filter to
+# the plan.
+@pytest.mark.parametrize(
+    "tables",
+    [
+        "v join categories c on c.id = v.{column}",
+        "categories c join v on c.id = v.{column}",
+    ],
+)
+@pytest.mark.parametrize("column", ["tag", "part"])
+def test_a_lookup_join_on_a_filter_column_returns_k_rows(fill, column, tables):
+    knn = (
+        f"select c.id, v.rowid from {tables.format(column=column)} "
+        "where v.embedding match '[0,0]' and v.k = 2"
+    )
+    assert [tuple(row) for row in fill.execute(knn)] == [(1, 1), (1, 2)]