From 05180724cb137e61b99d3dd5da4d656268a4dd05 Mon Sep 17 00:00:00 2001 From: mysma-9403 <64923976+mysma-9403@users.noreply.github.com> Date: Sat, 30 May 2026 15:51:20 +0200 Subject: [PATCH] perf(memory_tree): batch embedding fetch in topic rerank (eliminate N sequential SQLite round-trips) (#2927) Co-authored-by: oxoxDev --- .../memory_store/chunks/embeddings.rs | 111 ++++++++++++++++++ src/openhuman/memory_store/chunks/store.rs | 4 +- .../memory_store/chunks/store_tests.rs | 97 +++++++++++++++ src/openhuman/memory_store/trees/store.rs | 102 ++++++++++++++++ .../memory_store/trees/store_tests.rs | 84 +++++++++++++ src/openhuman/memory_tree/retrieval/topic.rs | 75 +++++++++--- 6 files changed, 455 insertions(+), 18 deletions(-) diff --git a/src/openhuman/memory_store/chunks/embeddings.rs b/src/openhuman/memory_store/chunks/embeddings.rs index 1ab137fce..533e68e73 100644 --- a/src/openhuman/memory_store/chunks/embeddings.rs +++ b/src/openhuman/memory_store/chunks/embeddings.rs @@ -3,6 +3,7 @@ use crate::openhuman::config::Config; use anyhow::{Context, Result}; use chrono::Utc; use rusqlite::{Connection, OptionalExtension}; +use std::collections::HashMap; // ── Phase 2: embedding column accessors ───────────────────────────────── @@ -287,3 +288,113 @@ fn embedding_from_blob(bytes: &[u8], dim: i64, label: &str) -> Result>` containing **only the chunks +/// that have a vector under `model_signature`**. Missing chunks are simply +/// absent from the map — callers handle them the same way as a `None` +/// return from the single-row [`get_chunk_embedding_for_signature`]. +/// +/// ## Why this exists +/// +/// The retrieval rerank path (`memory_tree::retrieval::topic:: +/// rerank_by_semantic_similarity`) used to call the single-row helper inside +/// a `for h in hits { spawn_blocking(...).await }` loop. With +/// `LOOKUP_HEADROOM = 200` that meant 200 sequential SQLite round-trips on +/// every entity-scoped query with `query=…`. This helper collapses that to +/// `ceil(n / MAX_EMBEDDING_BATCH)` round-trips while preserving the exact +/// `Option>` semantics of the per-row helper (missing row → absent +/// key → caller treats as `None`). +/// +/// Order of input ids is irrelevant; callers re-decorate from the returned +/// map by id, so the rerank loop preserves its original hit iteration +/// order and therefore its tie-break behaviour. +pub fn get_chunk_embeddings_for_signature_batch( + config: &Config, + chunk_ids: &[String], + model_signature: &str, +) -> Result>> { + if chunk_ids.is_empty() { + return Ok(HashMap::new()); + } + with_connection(config, |conn| { + let mut out: HashMap> = HashMap::with_capacity(chunk_ids.len()); + // Chunk the id list to stay safely under SQLite's + // SQLITE_MAX_VARIABLE_NUMBER cap. For the current + // LOOKUP_HEADROOM=200 callsite this loop executes exactly once; + // the chunking only kicks in if a future caller passes >500 + // ids in a single batch. + for window in chunk_ids.chunks(MAX_EMBEDDING_BATCH) { + // Build `IN (?,?,?,...)` with `window.len()` placeholders. + // model_signature is bound as the last parameter (?{n+1}). + let placeholders = std::iter::repeat_n("?", window.len()) + .collect::>() + .join(","); + let sql = format!( + "SELECT chunk_id, vector, dim + FROM mem_tree_chunk_embeddings + WHERE chunk_id IN ({placeholders}) + AND model_signature = ?{sig_idx}", + sig_idx = window.len() + 1, + ); + let mut stmt = conn + .prepare(&sql) + .context("prepare get_chunk_embeddings_for_signature_batch")?; + // Bind chunk_ids then model_signature. rusqlite wants + // ToSql trait objects in a uniform iterator. + let mut params: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(window.len() + 1); + for id in window { + params.push(id as &dyn rusqlite::ToSql); + } + params.push(&model_signature as &dyn rusqlite::ToSql); + let rows = stmt + .query_map(params.as_slice(), |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, Vec>(1)?, + row.get::<_, i64>(2)?, + )) + }) + .context("query get_chunk_embeddings_for_signature_batch")?; + for row in rows { + let (chunk_id, bytes, dim) = row?; + // Reuse the single-row decoder so corrupt-blob errors + // surface with the same diagnostic shape as the + // single-row path. + if let Some(v) = embedding_from_blob(&bytes, dim, "chunk embedding")? { + out.insert(chunk_id, v); + } + } + } + Ok(out) + }) +} + +/// Batched read of chunk embeddings under the **active** model signature. +/// Convenience wrapper mirroring [`get_chunk_embedding`] for the per-row +/// path: resolves `tree_active_signature(config)` exactly once and forwards +/// to [`get_chunk_embeddings_for_signature_batch`]. +pub fn get_chunk_embeddings_batch( + config: &Config, + chunk_ids: &[String], +) -> Result>> { + let signature = tree_active_signature(config); + get_chunk_embeddings_for_signature_batch(config, chunk_ids, &signature) +} diff --git a/src/openhuman/memory_store/chunks/store.rs b/src/openhuman/memory_store/chunks/store.rs index c03d39743..54277643b 100644 --- a/src/openhuman/memory_store/chunks/store.rs +++ b/src/openhuman/memory_store/chunks/store.rs @@ -1291,7 +1291,8 @@ fn add_column_if_missing(conn: &Connection, table: &str, name: &str, sql_type: & mod embeddings; pub use embeddings::{ clear_chunk_reembed_skipped, clear_reembed_skipped_for_signature, get_chunk_embedding, - get_chunk_embedding_for_signature, mark_chunk_reembed_skipped, set_chunk_embedding, + get_chunk_embedding_for_signature, get_chunk_embeddings_batch, + get_chunk_embeddings_for_signature_batch, mark_chunk_reembed_skipped, set_chunk_embedding, set_chunk_embedding_for_signature, }; #[cfg(test)] @@ -1300,6 +1301,7 @@ pub(crate) use embeddings::{ has_uncovered_reembed_work, set_chunk_embedding_for_signature_tx, tree_active_signature, validate_reembed_skip_key, }; +// ── Phase 2: embedding column accessors ───────────────────────────────── #[cfg(test)] #[path = "store_tests.rs"] diff --git a/src/openhuman/memory_store/chunks/store_tests.rs b/src/openhuman/memory_store/chunks/store_tests.rs index b6d963b26..b04fea460 100644 --- a/src/openhuman/memory_store/chunks/store_tests.rs +++ b/src/openhuman/memory_store/chunks/store_tests.rs @@ -962,3 +962,100 @@ fn validate_reembed_skip_key_rejects_empty_and_oversized() { "trimmed" ); } + +// ---------- get_chunk_embeddings_for_signature_batch ---------- +// +// Contract: equivalent to looping `get_chunk_embedding_for_signature` +// per id, but in O(ceil(n / MAX_EMBEDDING_BATCH)) round-trips instead +// of O(n). The map contains only ids that have a vector under the +// requested signature; absent rows are silently dropped (same as the +// per-row helper returning Ok(None)). + +#[test] +fn batch_embedding_lookup_returns_only_signature_scoped_rows() { + let (_tmp, cfg) = test_config(); + let c1 = sample_chunk("slack:#eng", 0, 1_700_000_000_000); + let c2 = sample_chunk("slack:#eng", 1, 1_700_000_000_000); + let c3 = sample_chunk("slack:#eng", 2, 1_700_000_000_000); + upsert_chunks(&cfg, &[c1.clone(), c2.clone(), c3.clone()]).unwrap(); + + let sig_a = "openai/text-embedding-3-small@1536"; + let sig_b = "local/bge-small@384"; + set_chunk_embedding_for_signature(&cfg, &c1.id, sig_a, &[0.1, 0.2]).unwrap(); + set_chunk_embedding_for_signature(&cfg, &c2.id, sig_a, &[0.3, 0.4]).unwrap(); + set_chunk_embedding_for_signature(&cfg, &c3.id, sig_b, &[0.5, 0.6, 0.7]).unwrap(); + + let ids = vec![c1.id.clone(), c2.id.clone(), c3.id.clone()]; + let map_a = get_chunk_embeddings_for_signature_batch(&cfg, &ids, sig_a).unwrap(); + assert_eq!(map_a.len(), 2, "only c1 and c2 are under sig_a"); + assert_eq!(map_a.get(&c1.id).cloned(), Some(vec![0.1, 0.2])); + assert_eq!(map_a.get(&c2.id).cloned(), Some(vec![0.3, 0.4])); + assert!(map_a.get(&c3.id).is_none(), "c3 has only sig_b"); + + let map_b = get_chunk_embeddings_for_signature_batch(&cfg, &ids, sig_b).unwrap(); + assert_eq!(map_b.len(), 1); + assert_eq!(map_b.get(&c3.id).cloned(), Some(vec![0.5, 0.6, 0.7])); +} + +#[test] +fn batch_embedding_lookup_empty_input_returns_empty_map() { + let (_tmp, cfg) = test_config(); + let map = get_chunk_embeddings_for_signature_batch(&cfg, &[], "any/sig@1").unwrap(); + assert!(map.is_empty()); +} + +#[test] +fn batch_embedding_lookup_unknown_ids_absent_from_map() { + // Pre-batch contract: per-row helper returned Ok(None) for missing + // chunks. Batch helper must mirror that — missing ids absent from + // the map, present ids carry their vector. The retrieval rerank + // path depends on this so absent rows get the + // (NEG_INFINITY, false) sink-to-bottom treatment. + let (_tmp, cfg) = test_config(); + let c = sample_chunk("slack:#eng", 0, 1_700_000_000_000); + upsert_chunks(&cfg, &[c.clone()]).unwrap(); + let sig = "openai/text-embedding-3-small@1536"; + set_chunk_embedding_for_signature(&cfg, &c.id, sig, &[0.1]).unwrap(); + + let ids = vec![ + c.id.clone(), + "ghost:no-such-chunk-1".into(), + "ghost:no-such-chunk-2".into(), + ]; + let map = get_chunk_embeddings_for_signature_batch(&cfg, &ids, sig).unwrap(); + assert_eq!(map.len(), 1); + assert_eq!(map.get(&c.id).cloned(), Some(vec![0.1])); +} + +#[test] +fn batch_embedding_lookup_splits_id_list_above_per_batch_threshold() { + // Validates the `chunks(MAX_EMBEDDING_BATCH)` window loop in + // `get_chunk_embeddings_for_signature_batch`. We pass > 500 ids in + // one call; the helper must internally split them into multiple + // `IN (...)` queries and merge results into a single map. 3 of the + // 501 ids actually carry embeddings; the other 498 are unknown + // strings and must be absent from the returned map (no error). + let (_tmp, cfg) = test_config(); + let c1 = sample_chunk("slack:#a", 0, 1_700_000_000_000); + let c2 = sample_chunk("slack:#b", 0, 1_700_000_000_000); + let c3 = sample_chunk("slack:#c", 0, 1_700_000_000_000); + upsert_chunks(&cfg, &[c1.clone(), c2.clone(), c3.clone()]).unwrap(); + let sig = "openai/text-embedding-3-small@1536"; + set_chunk_embedding_for_signature(&cfg, &c1.id, sig, &[1.0]).unwrap(); + set_chunk_embedding_for_signature(&cfg, &c2.id, sig, &[2.0]).unwrap(); + set_chunk_embedding_for_signature(&cfg, &c3.id, sig, &[3.0]).unwrap(); + + // Build 501 ids: 3 real + 498 ghosts. The 501-element vec crosses + // the 500-per-batch boundary, forcing two `IN (...)` queries. + let mut ids: Vec = (0..498).map(|i| format!("ghost:{i}")).collect(); + ids.push(c1.id.clone()); + ids.push(c2.id.clone()); + ids.push(c3.id.clone()); + assert_eq!(ids.len(), 501); + + let map = get_chunk_embeddings_for_signature_batch(&cfg, &ids, sig).unwrap(); + assert_eq!(map.len(), 3, "only the 3 real ids should be present"); + assert_eq!(map.get(&c1.id).cloned(), Some(vec![1.0])); + assert_eq!(map.get(&c2.id).cloned(), Some(vec![2.0])); + assert_eq!(map.get(&c3.id).cloned(), Some(vec![3.0])); +} diff --git a/src/openhuman/memory_store/trees/store.rs b/src/openhuman/memory_store/trees/store.rs index 3ca4e54bf..b74041a10 100644 --- a/src/openhuman/memory_store/trees/store.rs +++ b/src/openhuman/memory_store/trees/store.rs @@ -451,6 +451,108 @@ pub fn get_summary_embedding_for_signature( }) } +/// Per-batch cap on `?` placeholders. Mirrors `chunks::store:: +/// MAX_EMBEDDING_BATCH` — see that constant's doc for the rationale (well +/// below SQLite's `SQLITE_MAX_VARIABLE_NUMBER = 32766`, large enough that +/// the current `LOOKUP_HEADROOM = 200` callsite always fits in one +/// round-trip). The two sides are independent intentionally: the summary +/// and chunk tables can grow at different rates and the cap might want to +/// drift independently in the future. +const MAX_EMBEDDING_BATCH: usize = 500; + +/// Batched read of summary embeddings under a single `model_signature`. +/// +/// Returns a `HashMap>` containing **only the +/// summaries that have a vector under `model_signature`**. Summaries with +/// no row, with a `NULL` vector (pending re-embed), or with a corrupted +/// blob are simply absent from the map — semantically identical to the +/// per-row [`get_summary_embedding_for_signature`] returning `Ok(None)`. +/// +/// Mirror of `chunks::store::get_chunk_embeddings_for_signature_batch`. +/// See that helper's doc for the rerank-loop motivation. The summary +/// side has its own copy rather than a generic helper because the two +/// tables (`mem_tree_summary_embeddings` vs `mem_tree_chunk_embeddings`) +/// have different blob-nullability semantics: summaries can store an +/// explicit `NULL` vector to flag a pending re-embed (handled here via +/// `Option>` + `decode_signature_blob`), chunks cannot. +pub fn get_summary_embeddings_for_signature_batch( + config: &Config, + summary_ids: &[String], + model_signature: &str, +) -> Result>> { + if summary_ids.is_empty() { + return Ok(HashMap::new()); + } + with_connection(config, |conn| { + let mut out: HashMap> = HashMap::with_capacity(summary_ids.len()); + // Chunk to stay under SQLite's SQLITE_MAX_VARIABLE_NUMBER cap. + // For LOOKUP_HEADROOM=200 this loop runs exactly once; chunking + // only engages if a future caller passes >500 ids at a time. + for window in summary_ids.chunks(MAX_EMBEDDING_BATCH) { + // Build `IN (?,?,?,...)` with `window.len()` placeholders. + // model_signature is bound as the last parameter (?{n+1}). + let placeholders = std::iter::repeat_n("?", window.len()) + .collect::>() + .join(","); + let sql = format!( + "SELECT summary_id, vector, dim + FROM mem_tree_summary_embeddings + WHERE summary_id IN ({placeholders}) + AND model_signature = ?{sig_idx}", + sig_idx = window.len() + 1, + ); + let mut stmt = conn + .prepare(&sql) + .context("prepare get_summary_embeddings_for_signature_batch")?; + let mut bound: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(window.len() + 1); + for id in window { + bound.push(id as &dyn rusqlite::ToSql); + } + bound.push(&model_signature as &dyn rusqlite::ToSql); + let rows = stmt + .query_map(bound.as_slice(), |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, Option>>(1)?, + row.get::<_, i64>(2)?, + )) + }) + .context("query get_summary_embeddings_for_signature_batch")?; + for row in rows { + let (summary_id, blob, dim) = row?; + // Reuse the single-row decoder so NULL vectors (pending + // re-embed) and corrupt blobs surface with identical + // diagnostics to the per-row path. `Ok(None)` from the + // decoder is dropped: the map only carries materialised + // vectors, exactly mirroring the existing per-row + // contract. Length / dim-mismatch / negative-dim / + // non-multiple-of-4 are already enforced inside + // `decode_signature_blob` itself — no extra check here, + // matching the chunks side which delegates the same way + // to `embedding_from_blob`. + if let Some(v) = + decode_signature_blob(blob, dim, &format!("summary_id={summary_id}"))? + { + out.insert(summary_id, v); + } + } + } + Ok(out) + }) +} + +/// Batched read of summary embeddings under the **active** model +/// signature. Mirrors [`get_summary_embedding`] for the per-row path: +/// resolves `tree_active_signature` once, forwards to +/// [`get_summary_embeddings_for_signature_batch`]. +pub fn get_summary_embeddings_batch( + config: &Config, + summary_ids: &[String], +) -> Result>> { + let signature = crate::openhuman::memory_store::chunks::store::tree_active_signature(config); + get_summary_embeddings_for_signature_batch(config, summary_ids, &signature) +} + /// Fetch one summary by id. Soft-deleted rows are returned with /// `deleted = true` so callers can decide filtering policy. pub fn get_summary(config: &Config, id: &str) -> Result> { diff --git a/src/openhuman/memory_store/trees/store_tests.rs b/src/openhuman/memory_store/trees/store_tests.rs index 7ab9625c9..2e700b7ad 100644 --- a/src/openhuman/memory_store/trees/store_tests.rs +++ b/src/openhuman/memory_store/trees/store_tests.rs @@ -347,3 +347,87 @@ fn get_summaries_batch_empty_input_and_missing_ids() { assert_eq!(map.get("sum-a").unwrap(), &a); assert!(map.get("ghost:no-such").is_none()); } + +// ---------- get_summary_embeddings_for_signature_batch ---------- +// +// Contract mirror of the chunks-side batch helper: equivalent to looping +// `get_summary_embedding_for_signature` per id, but in +// O(ceil(n / MAX_EMBEDDING_BATCH)) round-trips instead of O(n). The map +// contains only ids that have a non-null vector under the requested +// signature; absent rows (no sidecar entry, or sidecar entry with NULL +// vector) are silently dropped (same as the per-row helper returning +// Ok(None)). Chunking-window behaviour is covered on the chunks side +// (`batch_embedding_lookup_splits_id_list_above_per_batch_threshold`); +// the implementations share the same `chunks(MAX_EMBEDDING_BATCH)` loop +// shape so re-validating it here would be pure duplication. + +fn seed_summary(cfg: &Config, tree_id: &str, summary_id: &str) { + insert_tree(cfg, &sample_tree(tree_id, &format!("scope:{tree_id}"))).ok(); + let node = sample_summary(summary_id, tree_id, 1); + with_connection(cfg, |conn| { + let tx = conn.unchecked_transaction()?; + insert_summary_tx(&tx, &node, None, "test")?; + tx.commit()?; + Ok(()) + }) + .unwrap(); +} + +#[test] +fn summary_batch_embedding_lookup_returns_only_signature_scoped_rows() { + let (_tmp, cfg) = test_config(); + seed_summary(&cfg, "tree-1", "sum-1"); + seed_summary(&cfg, "tree-1", "sum-2"); + seed_summary(&cfg, "tree-1", "sum-3"); + + let sig_a = "openai/text-embedding-3-small@1536"; + let sig_b = "local/bge-small@384"; + set_summary_embedding_for_signature(&cfg, "sum-1", sig_a, &[0.1, 0.2]).unwrap(); + set_summary_embedding_for_signature(&cfg, "sum-2", sig_a, &[0.3, 0.4]).unwrap(); + set_summary_embedding_for_signature(&cfg, "sum-3", sig_b, &[0.5, 0.6, 0.7]).unwrap(); + + let ids = vec![ + "sum-1".to_string(), + "sum-2".to_string(), + "sum-3".to_string(), + ]; + let map_a = get_summary_embeddings_for_signature_batch(&cfg, &ids, sig_a).unwrap(); + assert_eq!(map_a.len(), 2, "only sum-1 and sum-2 are under sig_a"); + assert_eq!(map_a.get("sum-1").cloned(), Some(vec![0.1, 0.2])); + assert_eq!(map_a.get("sum-2").cloned(), Some(vec![0.3, 0.4])); + assert!(map_a.get("sum-3").is_none(), "sum-3 has only sig_b"); + + let map_b = get_summary_embeddings_for_signature_batch(&cfg, &ids, sig_b).unwrap(); + assert_eq!(map_b.len(), 1); + assert_eq!(map_b.get("sum-3").cloned(), Some(vec![0.5, 0.6, 0.7])); +} + +#[test] +fn summary_batch_embedding_lookup_empty_input_returns_empty_map() { + let (_tmp, cfg) = test_config(); + let map = get_summary_embeddings_for_signature_batch(&cfg, &[], "any/sig@1").unwrap(); + assert!(map.is_empty()); +} + +#[test] +fn summary_batch_embedding_lookup_unknown_ids_absent_from_map() { + // Pre-batch contract: per-row helper returned Ok(None) for missing + // summaries OR for summaries whose sidecar row has a NULL vector + // (pending re-embed). The batch helper must mirror that — missing + // ids absent from the map, present ids carry their vector. The + // retrieval rerank path depends on this so absent rows get the + // (NEG_INFINITY, false) sink-to-bottom treatment. + let (_tmp, cfg) = test_config(); + seed_summary(&cfg, "tree-1", "sum-1"); + let sig = "openai/text-embedding-3-small@1536"; + set_summary_embedding_for_signature(&cfg, "sum-1", sig, &[0.1]).unwrap(); + + let ids = vec![ + "sum-1".to_string(), + "ghost:no-such-summary-1".to_string(), + "ghost:no-such-summary-2".to_string(), + ]; + let map = get_summary_embeddings_for_signature_batch(&cfg, &ids, sig).unwrap(); + assert_eq!(map.len(), 1); + assert_eq!(map.get("sum-1").cloned(), Some(vec![0.1])); +} diff --git a/src/openhuman/memory_tree/retrieval/topic.rs b/src/openhuman/memory_tree/retrieval/topic.rs index e13900a18..4cac9c3eb 100644 --- a/src/openhuman/memory_tree/retrieval/topic.rs +++ b/src/openhuman/memory_tree/retrieval/topic.rs @@ -175,9 +175,9 @@ async fn rerank_by_semantic_similarity( query: &str, hits: Vec, ) -> Result> { - use crate::openhuman::memory_store::chunks::store::get_chunk_embedding; + use crate::openhuman::memory_store::chunks::store::get_chunk_embeddings_batch; + use crate::openhuman::memory_store::trees::store::get_summary_embeddings_batch; use crate::openhuman::memory_tree::retrieval::types::NodeKind; - use crate::openhuman::memory_tree::tree::store as src_store; let embedder = build_embedder_from_config(config)?; let query_vec = embedder.embed(query).await?; @@ -187,28 +187,69 @@ async fn rerank_by_semantic_similarity( hits.len() ); - // Resolve each hit's embedding. spawn_blocking around the DB reads - // so the event loop stays healthy even for larger headroom pulls. - let mut decorated: Vec<(f32, bool, RetrievalHit)> = Vec::with_capacity(hits.len()); - for h in hits { - let node_id = h.node_id.clone(); - let node_kind = h.node_kind; - let config_owned = config.clone(); - let emb = tokio::task::spawn_blocking(move || -> Result>> { - match node_kind { - NodeKind::Summary => src_store::get_summary_embedding(&config_owned, &node_id), - NodeKind::Leaf => get_chunk_embedding(&config_owned, &node_id), - } + // Partition hit ids by node kind so each table gets a single batched + // `IN (...)` lookup. Summary embeddings live in + // `mem_tree_summary_embeddings`, leaf/chunk embeddings in + // `mem_tree_chunk_embeddings` — two tables, two batched queries. + // + // Why partition + batch instead of one query per hit: + // + // Previously this function looped over `hits` and ran + // `spawn_blocking(get_*_embedding).await` per element. That `.await` + // inside the `for` was *sequential* — N round-trips to SQLite, each + // paying its own prepare + bind + busy-wait. With LOOKUP_HEADROOM=200 + // the rerank path could fire 200 sequential SQL statements on every + // entity-scoped query carrying a `query=` arg, which is the common + // per-turn shape. + // + // The batched helpers (see `chunks::store:: + // get_chunk_embeddings_for_signature_batch` and `trees::store:: + // get_summary_embeddings_for_signature_batch`) collapse all + // same-kind lookups into one `IN (?,?,?,...) AND model_signature = ?` + // query (chunked internally to stay below SQLite's variable cap so + // large future headroom values stay safe). The hit list is decorated + // in its original order from the resulting maps, so sort stability, + // tie-break behaviour, and the existing `(NEG_INFINITY, false)` + // handling for missing embeddings are all preserved bit-for-bit. + let mut summary_ids: Vec = Vec::new(); + let mut chunk_ids: Vec = Vec::new(); + for h in &hits { + match h.node_kind { + NodeKind::Summary => summary_ids.push(h.node_id.clone()), + NodeKind::Leaf => chunk_ids.push(h.node_id.clone()), + } + } + + // Both fetches run under one `spawn_blocking` to keep the event loop + // free while the SQLite reads happen on the blocking pool. + let config_owned = config.clone(); + let (summary_embeddings, chunk_embeddings) = + tokio::task::spawn_blocking(move || -> Result<(_, _)> { + let s = get_summary_embeddings_batch(&config_owned, &summary_ids)?; + let c = get_chunk_embeddings_batch(&config_owned, &chunk_ids)?; + Ok((s, c)) }) .await - .map_err(|e| anyhow::anyhow!("embedding fetch join error: {e}"))??; + .map_err(|e| anyhow::anyhow!("embedding batch join error: {e}"))??; - match emb { + let mut decorated: Vec<(f32, bool, RetrievalHit)> = Vec::with_capacity(hits.len()); + for h in hits { + // Decorate in the original `hits` iteration order so two hits + // that tie on every ranked dimension still produce the same + // relative ordering as before this refactor. + let emb_lookup = match h.node_kind { + NodeKind::Summary => summary_embeddings.get(&h.node_id), + NodeKind::Leaf => chunk_embeddings.get(&h.node_id), + }; + match emb_lookup { Some(v) => { - let sim = cosine_similarity(&query_vec, &v); + let sim = cosine_similarity(&query_vec, v); decorated.push((sim, true, h)); } None => { + // Identical to the pre-batch path: absent embedding + // sinks the hit to the bottom of the rerank without + // dropping it from the result set. decorated.push((f32::NEG_INFINITY, false, h)); } }