diff --git a/sqlite-vec-rescore.c b/sqlite-vec-rescore.c index 6a472144..0bf649ab 100644 --- a/sqlite-vec-rescore.c +++ b/sqlite-vec-rescore.c @@ -602,6 +602,36 @@ static int rescore_knn(vec0_vtab *p, vec0_cursor *pCur, } } + // Apply any `distance` constraints from the WHERE clause. This must happen + // on the rescored float distances (the coarse quantized distances from + // phase 1 are not comparable to a user-supplied threshold), and before the + // top-k truncation below, so that a constraint like `distance > x` can + // still yield k rows instead of dropping the head of the result set. + i64 cand_kept = 0; + for (i64 j = 0; j < cand_used; j++) { + if (!vec0_distance_constraints_satisfied(float_distances[j], idxStr, argc, + argv)) { + continue; + } + cand_rowids[cand_kept] = cand_rowids[j]; + float_distances[cand_kept] = float_distances[j]; + cand_kept++; + } + cand_used = cand_kept; + + // Every candidate was filtered out: return an empty result set. Falling + // through would sqlite3_malloc(0), which returns NULL and would be + // misreported as SQLITE_NOMEM below. + if (cand_used == 0) { + knn_data->current_idx = 0; + knn_data->k = 0; + knn_data->rowids = NULL; + knn_data->distances = NULL; + knn_data->k_used = 0; + sqlite3_free(float_distances); + goto cleanup; + } + i64 result_k = min(k, cand_used); i64 *out_rowids = sqlite3_malloc(result_k * sizeof(i64)); f32 *out_distances = sqlite3_malloc(result_k * sizeof(f32)); diff --git a/sqlite-vec.c b/sqlite-vec.c index 7af3b6a7..a8eeea6a 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -6009,6 +6009,53 @@ typedef enum { VEC0_DISTANCE_CONSTRAINT_LE = 'd', } vec0_distance_constraint_operator; +/** + * @brief Test a single distance value against every distance constraint + * encoded in idxStr. + * + * vec0BestIndex() sets aConstraintUsage[].omit = 1 on `distance` constraints, + * which promises SQLite that the vtab applies them itself. Every KNN + * implementation (FLAT, rescore, IVF, ...) MUST therefore route its candidate + * distances through this predicate, otherwise the WHERE clause is silently + * dropped from the query plan and never re-checked by SQLite. + * + * @param distance - candidate distance to test + * @param idxStr - the xBestIndex/xFilter idxStr + * @param argc, argv - xFilter arguments, parallel to the idxStr entries + * @returns 1 if the distance satisfies all constraints, 0 otherwise. + */ +static int vec0_distance_constraints_satisfied(f32 distance, const char *idxStr, + int argc, + sqlite3_value **argv) { + for (int i = 0; i < argc; i++) { + int idx = 1 + (i * 4); + if (idxStr[idx + 0] != VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT) { + continue; + } + // TODO casts f64 to f32, is that a problem? + f32 target = (f32)sqlite3_value_double(argv[i]); + switch ((vec0_distance_constraint_operator)idxStr[idx + 1]) { + case VEC0_DISTANCE_CONSTRAINT_GT: + if (!(distance > target)) + return 0; + break; + case VEC0_DISTANCE_CONSTRAINT_GE: + if (!(distance >= target)) + return 0; + break; + case VEC0_DISTANCE_CONSTRAINT_LT: + if (!(distance < target)) + return 0; + break; + case VEC0_DISTANCE_CONSTRAINT_LE: + if (!(distance <= target)) + return 0; + break; + } + } + return 1; +} + static int vec0BestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdxInfo) { vec0_vtab *p = (vec0_vtab *)pVTab; /** @@ -7543,50 +7590,11 @@ int vec0Filter_knn_chunks_iter(vec0_vtab *p, sqlite3_stmt *stmtChunks, } if(hasDistanceConstraints) { - for(int i = 0; i < argc; i++) { - int idx = 1 + (i * 4); - char kind = idxStr[idx + 0]; - // TODO casts f64 to f32, is that a problem? - f32 target = (f32) sqlite3_value_double(argv[i]); - - if(kind != VEC0_IDXSTR_KIND_KNN_DISTANCE_CONSTRAINT) { - continue; - } - vec0_distance_constraint_operator op = idxStr[idx + 1]; - - switch(op) { - case VEC0_DISTANCE_CONSTRAINT_GE: { - for(int i = 0; i < p->chunk_size;i++) { - if(bitmap_get(b, i) && !(chunk_distances[i] >= target)) { - bitmap_set(b, i, 0); - } - } - break; - } - case VEC0_DISTANCE_CONSTRAINT_GT: { - for(int i = 0; i < p->chunk_size;i++) { - if(bitmap_get(b, i) && !(chunk_distances[i] > target)) { - bitmap_set(b, i, 0); - } - } - break; - } - case VEC0_DISTANCE_CONSTRAINT_LE: { - for(int i = 0; i < p->chunk_size;i++) { - if(bitmap_get(b, i) && !(chunk_distances[i] <= target)) { - bitmap_set(b, i, 0); - } - } - break; - } - case VEC0_DISTANCE_CONSTRAINT_LT: { - for(int i = 0; i < p->chunk_size;i++) { - if(bitmap_get(b, i) && !(chunk_distances[i] < target)) { - bitmap_set(b, i, 0); - } - } - break; - } + for(int i = 0; i < p->chunk_size; i++) { + if(bitmap_get(b, i) && + !vec0_distance_constraints_satisfied(chunk_distances[i], idxStr, + argc, argv)) { + bitmap_set(b, i, 0); } } } @@ -7796,6 +7804,22 @@ static int vec0Filter_knn_diskann( } } + // Apply any `distance` constraints from the WHERE clause. vec0BestIndex() + // omits these from the query plan, so SQLite will not re-check them. + { + int kept = 0; + for (int si = 0; si < resultCount; si++) { + if (!vec0_distance_constraints_satisfied(resultDistances[si], idxStr, + argc, argv)) { + continue; + } + resultRowids[kept] = resultRowids[si]; + resultDistances[kept] = resultDistances[si]; + kept++; + } + resultCount = kept; + } + knn_data->k = resultCount; knn_data->k_used = resultCount; knn_data->rowids = resultRowids; @@ -8071,6 +8095,22 @@ int vec0Filter_knn(vec0_cursor *pCur, vec0_vtab *p, int idxNum, if (rc != SQLITE_OK) { goto cleanup; } + // Apply any `distance` constraints from the WHERE clause. vec0BestIndex() + // omits these from the query plan, so SQLite will not re-check them. + { + i64 kept = 0; + for (i64 j = 0; j < knn_data->k_used; j++) { + if (!vec0_distance_constraints_satisfied(knn_data->distances[j], idxStr, + argc, argv)) { + continue; + } + knn_data->rowids[kept] = knn_data->rowids[j]; + knn_data->distances[kept] = knn_data->distances[j]; + kept++; + } + knn_data->k = kept; + knn_data->k_used = kept; + } pCur->knn_data = knn_data; pCur->query_plan = VEC0_QUERY_PLAN_KNN; rc = SQLITE_OK; diff --git a/tests/test-knn-distance-constraints.py b/tests/test-knn-distance-constraints.py index 1adcb0fc..0ccbf223 100644 --- a/tests/test-knn-distance-constraints.py +++ b/tests/test-knn-distance-constraints.py @@ -1,3 +1,6 @@ +import random + +import pytest import sqlite3 from helpers import exec @@ -39,6 +42,148 @@ def test_normal(db, snapshot): assert exec(db, BASE_KNN + "AND is_odd == TRUE AND distance BETWEEN 7 AND 10", ["[1]", 5]) == snapshot() +# vec0BestIndex() sets `omit = 1` on distance constraints, which promises +# SQLite that the vtab applies them itself. Every KNN backend must honor that +# promise -- if one doesn't, the WHERE clause is silently dropped from the +# query plan and rows that violate it are returned with no error at all. +def _has_ivf(): + db = sqlite3.connect(":memory:") + db.enable_load_extension(True) + db.load_extension("dist/vec0") + flags = db.execute("SELECT vec_debug()").fetchone()[0] + return "ivf" in flags.split("Build flags:")[-1].split() + + +ANN_INDEX_DEFS = [ + pytest.param("", id="flat"), + pytest.param( + "INDEXED BY rescore(quantizer=bit, oversample=16)", id="rescore-bit" + ), + pytest.param( + "INDEXED BY rescore(quantizer=int8, oversample=16)", id="rescore-int8" + ), + pytest.param( + "INDEXED BY diskann(neighbor_quantizer=int8)", id="diskann" + ), + pytest.param( + "INDEXED BY ivf(nlist=16, nprobe=8)", + id="ivf", + marks=pytest.mark.skipif( + not _has_ivf(), + reason="IVF not enabled (compile with -DSQLITE_VEC_EXPERIMENTAL_IVF_ENABLE=1)", + ), + ), +] + +# Indexes that apply the constraint to their candidate pool *before* the top-k +# truncation, and can therefore still return a full k rows. DiskANN and IVF +# filter their final result set instead, so a lower-bound constraint there +# legitimately yields fewer than k rows. +INDEX_DEFS_FILLING_K = [p for p in ANN_INDEX_DEFS if p.id in ("flat", "rescore-bit", "rescore-int8")] + +DIMENSIONS = 8 +NROWS = 200 + + +def _seed(db, index_def): + db.execute( + f"CREATE VIRTUAL TABLE v USING vec0(embedding float[{DIMENSIONS}] {index_def})" + ) + rng = random.Random(0) + rows = [ + (i, "[" + ",".join(str(rng.random()) for _ in range(DIMENSIONS)) + "]") + for i in range(1, NROWS + 1) + ] + db.executemany("INSERT INTO v(rowid, embedding) VALUES (?, ?)", rows) + return rows + + +@pytest.mark.parametrize("index_def", ANN_INDEX_DEFS) +@pytest.mark.parametrize( + "op,predicate", + [ + ("<=", lambda d, t: d <= t), + ("<", lambda d, t: d < t), + (">=", lambda d, t: d >= t), + (">", lambda d, t: d > t), + ], +) +def test_distance_constraint_is_honored_by_every_index(db, index_def, op, predicate): + """Regression test for #308. + + Distance constraints were only implemented in the FLAT chunk scan. The + rescore path (added later, in #276) never read them back out of idxStr, so + `AND distance <= x` was silently ignored and every top-k row was returned. + """ + rows = _seed(db, index_def) + query = rows[0][1] + + unfiltered = db.execute( + "SELECT rowid, distance FROM v WHERE embedding MATCH ? AND k = 20", + (query,), + ).fetchall() + assert len(unfiltered) == 20 + + # Pick a threshold in the middle of the observed distance range, so that + # the constraint is neither a no-op nor filters everything out. + distances = sorted(row["distance"] for row in unfiltered) + threshold = distances[len(distances) // 2] + + filtered = db.execute( + f"SELECT rowid, distance FROM v WHERE embedding MATCH ? AND k = 20 " + f"AND distance {op} ?", + (query, threshold), + ).fetchall() + + violations = [row["distance"] for row in filtered if not predicate(row["distance"], threshold)] + assert violations == [], ( + f"{len(violations)} row(s) violating `distance {op} {threshold}` were " + f"returned by index `{index_def or 'flat'}`" + ) + + +@pytest.mark.parametrize("index_def", ANN_INDEX_DEFS) +def test_distance_constraint_filtering_everything_is_not_an_error(db, index_def): + """An impossible constraint must yield an empty result set, not an error. + + On the rescore path this exercises the case where every rescored candidate + is filtered out: the result arrays are then zero-length, and a naive + sqlite3_malloc(0) returning NULL would be misreported as SQLITE_NOMEM. + """ + rows = _seed(db, index_def) + result = db.execute( + "SELECT rowid FROM v WHERE embedding MATCH ? AND k = 10 AND distance < ?", + (rows[0][1], -1.0), + ).fetchall() + assert result == [] + + +@pytest.mark.parametrize("index_def", INDEX_DEFS_FILLING_K) +def test_distance_constraint_lower_bound_still_fills_k(db, index_def): + """A `distance >` constraint must not shrink the result set. + + The constraint has to be applied to the candidate pool *before* the top-k + truncation. Applying it afterwards would chop off the head of the sorted + results and return fewer than k rows. + """ + rows = _seed(db, index_def) + query = rows[0][1] + + unfiltered = db.execute( + "SELECT distance FROM v WHERE embedding MATCH ? AND k = 5", (query,) + ).fetchall() + # Exclude the 5 nearest neighbors; there are plenty of rows left beyond them. + threshold = max(row["distance"] for row in unfiltered) + + filtered = db.execute( + "SELECT rowid, distance FROM v WHERE embedding MATCH ? AND k = 5 AND distance > ?", + (query, threshold), + ).fetchall() + + assert len(filtered) == 5 + assert all(row["distance"] > threshold for row in filtered) + + class Row: def __init__(self): pass