Read a quantized IVF column's vectors from _ivf_vectors

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

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

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