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 conn = db.get().map_err(|e| e.to_string())?;
|
||||||
let limit = params.limit.unwrap_or(64);
|
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()
|
let has_filters = params.folder_id.is_some()
|
||||||
|| params.media_kind.is_some()
|
|| params.media_kind.is_some()
|
||||||
|| params.favorites_only.unwrap_or(false)
|
|| params.favorites_only.unwrap_or(false)
|
||||||
|| params.rating_min.map_or(false, |r| r > 0);
|
|| 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())?;
|
.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())?;
|
let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?;
|
||||||
|
|
||||||
if let Some(folder_id) = params.folder_id {
|
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 {
|
if let Some(rating_min) = params.rating_min {
|
||||||
images.retain(|image| image.rating >= 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]
|
#[tauri::command]
|
||||||
|
|||||||
@@ -23,6 +23,12 @@ fn cache() -> &'static RwLock<Option<CachedHnswIndex>> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn build_index(conn: &Connection) -> Result<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 embeddings = vector::get_all_image_embeddings_with_ids(conn, None)?;
|
||||||
let max_elements = embeddings.len().max(1);
|
let max_elements = embeddings.len().max(1);
|
||||||
let max_layer = 16.min((max_elements as f32).ln().trunc() as usize).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.parallel_insert(&data_with_id);
|
||||||
hnsw.set_searching_mode(true);
|
hnsw.set_searching_mode(true);
|
||||||
|
|
||||||
Ok(CachedHnswIndex {
|
// If the revision is unchanged the index reflects a consistent snapshot.
|
||||||
revision: vector::get_embedding_revision(conn)?,
|
let revision_after = vector::get_embedding_revision(conn)?;
|
||||||
|
if revision_before == revision_after {
|
||||||
|
return Ok(CachedHnswIndex {
|
||||||
|
revision: revision_after,
|
||||||
image_ids_by_external,
|
image_ids_by_external,
|
||||||
external_by_image_id,
|
external_by_image_id,
|
||||||
hnsw,
|
hnsw,
|
||||||
})
|
});
|
||||||
|
}
|
||||||
|
// A concurrent write advanced the revision — discard this build and retry.
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn ensure_index(conn: &Connection) -> Result<()> {
|
fn ensure_index(conn: &Connection) -> Result<()> {
|
||||||
|
|||||||
@@ -33,6 +33,12 @@ pub fn migrate(conn: &Connection) -> Result<()> {
|
|||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
pub fn delete_embedding(conn: &Connection, image_id: i64) -> Result<()> {
|
pub fn delete_embedding(conn: &Connection, image_id: i64) -> Result<()> {
|
||||||
conn.execute("DELETE FROM image_vec WHERE image_id = ?1", [image_id])?;
|
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(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user