diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 4f62e16a..dff4d059 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -66,6 +66,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 any +hidden parameter column declared after `k`. 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. diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..372de41a 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -6009,6 +6009,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. 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; /** @@ -6035,6 +6052,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; diff --git a/tests/test-knn-joined-constraints.py b/tests/test-knn-joined-constraints.py new file mode 100644 index 00000000..ba774d70 --- /dev/null +++ b/tests/test-knn-joined-constraints.py @@ -0,0 +1,78 @@ +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 and 4 and 5 among tag 2. + + +@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)") + db.execute("insert into w values ('[0,0]', 2)") + db.execute("create table categories(id integer primary key, q text, n integer)") + db.execute("insert into categories values (1, '[0,0]', 2), (2, '[0,0]', 2)") + 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] + + +# A table that supplies the query vector or k 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", + ], +) +@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)]