Add release cached stmts command to flush DiskANN prepared statements without disconnect

sqlite-vec.c

@@ -10271,6 +10271,24 @@ int vec0Update_Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv) {
   return SQLITE_OK;
 }
 
+/*
+ * Index-type-agnostic vec0 control commands dispatched via the FTS5-style
+ * command column. Returns SQLITE_OK if handled, SQLITE_EMPTY otherwise.
+ *
+ * Recognized:
+ *   "release-cached-stmts"
+ *      Finalize this vec0 vtab's cached prepared statements
+ *      (stmtRowidsInsertRowid, stmtDiskannNodeRead, etc.) without renaming
+ *      or destroying the table. They are re-prepared lazily on next use.
+ */
+static int vec0_handle_general_command(vec0_vtab *p, const char *cmd) {
+  if (strcmp(cmd, "release-cached-stmts") == 0) {
+    vec0_free_resources(p);
+    return SQLITE_OK;
+  }
+  return SQLITE_EMPTY;
+}
+
 static int vec0Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv,
                       sqlite_int64 *pRowid) {
   // DELETE operation
@@ -10297,6 +10315,8 @@ static int vec0Update(sqlite3_vtab *pVTab, int argc, sqlite3_value **argv,
         if (cmdRc == SQLITE_EMPTY)
           cmdRc = diskann_handle_command(p, cmd);
 #endif
+        if (cmdRc == SQLITE_EMPTY)
+          cmdRc = vec0_handle_general_command(p, cmd);
         if (cmdRc == SQLITE_EMPTY) {
           vtab_set_error(pVTab, "unknown vec0 command: '%s'", cmd);
           return SQLITE_ERROR;

tests/test-release-cached-stmts.py

@@ -0,0 +1,56 @@
+import sqlite3
+import pytest
+from helpers import _f32
+
+
+def test_release_cached_stmts_basic(db):
+    """The release-cached-stmts command should succeed and not affect data."""
+    db.execute("create virtual table v using vec0(a float[2], chunk_size=8)")
+    db.execute("insert into v(rowid, a) values (1, ?)", [_f32([0.1, 0.2])])
+    db.execute("insert into v(rowid, a) values (2, ?)", [_f32([0.3, 0.4])])
+
+    db.execute("insert into v(v) values ('release-cached-stmts')")
+
+    # Data should still be there and queryable; vec0's cached statements
+    # are re-prepared on demand.
+    rows = db.execute(
+        "select rowid from v where a match ? and k=10",
+        [_f32([0.1, 0.2])],
+    ).fetchall()
+    assert sorted(r[0] for r in rows) == [1, 2]
+
+
+def test_release_cached_stmts_before_any_use(db):
+    """Issuing the command before any inserts should be a no-op."""
+    db.execute("create virtual table v using vec0(a float[2], chunk_size=8)")
+    db.execute("insert into v(v) values ('release-cached-stmts')")
+    # Inserts and queries still work after.
+    db.execute("insert into v(rowid, a) values (1, ?)", [_f32([0.1, 0.2])])
+    rows = db.execute(
+        "select rowid from v where a match ? and k=10",
+        [_f32([0.1, 0.2])],
+    ).fetchall()
+    assert rows[0][0] == 1
+
+
+def test_release_cached_stmts_diskann(db):
+    """Works on DiskANN-indexed tables (the original motivating case)."""
+    db.execute("""
+        create virtual table v using vec0(
+            a float[8] indexed by diskann(neighbor_quantizer=binary)
+        )
+    """)
+    db.execute("insert into v(rowid, a) values (1, ?)", [_f32([0.1] * 8)])
+    db.execute("insert into v(v) values ('release-cached-stmts')")
+    rows = db.execute(
+        "select rowid from v where a match ? and k=10",
+        [_f32([0.1] * 8)],
+    ).fetchall()
+    assert rows[0][0] == 1
+
+
+def test_release_cached_stmts_unknown_subcommand(db):
+    """Unknown vec0 commands should still error as before."""
+    db.execute("create virtual table v using vec0(a float[2], chunk_size=8)")
+    with pytest.raises(sqlite3.OperationalError, match="unknown vec0 command"):
+        db.execute("insert into v(v) values ('not-a-real-command')")