Apply KNN distance constraints in every index backend

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));

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;

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