From 94829af41f5a9a31e2a7f3f07b6bd6a1230d50c4 Mon Sep 17 00:00:00 2001 From: realgetOff Date: Tue, 4 Aug 2026 12:36:02 +0200 Subject: [PATCH] Apply KNN distance constraints in every index backend vec0BestIndex() sets aConstraintUsage[].omit = 1 on constraints against the `distance` column, which tells SQLite the vtab will enforce them itself and lets SQLite drop the term from the query plan. Only the FLAT chunk scan ever read those constraints back out of idxStr, so on rescore, DiskANN and IVF columns the predicate was silently discarded and rows violating it were returned with no error. This was a regression: #166 added the distance constraints, and the ANN backends merged six weeks later (#276, #277, #278) never wired them up. Introduce vec0_distance_constraints_satisfied() as the single point of truth for the predicate and route every backend through it: - FLAT keeps its pre-filter semantics; its four nested switch cases collapse into one loop over the chunk bitmap (net -44 lines, snapshots unchanged). - rescore filters the rescored float distances before the top-k truncation, so a lower-bound constraint still yields k rows. Coarse quantized distances from phase 1 are not comparable to a user-supplied threshold, so the filter cannot be pushed down into the quantized scan. - DiskANN and IVF compact their final result sets. Also fixes a latent issue on the rescore path: when the constraint filters out every candidate, result_k becomes 0 and sqlite3_malloc(0) returns NULL, which the existing check would have misreported as SQLITE_NOMEM. Adds parametrized coverage over flat / rescore-bit / rescore-int8 / diskann, plus ivf behind a build-flag skipif mirroring tests/conftest.py. 17 of the new assertions fail without this change on a default build (22 with IVF enabled) and all pass with it. Fixes #308 --- sqlite-vec-rescore.c | 30 +++++ sqlite-vec.c | 128 ++++++++++++++-------- tests/test-knn-distance-constraints.py | 145 +++++++++++++++++++++++++ 3 files changed, 259 insertions(+), 44 deletions(-) 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