diff --git a/sqlite-vec.c b/sqlite-vec.c index ab77083c..7a52250a 100644 --- a/sqlite-vec.c +++ b/sqlite-vec.c @@ -10871,6 +10871,15 @@ struct vec_static_blob_data { static_blob static_blobs[MAX_STATIC_BLOBS]; }; +/** Frees the registry and the names vec_static_blobsUpdate() copied in. */ +static void vec_static_blob_data_free(void *p) { + vec_static_blob_data *data = p; + for (int i = 0; i < MAX_STATIC_BLOBS; i++) { + sqlite3_free(data->static_blobs[i].name); + } + sqlite3_free(data); +} + typedef struct vec_static_blobs_vtab vec_static_blobs_vtab; struct vec_static_blobs_vtab { sqlite3_vtab base; @@ -11177,7 +11186,10 @@ static int vec_static_blob_entriesOpen(sqlite3_vtab *p, static int vec_static_blob_entriesClose(sqlite3_vtab_cursor *cur) { vec_static_blob_entries_cursor *pCur = (vec_static_blob_entries_cursor *)cur; - sqlite3_free(pCur->knn_data); + if (pCur->knn_data) { + sbe_query_knn_data_clear(pCur->knn_data); + sqlite3_free(pCur->knn_data); + } sqlite3_free(pCur); return SQLITE_OK; } @@ -11358,7 +11370,7 @@ static int vec_static_blob_entriesFilter(sqlite3_vtab_cursor *pVtabCursor, } bitmap_fill(candidates, bsize); - for (size_t i = bsize; i >= p->blob->nvectors; i--) { + for (size_t i = p->blob->nvectors; i < bsize; i++) { bitmap_set(candidates, i, 0); } i32 k_used = 0; @@ -11657,7 +11669,7 @@ sqlite3_vec_static_blobs_init(sqlite3 *db, char **pzErrMsg, return rc; rc = sqlite3_create_module_v2(db, "vec_static_blobs", &vec_static_blobsModule, - static_blob_data, sqlite3_free); + static_blob_data, vec_static_blob_data_free); if (rc != SQLITE_OK) return rc; rc = sqlite3_create_module_v2(db, "vec_static_blob_entries", diff --git a/tests/test-loadable.py b/tests/test-loadable.py index 0ef935c5..9f64406a 100644 --- a/tests/test-loadable.py +++ b/tests/test-loadable.py @@ -264,6 +264,19 @@ def test_vec_static_blob_entries(): ] +def test_vec_static_blob_entries_knn_with_whole_candidate_bytes(): + # 64 vectors fill the KNN's candidate bitmap to its last byte. The blob + # reads the array's memory, so the array stays referenced while it is queried. + db = connect(EXT_PATH, extra_entrypoint="sqlite3_vec_static_blobs_init") + vectors = np.array([[i / 64] * 4 for i in range(64)], np.float32) + register_numpy(db, "w", vectors) + rows = db.execute( + "select rowid from w where vector match ? and k = 3 order by distance", + [np.array([0.5] * 4, dtype=np.float32)], + ).fetchall() + assert [r[0] for r in rows] == [32, 33, 31] + + def test_limits(): db = connect(EXT_PATH) with _raises(