From ae3d6e10343e7484ed4da465e4f4da70ded9f3a3 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 7 Jun 2026 22:46:05 +0100 Subject: [PATCH] fix(search): close three P2 correctness gaps from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- src-tauri/src/commands.rs | 55 +++++++++++++++++---------- src-tauri/src/hnsw_index.rs | 76 +++++++++++++++++++++---------------- src-tauri/src/vector.rs | 6 +++ 3 files changed, 84 insertions(+), 53 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 647d5d1..cd239ad 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -532,34 +532,47 @@ 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) - .map_err(|e| e.to_string())?; - let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?; + 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()); + } - if let Some(folder_id) = params.folder_id { - images.retain(|image| image.folder_id == folder_id); - } - if let Some(media_kind) = params.media_kind.as_deref() { - images.retain(|image| image.media_kind == media_kind); - } - if params.favorites_only.unwrap_or(false) { - images.retain(|image| image.favorite); - } - if let Some(rating_min) = params.rating_min { - images.retain(|image| image.rating >= rating_min); - } - images.truncate(limit); + // 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())?; - Ok(images) + if let Some(folder_id) = params.folder_id { + images.retain(|image| image.folder_id == folder_id); + } + if let Some(media_kind) = params.media_kind.as_deref() { + images.retain(|image| image.media_kind == media_kind); + } + if params.favorites_only.unwrap_or(false) { + images.retain(|image| image.favorite); + } + if let Some(rating_min) = params.rating_min { + images.retain(|image| image.rating >= rating_min); + } + + if images.len() >= limit || exhausted { + images.truncate(limit); + return Ok(images); + } + batch = (batch * 2).min(8192); + } } #[tauri::command] diff --git a/src-tauri/src/hnsw_index.rs b/src-tauri/src/hnsw_index.rs index b7ef6c7..2cdc110 100644 --- a/src-tauri/src/hnsw_index.rs +++ b/src-tauri/src/hnsw_index.rs @@ -23,41 +23,53 @@ fn cache() -> &'static RwLock> { } fn build_index(conn: &Connection) -> Result { - 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); - let mut hnsw = Hnsw::::new( - HNSW_MAX_CONNECTIONS, - max_elements, - max_layer, - HNSW_EF_CONSTRUCTION, - DistCosine {}, - ); + // 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 image_ids_by_external = embeddings - .iter() - .map(|(image_id, _)| *image_id) - .collect::>(); - let external_by_image_id = image_ids_by_external - .iter() - .enumerate() - .map(|(external_id, image_id)| (*image_id, external_id)) - .collect::>(); - let data_with_id = embeddings - .iter() - .enumerate() - .map(|(external_id, (_, embedding))| (embedding, external_id)) - .collect::>(); + 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); + let mut hnsw = Hnsw::::new( + HNSW_MAX_CONNECTIONS, + max_elements, + max_layer, + HNSW_EF_CONSTRUCTION, + DistCosine {}, + ); - hnsw.parallel_insert(&data_with_id); - hnsw.set_searching_mode(true); + let image_ids_by_external = embeddings + .iter() + .map(|(image_id, _)| *image_id) + .collect::>(); + let external_by_image_id = image_ids_by_external + .iter() + .enumerate() + .map(|(external_id, image_id)| (*image_id, external_id)) + .collect::>(); + let data_with_id = embeddings + .iter() + .enumerate() + .map(|(external_id, (_, embedding))| (embedding, external_id)) + .collect::>(); - Ok(CachedHnswIndex { - revision: vector::get_embedding_revision(conn)?, - image_ids_by_external, - external_by_image_id, - hnsw, - }) + hnsw.parallel_insert(&data_with_id); + hnsw.set_searching_mode(true); + + // 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<()> { diff --git a/src-tauri/src/vector.rs b/src-tauri/src/vector.rs index 25a3220..98b5349 100644 --- a/src-tauri/src/vector.rs +++ b/src-tauri/src/vector.rs @@ -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(()) }