Apply rowid filters to IVF KNN queries

sqlite-vec-ivf.c

@@ -670,13 +670,54 @@ static int ivf_delete(vec0_vtab *p, int col_idx, i64 rowid) {
 // Point query
 // ============================================================================
 
+/**
+ * Read a quantized column's full-precision vector from _ivf_vectors.
+ */
+static int ivf_full_vector_read(vec0_vtab *p, i64 rowid, int col_idx,
+                                 void **outVector, int *outVectorSize) {
+  int vecSize = ivf_full_vec_size(p, col_idx);
+  int rc = ivf_ensure_stmt(p, &p->stmtIvfVectorsRead[col_idx],
+      "SELECT vector FROM " VEC0_SHADOW_IVF_VECTORS_NAME " WHERE rowid = ?",
+      col_idx);
+  if (rc != SQLITE_OK) return rc;
+  sqlite3_stmt *s = p->stmtIvfVectorsRead[col_idx];
+  sqlite3_reset(s);
+  sqlite3_bind_int64(s, 1, rowid);
+  rc = sqlite3_step(s);
+  if (rc == SQLITE_ROW) {
+    void *buf = NULL;
+    if (sqlite3_column_bytes(s, 0) != vecSize) {
+      rc = SQLITE_CORRUPT;
+    } else if (!(buf = sqlite3_malloc(vecSize))) {
+      rc = SQLITE_NOMEM;
+    } else {
+      memcpy(buf, sqlite3_column_blob(s, 0), vecSize);
+      *outVector = buf;
+      if (outVectorSize) *outVectorSize = vecSize;
+      rc = SQLITE_OK;
+    }
+  } else if (rc == SQLITE_DONE) {
+    rc = SQLITE_EMPTY;
+  }
+  // A cached statement left on a row stays in progress until the next commit,
+  // and a connection with a statement in progress refuses VACUUM.
+  sqlite3_reset(s);
+  return rc;
+}
+
 static int ivf_get_vector_data(vec0_vtab *p, i64 rowid, int col_idx,
                                 void **outVector, int *outVectorSize) {
   int rc;
   int vecSize = ivf_vec_size(p, col_idx);
   i64 cell_id = 0;
   int slot = -1;
 
+  // A quantized column's cells hold its quantized vectors, which are not the
+  // vectors it was given.
+  if (p->vector_columns[col_idx].ivf.quantizer != VEC0_IVF_QUANTIZER_NONE) {
+    return ivf_full_vector_read(p, rowid, col_idx, outVector, outVectorSize);
+  }
+
   rc = ivf_ensure_stmt(p, &p->stmtIvfRowidMapLookup[col_idx],
       "SELECT cell_id, slot FROM " VEC0_SHADOW_IVF_ROWID_MAP_NAME
       " WHERE rowid = ?", col_idx);
@@ -1222,8 +1263,38 @@ static int ivf_scan_cells_from_stmt(vec0_vtab *p, int col_idx,
   return SQLITE_OK;
 }
 
+/**
+ * Rank the rows a `rowid in (...)` constraint names by their exact distances
+ * to the query. A rowid the table does not have is skipped.
+ */
+static int ivf_rank_rowids(vec0_vtab *p, int col_idx, const void *queryVector,
+                            struct Array *rowidsIn,
+                            struct IvfCandidate **candidates,
+                            int *nCandidates, int *cap) {
+  for (size_t i = 0; i < rowidsIn->length; i++) {
+    i64 rowid = ((const i64 *)rowidsIn->z)[i];
+    void *vec = NULL;
+    int rc = ivf_get_vector_data(p, rowid, col_idx, &vec, NULL);
+    if (rc == SQLITE_EMPTY) continue;
+    if (rc != SQLITE_OK) return rc;
+    if (*nCandidates >= *cap) {
+      *cap *= 2;
+      struct IvfCandidate *tmp = sqlite3_realloc64(*candidates, (i64)*cap * sizeof(struct IvfCandidate));
+      if (!tmp) { sqlite3_free(vec); return SQLITE_NOMEM; }
+      *candidates = tmp;
+    }
+    (*candidates)[*nCandidates].rowid = rowid;
+    (*candidates)[*nCandidates].distance = ivf_distance_float(p, col_idx,
+        (const float *)queryVector, (const float *)vec);
+    (*nCandidates)++;
+    sqlite3_free(vec);
+  }
+  return SQLITE_OK;
+}
+
 static int ivf_query_knn(vec0_vtab *p, int col_idx,
                           const void *queryVector, int queryVectorSize,
+                          struct Array *arrayRowidsIn,
                           i64 k, struct vec0_query_knn_data *knn_data) {
   UNUSED_PARAMETER(queryVectorSize);
   int rc;
@@ -1246,7 +1317,14 @@ static int ivf_query_knn(vec0_vtab *p, int col_idx,
   struct IvfCandidate *candidates = sqlite3_malloc64((i64)cap * sizeof(struct IvfCandidate));
   if (!candidates) { sqlite3_free(queryQ); return SQLITE_NOMEM; }
 
-  if (trained) {
+  if (arrayRowidsIn) {
+    // A rowid filter names its candidates, so each is ranked by its exact
+    // distance; the probed cells hold only the rows near the query, which
+    // the filter may not name.
+    rc = ivf_rank_rowids(p, col_idx, queryVector, arrayRowidsIn,
+                         &candidates, &nCandidates, &cap);
+    if (rc != SQLITE_OK) { sqlite3_free(queryQ); sqlite3_free(candidates); return rc; }
+  } else if (trained) {
     // Find top nprobe centroids using quantized distance
     int nlist = 0;
     rc = ivf_ensure_stmt(p, &p->stmtIvfCentroidsAll[col_idx],
@@ -1333,7 +1411,8 @@ static int ivf_query_knn(vec0_vtab *p, int col_idx,
   qsort(candidates, nCandidates, sizeof(struct IvfCandidate), ivf_candidate_cmp);
 
   // Oversample re-ranking: re-score top (oversample*k) with full-precision vectors
-  if (oversample > 1 && quantizer != VEC0_IVF_QUANTIZER_NONE && nCandidates > 0) {
+  if (oversample > 1 && quantizer != VEC0_IVF_QUANTIZER_NONE && nCandidates > 0 &&
+      !arrayRowidsIn) {
     i64 rescore_n = collect_k < nCandidates ? collect_k : nCandidates;
     sqlite3_stmt *stmtVec = NULL;
     char *zSql = sqlite3_mprintf(

sqlite-vec.c

@@ -3581,6 +3581,7 @@ struct vec0_vtab {
   sqlite3_stmt *stmtIvfRowidMapLookup[VEC0_MAX_VECTOR_COLUMNS]; // SELECT cell_id,slot FROM rowid_map WHERE rowid=?
   sqlite3_stmt *stmtIvfRowidMapDelete[VEC0_MAX_VECTOR_COLUMNS]; // DELETE FROM rowid_map WHERE rowid=?
   sqlite3_stmt *stmtIvfCentroidsAll[VEC0_MAX_VECTOR_COLUMNS];   // SELECT centroid_id,centroid FROM centroids
+  sqlite3_stmt *stmtIvfVectorsRead[VEC0_MAX_VECTOR_COLUMNS];    // SELECT vector FROM ivf_vectors WHERE rowid=?
 #endif
 
   // select latest chunk from _chunks, getting chunk_id
@@ -3695,6 +3696,7 @@ void vec0_free_resources(vec0_vtab *p) {
     sqlite3_finalize(p->stmtIvfRowidMapLookup[i]); p->stmtIvfRowidMapLookup[i] = NULL;
     sqlite3_finalize(p->stmtIvfRowidMapDelete[i]); p->stmtIvfRowidMapDelete[i] = NULL;
     sqlite3_finalize(p->stmtIvfCentroidsAll[i]); p->stmtIvfCentroidsAll[i] = NULL;
+    sqlite3_finalize(p->stmtIvfVectorsRead[i]); p->stmtIvfVectorsRead[i] = NULL;
   }
 #endif
 #if SQLITE_VEC_ENABLE_DISKANN
@@ -8067,7 +8069,8 @@ int vec0Filter_knn(vec0_cursor *pCur, vec0_vtab *p, int idxNum,
   // IVF dispatch: if vector column has IVF, use IVF query instead of chunk scan
   if (vector_column->index_type == VEC0_INDEX_TYPE_IVF) {
     rc = ivf_query_knn(p, vectorColumnIdx, queryVector,
-                       (int)vector_column_byte_size(*vector_column), k, knn_data);
+                       (int)vector_column_byte_size(*vector_column),
+                       arrayRowidsIn, k, knn_data);
     if (rc != SQLITE_OK) {
       goto cleanup;
     }

tests/test-ivf-quantization.py

@@ -255,6 +255,32 @@ def test_ivf_quantized_delete(db):
     assert db.execute("SELECT count(*) FROM t_ivf_vectors00").fetchone()[0] == 9
 
 
+@pytest.mark.parametrize("quantizer", ["int8", "binary"])
+def test_ivf_quantized_vectors_read_back(db, quantizer):
+    """A quantized column returns the vectors it was given from a point
+    query, a full scan and a KNN query, before and after training."""
+    db.execute(
+        "CREATE VIRTUAL TABLE t USING vec0("
+        f"v float[16] indexed by ivf(nlist=2, quantizer={quantizer}))"
+    )
+    vecs = {
+        i: _f32([((i * 7 + j * 3) % 11) / 10 - 0.5 for j in range(16)])
+        for i in range(1, 11)
+    }
+    for rowid, v in vecs.items():
+        db.execute("INSERT INTO t(rowid, v) VALUES (?, ?)", [rowid, v])
+
+    for trained in [False, True]:
+        if trained:
+            db.execute("INSERT INTO t(t) VALUES ('compute-centroids')")
+        assert db.execute("SELECT v FROM t WHERE rowid = 3").fetchone()[0] == vecs[3]
+        assert {r[0]: r[1] for r in db.execute("SELECT rowid, v FROM t")} == vecs
+        knn = db.execute(
+            "SELECT rowid, v FROM t WHERE v MATCH ? AND k = 10", [vecs[3]]
+        ).fetchall()
+        assert {r[0]: r[1] for r in knn} == vecs
+
+
 def test_ivf_binary_rejects_non_multiple_of_8_dims(db):
     """Binary quantizer requires dimensions divisible by 8."""
     with pytest.raises(sqlite3.OperationalError):

tests/test-ivf.py

@@ -2,6 +2,7 @@
 import sqlite3
 import struct
 import math
+import random
 from helpers import _f32, exec
 
 
@@ -319,6 +320,35 @@ def test_ivf_knn_k_larger_than_n(db):
     assert len(rows) == 5
 
 
+@pytest.mark.parametrize("quantizer", ["none", "int8", "binary"])
+def test_ivf_knn_rowid_in(db, quantizer):
+    """A rowid filter returns the k nearest of the rows it names, as a flat
+    column does, whether or not their cells are probed."""
+    db.execute(
+        "CREATE VIRTUAL TABLE t USING vec0("
+        f"v float[8] indexed by ivf(nlist=4, nprobe=1, quantizer={quantizer}))"
+    )
+    db.execute("CREATE VIRTUAL TABLE flat USING vec0(v float[8])")
+    rng = random.Random(7)
+    for i in range(40):
+        v = _f32([rng.uniform(-1, 1) for _ in range(8)])
+        db.execute("INSERT INTO t(rowid, v) VALUES (?, ?)", [i, v])
+        db.execute("INSERT INTO flat(rowid, v) VALUES (?, ?)", [i, v])
+    query = db.execute("SELECT v FROM flat WHERE rowid = 3").fetchone()[0]
+
+    for trained in [False, True]:
+        if trained:
+            db.execute("INSERT INTO t(t) VALUES ('compute-centroids')")
+        for rowids in ["1, 2", "0, 20, 39", "5, 17, 30, 38, 99"]:
+            sql = (
+                "SELECT rowid, distance FROM {} WHERE v MATCH ? AND k = 3"
+                f" AND rowid IN ({rowids})"
+            )
+            got = [tuple(r) for r in db.execute(sql.format("t"), [query])]
+            want = [tuple(r) for r in db.execute(sql.format("flat"), [query])]
+            assert got == want
+
+
 # ============================================================================
 # Manual centroid import (set-centroid, assign-vectors)
 # ============================================================================