fix(search): close three P2 correctness gaps from code review
- vector: increment embedding_revision whenever delete_embedding is called so the HNSW cache is invalidated after image/folder deletions, not just after inserts - hnsw_index: fix race in build_index — read the revision before fetching embeddings and again after the parallel_insert; retry the whole build if the revision advanced during construction, ensuring the cached index always reflects a consistent snapshot - commands: replace the fixed 4× over-fetch in semantic_search_images with a progressive doubling loop; start at exactly `limit` candidates and double the batch (capped at 8192) until the page is filled or the vector table is exhausted, preventing both under- and over-fetching
This commit is contained in:
@@ -532,17 +532,26 @@ pub async fn semantic_search_images(
|
||||
let conn = db.get().map_err(|e| e.to_string())?;
|
||||
let limit = params.limit.unwrap_or(64);
|
||||
|
||||
// When post-query filters are active the vector search may discard many
|
||||
// candidates ranked just beyond the bare limit. Over-fetch by 4× so that
|
||||
// well-rated or folder-scoped matches are not silently omitted.
|
||||
let has_filters = params.folder_id.is_some()
|
||||
|| params.media_kind.is_some()
|
||||
|| params.favorites_only.unwrap_or(false)
|
||||
|| params.rating_min.map_or(false, |r| r > 0);
|
||||
let fetch_limit = if has_filters { limit * 4 } else { limit };
|
||||
|
||||
let ids = vector::search_image_ids_by_embedding(&conn, &embedding, fetch_limit)
|
||||
if !has_filters {
|
||||
// No post-query filtering — a single fetch of exactly `limit` results is optimal.
|
||||
let ids = vector::search_image_ids_by_embedding(&conn, &embedding, limit)
|
||||
.map_err(|e| e.to_string())?;
|
||||
return db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string());
|
||||
}
|
||||
|
||||
// Post-query filters are active. Progressively double the fetch batch until we
|
||||
// collect enough results or exhaust the vector table, so we never over-fetch
|
||||
// more than necessary while still filling the requested page.
|
||||
let mut batch = limit;
|
||||
loop {
|
||||
let ids = vector::search_image_ids_by_embedding(&conn, &embedding, batch)
|
||||
.map_err(|e| e.to_string())?;
|
||||
let exhausted = ids.len() < batch;
|
||||
let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?;
|
||||
|
||||
if let Some(folder_id) = params.folder_id {
|
||||
@@ -557,9 +566,13 @@ pub async fn semantic_search_images(
|
||||
if let Some(rating_min) = params.rating_min {
|
||||
images.retain(|image| image.rating >= rating_min);
|
||||
}
|
||||
images.truncate(limit);
|
||||
|
||||
Ok(images)
|
||||
if images.len() >= limit || exhausted {
|
||||
images.truncate(limit);
|
||||
return Ok(images);
|
||||
}
|
||||
batch = (batch * 2).min(8192);
|
||||
}
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
|
||||
@@ -23,6 +23,12 @@ fn cache() -> &'static RwLock<Option<CachedHnswIndex>> {
|
||||
}
|
||||
|
||||
fn build_index(conn: &Connection) -> Result<CachedHnswIndex> {
|
||||
// Read the revision *before* fetching embeddings so we can detect any write
|
||||
// that races with the build. If the revision advances while we are building,
|
||||
// the resulting index would be stale — retry until it is stable.
|
||||
loop {
|
||||
let revision_before = vector::get_embedding_revision(conn)?;
|
||||
|
||||
let embeddings = vector::get_all_image_embeddings_with_ids(conn, None)?;
|
||||
let max_elements = embeddings.len().max(1);
|
||||
let max_layer = 16.min((max_elements as f32).ln().trunc() as usize).max(1);
|
||||
@@ -52,12 +58,18 @@ fn build_index(conn: &Connection) -> Result<CachedHnswIndex> {
|
||||
hnsw.parallel_insert(&data_with_id);
|
||||
hnsw.set_searching_mode(true);
|
||||
|
||||
Ok(CachedHnswIndex {
|
||||
revision: vector::get_embedding_revision(conn)?,
|
||||
// If the revision is unchanged the index reflects a consistent snapshot.
|
||||
let revision_after = vector::get_embedding_revision(conn)?;
|
||||
if revision_before == revision_after {
|
||||
return Ok(CachedHnswIndex {
|
||||
revision: revision_after,
|
||||
image_ids_by_external,
|
||||
external_by_image_id,
|
||||
hnsw,
|
||||
})
|
||||
});
|
||||
}
|
||||
// A concurrent write advanced the revision — discard this build and retry.
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_index(conn: &Connection) -> Result<()> {
|
||||
|
||||
@@ -33,6 +33,12 @@ pub fn migrate(conn: &Connection) -> Result<()> {
|
||||
#[allow(dead_code)]
|
||||
pub fn delete_embedding(conn: &Connection, image_id: i64) -> Result<()> {
|
||||
conn.execute("DELETE FROM image_vec WHERE image_id = ?1", [image_id])?;
|
||||
// Advance the revision so any cached HNSW index is invalidated after deletions.
|
||||
conn.execute(
|
||||
"INSERT INTO app_kv (key, value) VALUES ('embedding_revision', 1)
|
||||
ON CONFLICT(key) DO UPDATE SET value = value + 1",
|
||||
[],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user