perf(agent): store embeddings once, not twice
MemoryCache held every embedding in two places: a `Vec<Vec<f32>>` and a flattened copy for the batched kernels, kept in lock-step on every push, update and compaction. A store loaded from disk therefore carried the corpus twice, plus one heap allocation per entry. A new `cache::Embeddings` owns just the flat `[N x dim]` buffer and indexes into it, so `embeddings[i]` still reads as a `&[f32]` row. The batch kernels take a `VectorSet` (implemented for both `Embeddings` and `Vec<Vec<f32>>`) instead of `&[Vec<f32>]`, so their callers and tests are unchanged. Loading no longer unflattens what it just read. 100k 384-dim entries, reopened from disk: 505 -> 357 MiB, 3.44x -> 2.43x the raw vectors. Recall (1.0000 at ef=64) and query latency are unchanged. Rows are now always exactly `dim` long, shorter ones zero-padded. The old representation allowed ragged rows, which silently misaligned the flattened copy — every row after a wrong-length embedding — and `update` carried a comment about falling back to a rebuild to avoid exactly that. It is now unrepresentable. A record saved without an embedding holds a zero row and is told apart by its norm, which is what `total_embeddings` now counts. Measured with a counting allocator rather than RSS: freeing a structure returns its pages to the allocator's pool, not the OS, so an RSS reading from inside the process showed the two representations as identical. Breaking: MemoryCache::embeddings changes type, embeddings_flat is replaced by flat_embeddings(), rebuild_flat() is a deprecated no-op. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
@@ -28,7 +28,7 @@ use crate::vector_search;
|
||||
pub fn hybrid_search(
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||
chunks: &[String],
|
||||
tombstones: &[u8],
|
||||
bm25_index: &BM25Index,
|
||||
@@ -56,7 +56,7 @@ pub fn hybrid_search(
|
||||
pub fn hybrid_search_fused(
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||
_chunks: &[String],
|
||||
tombstones: &[u8],
|
||||
bm25_index: &BM25Index,
|
||||
@@ -69,12 +69,12 @@ pub fn hybrid_search_fused(
|
||||
let vec_scores = {
|
||||
#[cfg(feature = "parallel")]
|
||||
{
|
||||
if vectors.len() > 10_000 {
|
||||
if vectors.count() > 10_000 {
|
||||
vector_search::parallel_cosine_batch(
|
||||
query_embedding,
|
||||
vectors,
|
||||
tombstones,
|
||||
vectors.len(),
|
||||
vectors.count(),
|
||||
)
|
||||
} else {
|
||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||
@@ -270,7 +270,7 @@ fn normalize_scores(scores: &[(usize, f32)]) -> Vec<(usize, f32)> {
|
||||
pub fn rrf_hybrid_search(
|
||||
query_embedding: &[f32],
|
||||
query_text: &str,
|
||||
vectors: &[Vec<f32>],
|
||||
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||
_chunks: &[String],
|
||||
tombstones: &[u8],
|
||||
bm25_index: &BM25Index,
|
||||
@@ -282,12 +282,12 @@ pub fn rrf_hybrid_search(
|
||||
let mut vec_scores = {
|
||||
#[cfg(feature = "parallel")]
|
||||
{
|
||||
if vectors.len() > 10_000 {
|
||||
if vectors.count() > 10_000 {
|
||||
vector_search::parallel_cosine_batch(
|
||||
query_embedding,
|
||||
vectors,
|
||||
tombstones,
|
||||
vectors.len(),
|
||||
vectors.count(),
|
||||
)
|
||||
} else {
|
||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||
@@ -298,7 +298,7 @@ pub fn rrf_hybrid_search(
|
||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||
}
|
||||
};
|
||||
let mut kw_scores = bm25_index.search(query_text, vectors.len());
|
||||
let mut kw_scores = bm25_index.search(query_text, vectors.count());
|
||||
|
||||
// Sort both lists descending so rank 1 = best.
|
||||
vec_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
Reference in New Issue
Block a user