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:
+34
-21
@@ -532,34 +532,47 @@ 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 {
|
||||||
.map_err(|e| e.to_string())?;
|
// No post-query filtering — a single fetch of exactly `limit` results is optimal.
|
||||||
let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?;
|
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 {
|
// Post-query filters are active. Progressively double the fetch batch until we
|
||||||
images.retain(|image| image.folder_id == folder_id);
|
// collect enough results or exhaust the vector table, so we never over-fetch
|
||||||
}
|
// more than necessary while still filling the requested page.
|
||||||
if let Some(media_kind) = params.media_kind.as_deref() {
|
let mut batch = limit;
|
||||||
images.retain(|image| image.media_kind == media_kind);
|
loop {
|
||||||
}
|
let ids = vector::search_image_ids_by_embedding(&conn, &embedding, batch)
|
||||||
if params.favorites_only.unwrap_or(false) {
|
.map_err(|e| e.to_string())?;
|
||||||
images.retain(|image| image.favorite);
|
let exhausted = ids.len() < batch;
|
||||||
}
|
let mut images = db::get_images_by_ids(&conn, &ids).map_err(|e| e.to_string())?;
|
||||||
if let Some(rating_min) = params.rating_min {
|
|
||||||
images.retain(|image| image.rating >= rating_min);
|
|
||||||
}
|
|
||||||
images.truncate(limit);
|
|
||||||
|
|
||||||
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]
|
#[tauri::command]
|
||||||
|
|||||||
+44
-32
@@ -23,41 +23,53 @@ fn cache() -> &'static RwLock<Option<CachedHnswIndex>> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn build_index(conn: &Connection) -> Result<CachedHnswIndex> {
|
fn build_index(conn: &Connection) -> Result<CachedHnswIndex> {
|
||||||
let embeddings = vector::get_all_image_embeddings_with_ids(conn, None)?;
|
// Read the revision *before* fetching embeddings so we can detect any write
|
||||||
let max_elements = embeddings.len().max(1);
|
// that races with the build. If the revision advances while we are building,
|
||||||
let max_layer = 16.min((max_elements as f32).ln().trunc() as usize).max(1);
|
// the resulting index would be stale — retry until it is stable.
|
||||||
let mut hnsw = Hnsw::<f32, DistCosine>::new(
|
loop {
|
||||||
HNSW_MAX_CONNECTIONS,
|
let revision_before = vector::get_embedding_revision(conn)?;
|
||||||
max_elements,
|
|
||||||
max_layer,
|
|
||||||
HNSW_EF_CONSTRUCTION,
|
|
||||||
DistCosine {},
|
|
||||||
);
|
|
||||||
|
|
||||||
let image_ids_by_external = embeddings
|
let embeddings = vector::get_all_image_embeddings_with_ids(conn, None)?;
|
||||||
.iter()
|
let max_elements = embeddings.len().max(1);
|
||||||
.map(|(image_id, _)| *image_id)
|
let max_layer = 16.min((max_elements as f32).ln().trunc() as usize).max(1);
|
||||||
.collect::<Vec<_>>();
|
let mut hnsw = Hnsw::<f32, DistCosine>::new(
|
||||||
let external_by_image_id = image_ids_by_external
|
HNSW_MAX_CONNECTIONS,
|
||||||
.iter()
|
max_elements,
|
||||||
.enumerate()
|
max_layer,
|
||||||
.map(|(external_id, image_id)| (*image_id, external_id))
|
HNSW_EF_CONSTRUCTION,
|
||||||
.collect::<HashMap<_, _>>();
|
DistCosine {},
|
||||||
let data_with_id = embeddings
|
);
|
||||||
.iter()
|
|
||||||
.enumerate()
|
|
||||||
.map(|(external_id, (_, embedding))| (embedding, external_id))
|
|
||||||
.collect::<Vec<_>>();
|
|
||||||
|
|
||||||
hnsw.parallel_insert(&data_with_id);
|
let image_ids_by_external = embeddings
|
||||||
hnsw.set_searching_mode(true);
|
.iter()
|
||||||
|
.map(|(image_id, _)| *image_id)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let external_by_image_id = image_ids_by_external
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(external_id, image_id)| (*image_id, external_id))
|
||||||
|
.collect::<HashMap<_, _>>();
|
||||||
|
let data_with_id = embeddings
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(external_id, (_, embedding))| (embedding, external_id))
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
|
||||||
Ok(CachedHnswIndex {
|
hnsw.parallel_insert(&data_with_id);
|
||||||
revision: vector::get_embedding_revision(conn)?,
|
hnsw.set_searching_mode(true);
|
||||||
image_ids_by_external,
|
|
||||||
external_by_image_id,
|
// If the revision is unchanged the index reflects a consistent snapshot.
|
||||||
hnsw,
|
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<()> {
|
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