perf(memory_tree): batch embedding fetch in topic rerank (eliminate N sequential SQLite round-trips) (#2927)

Co-authored-by: oxoxDev <nikhil@tinyhumans.ai>

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<Option<Vec
     }
     Ok(Some(floats))
 }
+/// SQLite's compile-time hard cap on `?` parameters per prepared statement is
+/// `SQLITE_MAX_VARIABLE_NUMBER` (32 766 since SQLite 3.32, bundled with
+/// rusqlite). We chunk well below that for two reasons:
+///
+/// 1. **Headroom against schema growth.** The batched query binds one `?` per
+///    chunk_id *plus* the `model_signature` parameter. Picking 500 as the
+///    per-chunk cap gives ~65× safety margin even if a future caller bumps
+///    `LOOKUP_HEADROOM` from 200 into the thousands — the batch helper
+///    silently splits the request, the caller sees no semantic change.
+/// 2. **Sane SQL string length.** A 500-`?` `IN(...)` clause is ~1 KB of SQL
+///    text — small enough that prepare-and-discard per chunk is cheap and the
+///    planner produces a simple covered-index lookup.
+///
+/// Lowering this constant is safe (just more round-trips). Raising it past
+/// the SQLite limit will fail at prepare time with `too many SQL variables`.
+const MAX_EMBEDDING_BATCH: usize = 500;
+
+/// Batched read of chunk embeddings under a single `model_signature`.
+///
+/// Returns a `HashMap<chunk_id, Vec<f32>>` 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<Vec<f32>>` 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<HashMap<String, Vec<f32>>> {
+    if chunk_ids.is_empty() {
+        return Ok(HashMap::new());
+    }
+    with_connection(config, |conn| {
+        let mut out: HashMap<String, Vec<f32>> = 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::<Vec<_>>()
+                .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<u8>>(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<HashMap<String, Vec<f32>>> {
+    let signature = tree_active_signature(config);
+    get_chunk_embeddings_for_signature_batch(config, chunk_ids, &signature)
+}

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"]

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<String> = (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]));
+}

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<summary_id, Vec<f32>>` 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<Vec<u8>>` + `decode_signature_blob`), chunks cannot.
+pub fn get_summary_embeddings_for_signature_batch(
+    config: &Config,
+    summary_ids: &[String],
+    model_signature: &str,
+) -> Result<HashMap<String, Vec<f32>>> {
+    if summary_ids.is_empty() {
+        return Ok(HashMap::new());
+    }
+    with_connection(config, |conn| {
+        let mut out: HashMap<String, Vec<f32>> = 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::<Vec<_>>()
+                .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<Vec<u8>>>(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<HashMap<String, Vec<f32>>> {
+    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<Option<SummaryNode>> {

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]));
+}

src/openhuman/memory_tree/retrieval/topic.rs

@@ -175,9 +175,9 @@ async fn rerank_by_semantic_similarity(
     query: &str,
     hits: Vec<RetrievalHit>,
 ) -> Result<Vec<RetrievalHit>> {
-    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<Option<Vec<f32>>> {
-            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<String> = Vec::new();
+    let mut chunk_ids: Vec<String> = 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));
             }
         }