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,6 +28,35 @@
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
## Memory footprint
|
||||||
|
|
||||||
|
`cargo run --release -p clawhdf5-bench --bin search_harness -- --footprint --full`,
|
||||||
|
384-dim `f32`. The figure that matters is **reopened**: a store loaded from
|
||||||
|
disk, which is what a long-lived process holds.
|
||||||
|
|
||||||
|
Measured with a counting global allocator, not RSS. RSS cannot see this from
|
||||||
|
inside one process — freeing a large structure returns its pages to the
|
||||||
|
allocator's pool rather than to the OS, so allocating the next one shows no
|
||||||
|
change at all. Measured that way a store holding the corpus twice and one
|
||||||
|
holding it once came out *identical* (1.00x both), which is how the first
|
||||||
|
attempt at this measurement went.
|
||||||
|
|
||||||
|
| N | vectors (raw) | reopened, before | reopened, after |
|
||||||
|
|---:|---:|---:|---:|
|
||||||
|
| 1 000 | 1 MiB | 5 MiB (3.41x) | 4 MiB (2.39x) |
|
||||||
|
| 10 000 | 15 MiB | 50 MiB (3.43x) | 35 MiB (2.42x) |
|
||||||
|
| 100 000 | 146 MiB | 505 MiB (3.44x) | **357 MiB (2.43x)** |
|
||||||
|
|
||||||
|
The cache stored every embedding twice — once as a `Vec<Vec<f32>>` and once
|
||||||
|
flattened for the batched kernels, kept in lock-step on every push, update and
|
||||||
|
compaction. Storing only the flat buffer and indexing into it gives back
|
||||||
|
almost exactly one copy of the corpus (148 MiB at 100k) and one heap
|
||||||
|
allocation per entry. Recall and query latency are unchanged.
|
||||||
|
|
||||||
|
What remains at 2.43x: the flat vectors (1.0x), the HNSW index's own copy of
|
||||||
|
them (1.0x), and text, ids and graph (~0.4x). The index copy is the next
|
||||||
|
target — it is what a quantised or borrowed representation would address.
|
||||||
|
|
||||||
## Read harness
|
## Read harness
|
||||||
|
|
||||||
Produced by `cargo run --release -p clawhdf5-bench --bin read_harness`: a 4096 x
|
Produced by `cargo run --release -p clawhdf5-bench --bin read_harness`: a 4096 x
|
||||||
|
|||||||
@@ -2,6 +2,24 @@
|
|||||||
|
|
||||||
## Unreleased
|
## Unreleased
|
||||||
|
|
||||||
|
### Memory
|
||||||
|
- `clawhdf5-agent`: **a loaded store holds ~30% less memory** (100k 384-dim
|
||||||
|
entries: 505 -> 357 MiB, 3.44x -> 2.43x the raw vectors). The cache kept
|
||||||
|
every embedding twice — a `Vec<Vec<f32>>` and a flattened copy for the
|
||||||
|
batched kernels, maintained in lock-step — so it now stores only the flat
|
||||||
|
buffer and indexes into it. Recall and query latency are unchanged.
|
||||||
|
**Breaking:** `MemoryCache::embeddings` is a `cache::Embeddings` rather than
|
||||||
|
a `Vec<Vec<f32>>` (indexing still yields a `&[f32]` row); `embeddings_flat`
|
||||||
|
is gone, replaced by `flat_embeddings()`; `rebuild_flat()` is a deprecated
|
||||||
|
no-op. Rows are now always exactly `dim` long — shorter ones are
|
||||||
|
zero-padded — which makes the ragged-row case that used to silently
|
||||||
|
misalign the flattened copy unrepresentable.
|
||||||
|
- `clawhdf5-bench`: `search_harness --footprint` reports live heap use per
|
||||||
|
stage, measured with a counting allocator (RSS cannot see a structure freed
|
||||||
|
into the allocator's own pool).
|
||||||
|
|
||||||
|
## Unreleased
|
||||||
|
|
||||||
### Retrieval quality
|
### Retrieval quality
|
||||||
- `clawhdf5-agent`: **re-ranking discarded the retrieval score.**
|
- `clawhdf5-agent`: **re-ranking discarded the retrieval score.**
|
||||||
`reranker::rerank` built its combined score from temporal decay, source
|
`reranker::rerank` built its combined score from temporal decay, source
|
||||||
|
|||||||
@@ -2,16 +2,143 @@
|
|||||||
|
|
||||||
use crate::vector_search;
|
use crate::vector_search;
|
||||||
|
|
||||||
|
/// Every entry's embedding, in one contiguous `[N x dim]` buffer.
|
||||||
|
///
|
||||||
|
/// Rows are always exactly `dim` long: a shorter one is zero-padded, a longer
|
||||||
|
/// one truncated. The previous `Vec<Vec<f32>>` allowed ragged rows, which
|
||||||
|
/// silently misaligned the flattened copy that the batched kernels read — a
|
||||||
|
/// single wrong-length embedding shifted every row after it. Padding makes
|
||||||
|
/// that unrepresentable. A record stored without an embedding therefore holds
|
||||||
|
/// a zero row, and is told apart by its norm being zero rather than by length.
|
||||||
|
///
|
||||||
|
/// This used to be two fields — a `Vec<Vec<f32>>` and a flattened copy kept in
|
||||||
|
/// lock-step — which stored the whole corpus twice and cost one heap
|
||||||
|
/// allocation per entry on top. At 100k 384-dim entries that duplicate was
|
||||||
|
/// ~150 MiB. Indexing yields a `&[f32]` row, so `embeddings[i]` still reads
|
||||||
|
/// the same way.
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct Embeddings {
|
||||||
|
flat: Vec<f32>,
|
||||||
|
dim: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Embeddings {
|
||||||
|
pub fn new(dim: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
flat: Vec::new(),
|
||||||
|
dim,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Number of embeddings.
|
||||||
|
pub fn len(&self) -> usize {
|
||||||
|
self.flat.len().checked_div(self.dim).unwrap_or(0)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.len() == 0
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The whole buffer, `[N x dim]` row-major — what batched kernels read.
|
||||||
|
pub fn as_flat(&self) -> &[f32] {
|
||||||
|
&self.flat
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn dim(&self) -> usize {
|
||||||
|
self.dim
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Row `i`, or `None` if out of range.
|
||||||
|
pub fn get(&self, i: usize) -> Option<&[f32]> {
|
||||||
|
let start = i.checked_mul(self.dim)?;
|
||||||
|
self.flat.get(start..start.checked_add(self.dim)?)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn iter(&self) -> impl ExactSizeIterator<Item = &[f32]> {
|
||||||
|
self.flat.chunks_exact(self.dim.max(1))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Append one embedding. A row whose length doesn't match `dim` is padded
|
||||||
|
/// or truncated, so the buffer stays rectangular whatever a caller passes.
|
||||||
|
pub fn push(&mut self, embedding: &[f32]) {
|
||||||
|
if self.dim == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let take = embedding.len().min(self.dim);
|
||||||
|
self.flat.extend_from_slice(&embedding[..take]);
|
||||||
|
self.flat.resize(self.flat.len() + (self.dim - take), 0.0);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Replace row `i`. Out-of-range indices are ignored.
|
||||||
|
pub fn set(&mut self, i: usize, embedding: &[f32]) {
|
||||||
|
let Some(start) = i.checked_mul(self.dim) else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
if start + self.dim > self.flat.len() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let take = embedding.len().min(self.dim);
|
||||||
|
self.flat[start..start + take].copy_from_slice(&embedding[..take]);
|
||||||
|
self.flat[start + take..start + self.dim].fill(0.0);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Keep only the rows `keep` returns true for, preserving order.
|
||||||
|
pub fn retain(&mut self, mut keep: impl FnMut(usize) -> bool) {
|
||||||
|
if self.dim == 0 {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let mut write = 0usize;
|
||||||
|
for read in 0..self.len() {
|
||||||
|
if keep(read) {
|
||||||
|
if write != read {
|
||||||
|
let (dst, src) = (write * self.dim, read * self.dim);
|
||||||
|
self.flat.copy_within(src..src + self.dim, dst);
|
||||||
|
}
|
||||||
|
write += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
self.flat.truncate(write * self.dim);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Replace the contents with `rows`.
|
||||||
|
pub fn reset_from(&mut self, dim: usize, rows: impl IntoIterator<Item = Vec<f32>>) {
|
||||||
|
self.dim = dim;
|
||||||
|
self.flat.clear();
|
||||||
|
for row in rows {
|
||||||
|
self.push(&row);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Adopt an already-flat buffer, trimming any partial trailing row.
|
||||||
|
pub fn set_flat(&mut self, dim: usize, mut flat: Vec<f32>) {
|
||||||
|
self.dim = dim;
|
||||||
|
match flat.len().checked_div(dim) {
|
||||||
|
Some(rows) => flat.truncate(rows * dim),
|
||||||
|
None => flat.clear(),
|
||||||
|
}
|
||||||
|
self.flat = flat;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PartialEq for Embeddings {
|
||||||
|
fn eq(&self, other: &Self) -> bool {
|
||||||
|
self.dim == other.dim && self.flat == other.flat
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::ops::Index<usize> for Embeddings {
|
||||||
|
type Output = [f32];
|
||||||
|
|
||||||
|
fn index(&self, i: usize) -> &[f32] {
|
||||||
|
self.get(i).expect("embedding index out of range")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// In-memory cache for the /memory group data.
|
/// In-memory cache for the /memory group data.
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct MemoryCache {
|
pub struct MemoryCache {
|
||||||
pub chunks: Vec<String>,
|
pub chunks: Vec<String>,
|
||||||
pub embeddings: Vec<Vec<f32>>,
|
pub embeddings: Embeddings,
|
||||||
/// `embeddings` flattened into one contiguous `[N × embedding_dim]`
|
|
||||||
/// buffer, maintained incrementally alongside `embeddings` (push/update/
|
|
||||||
/// compact) so BLAS/Accelerate batch search can read it directly instead
|
|
||||||
/// of re-flattening the whole corpus on every query.
|
|
||||||
pub embeddings_flat: Vec<f32>,
|
|
||||||
pub source_channels: Vec<String>,
|
pub source_channels: Vec<String>,
|
||||||
pub timestamps: Vec<f64>,
|
pub timestamps: Vec<f64>,
|
||||||
pub session_ids: Vec<String>,
|
pub session_ids: Vec<String>,
|
||||||
@@ -28,8 +155,7 @@ impl MemoryCache {
|
|||||||
pub fn new(embedding_dim: usize) -> Self {
|
pub fn new(embedding_dim: usize) -> Self {
|
||||||
Self {
|
Self {
|
||||||
chunks: Vec::new(),
|
chunks: Vec::new(),
|
||||||
embeddings: Vec::new(),
|
embeddings: Embeddings::new(embedding_dim),
|
||||||
embeddings_flat: Vec::new(),
|
|
||||||
source_channels: Vec::new(),
|
source_channels: Vec::new(),
|
||||||
timestamps: Vec::new(),
|
timestamps: Vec::new(),
|
||||||
session_ids: Vec::new(),
|
session_ids: Vec::new(),
|
||||||
@@ -41,15 +167,14 @@ impl MemoryCache {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Rebuild `embeddings_flat` from `embeddings` from scratch. Callers that
|
/// Kept for callers that used to have to re-flatten after a bulk load.
|
||||||
/// populate `embeddings` directly (bulk loads) must call this afterward.
|
/// The buffer is always flat now, so there is nothing to rebuild.
|
||||||
pub fn rebuild_flat(&mut self) {
|
#[deprecated(note = "embeddings are stored flat; this is a no-op")]
|
||||||
self.embeddings_flat.clear();
|
pub fn rebuild_flat(&mut self) {}
|
||||||
self.embeddings_flat
|
|
||||||
.reserve(self.embeddings.len() * self.embedding_dim);
|
/// The embeddings as one contiguous `[N x dim]` buffer.
|
||||||
for emb in &self.embeddings {
|
pub fn flat_embeddings(&self) -> &[f32] {
|
||||||
self.embeddings_flat.extend_from_slice(emb);
|
self.embeddings.as_flat()
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Total number of entries (including tombstoned).
|
/// Total number of entries (including tombstoned).
|
||||||
@@ -79,8 +204,7 @@ impl MemoryCache {
|
|||||||
let idx = self.chunks.len();
|
let idx = self.chunks.len();
|
||||||
let norm = vector_search::compute_norm(&embedding);
|
let norm = vector_search::compute_norm(&embedding);
|
||||||
self.chunks.push(chunk);
|
self.chunks.push(chunk);
|
||||||
self.embeddings_flat.extend_from_slice(&embedding);
|
self.embeddings.push(&embedding);
|
||||||
self.embeddings.push(embedding);
|
|
||||||
self.source_channels.push(source_channel);
|
self.source_channels.push(source_channel);
|
||||||
self.timestamps.push(timestamp);
|
self.timestamps.push(timestamp);
|
||||||
self.session_ids.push(session_id);
|
self.session_ids.push(session_id);
|
||||||
@@ -118,20 +242,7 @@ impl MemoryCache {
|
|||||||
if idx < self.chunks.len() {
|
if idx < self.chunks.len() {
|
||||||
let norm = vector_search::compute_norm(&embedding);
|
let norm = vector_search::compute_norm(&embedding);
|
||||||
self.chunks[idx] = chunk;
|
self.chunks[idx] = chunk;
|
||||||
let dim = self.embedding_dim;
|
self.embeddings.set(idx, &embedding);
|
||||||
let flat_start = idx * dim;
|
|
||||||
let matches_dim =
|
|
||||||
embedding.len() == dim && flat_start + dim <= self.embeddings_flat.len();
|
|
||||||
self.embeddings[idx] = embedding;
|
|
||||||
if matches_dim {
|
|
||||||
self.embeddings_flat[flat_start..flat_start + dim]
|
|
||||||
.copy_from_slice(&self.embeddings[idx]);
|
|
||||||
} else {
|
|
||||||
// Embedding length doesn't match embedding_dim (shouldn't
|
|
||||||
// happen in practice) — fall back to a full rebuild rather
|
|
||||||
// than leave embeddings_flat misaligned with embeddings.
|
|
||||||
self.rebuild_flat();
|
|
||||||
}
|
|
||||||
self.source_channels[idx] = source_channel;
|
self.source_channels[idx] = source_channel;
|
||||||
self.timestamps[idx] = timestamp;
|
self.timestamps[idx] = timestamp;
|
||||||
self.session_ids[idx] = session_id;
|
self.session_ids[idx] = session_id;
|
||||||
@@ -183,7 +294,7 @@ impl MemoryCache {
|
|||||||
new_idx += 1;
|
new_idx += 1;
|
||||||
let norm = vector_search::compute_norm(&self.embeddings[i]);
|
let norm = vector_search::compute_norm(&self.embeddings[i]);
|
||||||
new_chunks.push(self.chunks[i].clone());
|
new_chunks.push(self.chunks[i].clone());
|
||||||
new_embeddings.push(self.embeddings[i].clone());
|
new_embeddings.push(self.embeddings[i].to_vec());
|
||||||
new_source_channels.push(self.source_channels[i].clone());
|
new_source_channels.push(self.source_channels[i].clone());
|
||||||
new_timestamps.push(self.timestamps[i]);
|
new_timestamps.push(self.timestamps[i]);
|
||||||
new_session_ids.push(self.session_ids[i].clone());
|
new_session_ids.push(self.session_ids[i].clone());
|
||||||
@@ -196,7 +307,8 @@ impl MemoryCache {
|
|||||||
|
|
||||||
let removed = old_len - new_chunks.len();
|
let removed = old_len - new_chunks.len();
|
||||||
self.chunks = new_chunks;
|
self.chunks = new_chunks;
|
||||||
self.embeddings = new_embeddings;
|
self.embeddings
|
||||||
|
.reset_from(self.embedding_dim, new_embeddings);
|
||||||
self.source_channels = new_source_channels;
|
self.source_channels = new_source_channels;
|
||||||
self.timestamps = new_timestamps;
|
self.timestamps = new_timestamps;
|
||||||
self.session_ids = new_session_ids;
|
self.session_ids = new_session_ids;
|
||||||
@@ -204,16 +316,14 @@ impl MemoryCache {
|
|||||||
self.tombstones = new_tombstones;
|
self.tombstones = new_tombstones;
|
||||||
self.norms = new_norms;
|
self.norms = new_norms;
|
||||||
self.activation_weights = new_activation_weights;
|
self.activation_weights = new_activation_weights;
|
||||||
self.rebuild_flat();
|
|
||||||
|
|
||||||
(removed, index_map)
|
(removed, index_map)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Flatten all embeddings into a single Vec<f32> for HDF5 storage.
|
/// All embeddings as one owned `[N x dim]` buffer, for HDF5 storage.
|
||||||
/// `embeddings_flat` is already maintained incrementally, so this just
|
/// Prefer [`MemoryCache::flat_embeddings`] where a borrow will do.
|
||||||
/// clones it — kept as a method for callers that want an owned copy.
|
pub fn flat_embeddings_owned(&self) -> Vec<f32> {
|
||||||
pub fn flat_embeddings(&self) -> Vec<f32> {
|
self.embeddings.as_flat().to_vec()
|
||||||
self.embeddings_flat.clone()
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -224,7 +334,7 @@ mod tests {
|
|||||||
/// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
|
/// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
|
||||||
fn assert_flat_in_sync(cache: &MemoryCache) {
|
fn assert_flat_in_sync(cache: &MemoryCache) {
|
||||||
let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
|
let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
|
||||||
assert_eq!(cache.embeddings_flat, expected);
|
assert_eq!(cache.embeddings.as_flat(), expected);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -247,7 +357,10 @@ mod tests {
|
|||||||
String::new(),
|
String::new(),
|
||||||
);
|
);
|
||||||
assert_flat_in_sync(&cache);
|
assert_flat_in_sync(&cache);
|
||||||
assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
|
assert_eq!(
|
||||||
|
cache.embeddings.as_flat(),
|
||||||
|
vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -279,7 +392,7 @@ mod tests {
|
|||||||
);
|
);
|
||||||
assert_flat_in_sync(&cache);
|
assert_flat_in_sync(&cache);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
cache.embeddings_flat,
|
cache.embeddings.as_flat(),
|
||||||
vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
|
vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
|
||||||
"update must overwrite the correct flat slice, not just append"
|
"update must overwrite the correct flat slice, not just append"
|
||||||
);
|
);
|
||||||
@@ -315,14 +428,15 @@ mod tests {
|
|||||||
cache.mark_deleted(1);
|
cache.mark_deleted(1);
|
||||||
cache.compact();
|
cache.compact();
|
||||||
assert_flat_in_sync(&cache);
|
assert_flat_in_sync(&cache);
|
||||||
assert_eq!(cache.embeddings_flat, vec![1.0, 1.0, 3.0, 3.0]);
|
assert_eq!(cache.embeddings.as_flat(), vec![1.0, 1.0, 3.0, 3.0]);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rebuild_flat_matches_manual_flatten() {
|
fn rebuild_flat_matches_manual_flatten() {
|
||||||
let mut cache = MemoryCache::new(2);
|
let mut cache = MemoryCache::new(2);
|
||||||
cache.embeddings = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
cache
|
||||||
cache.rebuild_flat();
|
.embeddings
|
||||||
assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0]);
|
.reset_from(2, vec![vec![1.0, 2.0], vec![3.0, 4.0]]);
|
||||||
|
assert_eq!(cache.embeddings.as_flat(), vec![1.0, 2.0, 3.0, 4.0]);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -28,7 +28,7 @@ use crate::vector_search;
|
|||||||
pub fn hybrid_search(
|
pub fn hybrid_search(
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||||
chunks: &[String],
|
chunks: &[String],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
bm25_index: &BM25Index,
|
bm25_index: &BM25Index,
|
||||||
@@ -56,7 +56,7 @@ pub fn hybrid_search(
|
|||||||
pub fn hybrid_search_fused(
|
pub fn hybrid_search_fused(
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||||
_chunks: &[String],
|
_chunks: &[String],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
bm25_index: &BM25Index,
|
bm25_index: &BM25Index,
|
||||||
@@ -69,12 +69,12 @@ pub fn hybrid_search_fused(
|
|||||||
let vec_scores = {
|
let vec_scores = {
|
||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
{
|
{
|
||||||
if vectors.len() > 10_000 {
|
if vectors.count() > 10_000 {
|
||||||
vector_search::parallel_cosine_batch(
|
vector_search::parallel_cosine_batch(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
vectors,
|
vectors,
|
||||||
tombstones,
|
tombstones,
|
||||||
vectors.len(),
|
vectors.count(),
|
||||||
)
|
)
|
||||||
} else {
|
} else {
|
||||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
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(
|
pub fn rrf_hybrid_search(
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
|
||||||
_chunks: &[String],
|
_chunks: &[String],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
bm25_index: &BM25Index,
|
bm25_index: &BM25Index,
|
||||||
@@ -282,12 +282,12 @@ pub fn rrf_hybrid_search(
|
|||||||
let mut vec_scores = {
|
let mut vec_scores = {
|
||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
{
|
{
|
||||||
if vectors.len() > 10_000 {
|
if vectors.count() > 10_000 {
|
||||||
vector_search::parallel_cosine_batch(
|
vector_search::parallel_cosine_batch(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
vectors,
|
vectors,
|
||||||
tombstones,
|
tombstones,
|
||||||
vectors.len(),
|
vectors.count(),
|
||||||
)
|
)
|
||||||
} else {
|
} else {
|
||||||
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
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)
|
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.
|
// 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));
|
vec_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
|||||||
@@ -565,7 +565,9 @@ impl HDF5Memory {
|
|||||||
if u64::from_le_bytes(stamp.try_into().ok()?) != generation {
|
if u64::from_le_bytes(stamp.try_into().ok()?) != generation {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
let vectors = cache.embeddings.get(..n_checkpoint)?.to_vec();
|
let vectors: Vec<Vec<f32>> = (0..n_checkpoint)
|
||||||
|
.map(|i| cache.embeddings.get(i).map(<[f32]>::to_vec))
|
||||||
|
.collect::<Option<_>>()?;
|
||||||
let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?;
|
let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?;
|
||||||
if index.dimension() != cache.embedding_dim {
|
if index.dimension() != cache.embedding_dim {
|
||||||
return None;
|
return None;
|
||||||
@@ -573,7 +575,7 @@ impl HDF5Memory {
|
|||||||
// Records appended since (replayed from the WAL) join incrementally.
|
// Records appended since (replayed from the WAL) join incrementally.
|
||||||
for id in n_checkpoint..cache.embeddings.len() {
|
for id in n_checkpoint..cache.embeddings.len() {
|
||||||
if cache.embeddings[id].len() != index.dimension()
|
if cache.embeddings[id].len() != index.dimension()
|
||||||
|| index.insert(cache.embeddings[id].clone()) != id
|
|| index.insert(cache.embeddings[id].to_vec()) != id
|
||||||
{
|
{
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
@@ -816,8 +818,11 @@ impl HDF5Memory {
|
|||||||
if self.cache.embeddings.iter().any(|e| e.len() != dim) {
|
if self.cache.embeddings.iter().any(|e| e.len() != dim) {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
// The index owns its vectors, so it needs rows rather than the cache's
|
||||||
|
// flat buffer. This copy is the index's own; the cache keeps one.
|
||||||
|
let rows: Vec<Vec<f32>> = self.cache.embeddings.iter().map(<[f32]>::to_vec).collect();
|
||||||
let mut index = HnswIndex::build_with_metric(
|
let mut index = HnswIndex::build_with_metric(
|
||||||
&self.cache.embeddings,
|
&rows,
|
||||||
HNSW_M,
|
HNSW_M,
|
||||||
HNSW_EF_CONSTRUCTION,
|
HNSW_EF_CONSTRUCTION,
|
||||||
DistanceMetric::Cosine,
|
DistanceMetric::Cosine,
|
||||||
@@ -846,7 +851,7 @@ impl HDF5Memory {
|
|||||||
let dim = index.dimension();
|
let dim = index.dimension();
|
||||||
let appended = (self.hnsw_synced_len..n).all(|id| {
|
let appended = (self.hnsw_synced_len..n).all(|id| {
|
||||||
self.cache.embeddings[id].len() == dim
|
self.cache.embeddings[id].len() == dim
|
||||||
&& index.insert(self.cache.embeddings[id].clone()) == id
|
&& index.insert(self.cache.embeddings[id].to_vec()) == id
|
||||||
});
|
});
|
||||||
if appended {
|
if appended {
|
||||||
for id in self.hnsw_synced_len..n {
|
for id in self.hnsw_synced_len..n {
|
||||||
@@ -877,7 +882,7 @@ impl HDF5Memory {
|
|||||||
let emb_len = self.cache.embeddings[idx].len();
|
let emb_len = self.cache.embeddings[idx].len();
|
||||||
match self.hnsw.as_mut() {
|
match self.hnsw.as_mut() {
|
||||||
Some(index) if emb_len == index.dimension() => {
|
Some(index) if emb_len == index.dimension() => {
|
||||||
let id = index.insert(self.cache.embeddings[idx].clone());
|
let id = index.insert(self.cache.embeddings[idx].to_vec());
|
||||||
if id == idx {
|
if id == idx {
|
||||||
self.hnsw_synced_len = self.cache.embeddings.len();
|
self.hnsw_synced_len = self.cache.embeddings.len();
|
||||||
} else {
|
} else {
|
||||||
|
|||||||
@@ -466,7 +466,7 @@ impl ClawhdfBackend {
|
|||||||
let record = MemoryRecord {
|
let record = MemoryRecord {
|
||||||
id: i as u64,
|
id: i as u64,
|
||||||
chunk: cache.chunks[i].clone(),
|
chunk: cache.chunks[i].clone(),
|
||||||
embedding: cache.embeddings[i].clone(),
|
embedding: cache.embeddings[i].to_vec(),
|
||||||
tier: MemoryTier::Working,
|
tier: MemoryTier::Working,
|
||||||
importance: cache.activation_weights[i],
|
importance: cache.activation_weights[i],
|
||||||
access_count: 0,
|
access_count: 0,
|
||||||
@@ -717,11 +717,13 @@ impl MemoryBackend for ClawhdfBackend {
|
|||||||
|
|
||||||
let total_records = cache.count_active();
|
let total_records = cache.count_active();
|
||||||
|
|
||||||
|
// A record saved without an embedding occupies a zero row, so "has an
|
||||||
|
// embedding" is "has a non-zero norm" rather than "row is non-empty".
|
||||||
let total_embeddings = cache
|
let total_embeddings = cache
|
||||||
.embeddings
|
.norms
|
||||||
.iter()
|
.iter()
|
||||||
.enumerate()
|
.enumerate()
|
||||||
.filter(|(i, emb)| cache.tombstones[*i] == 0 && !emb.is_empty())
|
.filter(|(i, norm)| cache.tombstones[*i] == 0 && **norm > 0.0)
|
||||||
.count();
|
.count();
|
||||||
|
|
||||||
let file_size_bytes = std::fs::metadata(&self.hdf5_path)
|
let file_size_bytes = std::fs::metadata(&self.hdf5_path)
|
||||||
|
|||||||
@@ -153,7 +153,7 @@ fn build_memory_group(
|
|||||||
{
|
{
|
||||||
let ds = group
|
let ds = group
|
||||||
.create_dataset("embeddings")
|
.create_dataset("embeddings")
|
||||||
.with_f32_data(&flat)
|
.with_f32_data(flat)
|
||||||
.with_shape(&[n, d]);
|
.with_shape(&[n, d]);
|
||||||
|
|
||||||
// Chunk size tuning: target ~256KB per chunk for optimal I/O
|
// Chunk size tuning: target ~256KB per chunk for optimal I/O
|
||||||
@@ -563,12 +563,7 @@ fn load_memory_group(
|
|||||||
.collect(),
|
.collect(),
|
||||||
};
|
};
|
||||||
|
|
||||||
// Unflatten embeddings
|
// No unflattening: the cache stores the buffer as it is on disk.
|
||||||
let embeddings: Vec<Vec<f32>> = flat_embeddings
|
|
||||||
.chunks(embedding_dim)
|
|
||||||
.map(|c| c.to_vec())
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
// Read activation_weights if present, default to vec![1.0; N] for backward compat
|
// Read activation_weights if present, default to vec![1.0; N] for backward compat
|
||||||
let activation_weights = match read_f32_dataset(&group, "activation_weights") {
|
let activation_weights = match read_f32_dataset(&group, "activation_weights") {
|
||||||
Ok(w) if w.len() == n => w,
|
Ok(w) if w.len() == n => w,
|
||||||
@@ -576,7 +571,7 @@ fn load_memory_group(
|
|||||||
};
|
};
|
||||||
|
|
||||||
cache.chunks = chunks;
|
cache.chunks = chunks;
|
||||||
cache.embeddings = embeddings;
|
cache.embeddings.set_flat(embedding_dim, flat_embeddings);
|
||||||
cache.source_channels = source_channels;
|
cache.source_channels = source_channels;
|
||||||
cache.timestamps = timestamps;
|
cache.timestamps = timestamps;
|
||||||
cache.session_ids = session_ids;
|
cache.session_ids = session_ids;
|
||||||
@@ -584,7 +579,6 @@ fn load_memory_group(
|
|||||||
cache.tombstones = tombstones;
|
cache.tombstones = tombstones;
|
||||||
cache.norms = norms;
|
cache.norms = norms;
|
||||||
cache.activation_weights = activation_weights;
|
cache.activation_weights = activation_weights;
|
||||||
cache.rebuild_flat();
|
|
||||||
|
|
||||||
Ok(cache)
|
Ok(cache)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4,6 +4,44 @@
|
|||||||
//! `clawhdf5_accel`, with optional float16 support via the `half` crate.
|
//! `clawhdf5_accel`, with optional float16 support via the `half` crate.
|
||||||
//! Supports pre-computed norms for eliminating redundant norm computations.
|
//! Supports pre-computed norms for eliminating redundant norm computations.
|
||||||
|
|
||||||
|
/// A corpus of equal-length embeddings addressable by index.
|
||||||
|
///
|
||||||
|
/// Lets the batch kernels read either the cache's flat `[N x dim]` buffer or a
|
||||||
|
/// plain `Vec<Vec<f32>>` without either side owning a second copy.
|
||||||
|
pub trait VectorSet {
|
||||||
|
/// Number of embeddings.
|
||||||
|
fn count(&self) -> usize;
|
||||||
|
/// Embedding `i`; callers only index below [`VectorSet::count`].
|
||||||
|
fn row(&self, i: usize) -> &[f32];
|
||||||
|
}
|
||||||
|
|
||||||
|
impl VectorSet for [Vec<f32>] {
|
||||||
|
fn count(&self) -> usize {
|
||||||
|
self.len()
|
||||||
|
}
|
||||||
|
fn row(&self, i: usize) -> &[f32] {
|
||||||
|
&self[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl VectorSet for Vec<Vec<f32>> {
|
||||||
|
fn count(&self) -> usize {
|
||||||
|
self.len()
|
||||||
|
}
|
||||||
|
fn row(&self, i: usize) -> &[f32] {
|
||||||
|
&self[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl VectorSet for crate::cache::Embeddings {
|
||||||
|
fn count(&self) -> usize {
|
||||||
|
self.len()
|
||||||
|
}
|
||||||
|
fn row(&self, i: usize) -> &[f32] {
|
||||||
|
&self[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Compute cosine similarity between two f32 slices.
|
/// Compute cosine similarity between two f32 slices.
|
||||||
///
|
///
|
||||||
/// Returns 0.0 if either vector has zero magnitude.
|
/// Returns 0.0 if either vector has zero magnitude.
|
||||||
@@ -22,7 +60,7 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
|||||||
/// Returns `(index, score)` pairs sorted by score descending.
|
/// Returns `(index, score)` pairs sorted by score descending.
|
||||||
pub fn cosine_similarity_batch(
|
pub fn cosine_similarity_batch(
|
||||||
query: &[f32],
|
query: &[f32],
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl VectorSet + ?Sized),
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
let query_norm = clawhdf5_accel::vector_norm(query);
|
let query_norm = clawhdf5_accel::vector_norm(query);
|
||||||
@@ -30,7 +68,7 @@ pub fn cosine_similarity_batch(
|
|||||||
return Vec::new();
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
let n = vectors.len();
|
let n = vectors.count();
|
||||||
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
||||||
|
|
||||||
// Process 4 vectors at a time where possible
|
// Process 4 vectors at a time where possible
|
||||||
@@ -42,8 +80,9 @@ pub fn cosine_similarity_batch(
|
|||||||
if i < tombstones.len() && tombstones[i] != 0 {
|
if i < tombstones.len() && tombstones[i] != 0 {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
|
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
|
let score =
|
||||||
|
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||||
results.push((i, score));
|
results.push((i, score));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -53,8 +92,8 @@ pub fn cosine_similarity_batch(
|
|||||||
if i < tombstones.len() && tombstones[i] != 0 {
|
if i < tombstones.len() && tombstones[i] != 0 {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
|
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
|
let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||||
results.push((i, score));
|
results.push((i, score));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -68,7 +107,7 @@ pub fn cosine_similarity_batch(
|
|||||||
/// collections. Uses `score = dot(query, vec) / (query_norm * stored_norm)`.
|
/// collections. Uses `score = dot(query, vec) / (query_norm * stored_norm)`.
|
||||||
pub fn cosine_similarity_batch_prenorm(
|
pub fn cosine_similarity_batch_prenorm(
|
||||||
query: &[f32],
|
query: &[f32],
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl VectorSet + ?Sized),
|
||||||
norms: &[f32],
|
norms: &[f32],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
@@ -77,7 +116,7 @@ pub fn cosine_similarity_batch_prenorm(
|
|||||||
return Vec::new();
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
let n = vectors.len();
|
let n = vectors.count();
|
||||||
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
|
||||||
|
|
||||||
for i in 0..n {
|
for i in 0..n {
|
||||||
@@ -85,7 +124,7 @@ pub fn cosine_similarity_batch_prenorm(
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let vec_norm = norms[i];
|
let vec_norm = norms[i];
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
|
let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||||
results.push((i, score));
|
results.push((i, score));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -162,7 +201,7 @@ pub fn cosine_similarity_f16(
|
|||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
pub fn parallel_cosine_batch(
|
pub fn parallel_cosine_batch(
|
||||||
query: &[f32],
|
query: &[f32],
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl VectorSet + Sync + ?Sized),
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
@@ -174,24 +213,27 @@ pub fn parallel_cosine_batch(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let num_cores = rayon::current_num_threads().max(1);
|
let num_cores = rayon::current_num_threads().max(1);
|
||||||
let chunk_size = vectors.len().div_ceil(num_cores);
|
let chunk_size = vectors.count().div_ceil(num_cores);
|
||||||
if chunk_size == 0 {
|
if chunk_size == 0 {
|
||||||
return Vec::new();
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut all_results: Vec<(usize, f32)> = vectors
|
// Chunk over index ranges: the corpus may be one flat buffer rather than
|
||||||
.par_chunks(chunk_size)
|
// a slice of rows, so there is nothing to `par_chunks` over.
|
||||||
.enumerate()
|
let n = vectors.count();
|
||||||
.flat_map(|(chunk_idx, chunk)| {
|
let mut all_results: Vec<(usize, f32)> = (0..n.div_ceil(chunk_size))
|
||||||
|
.into_par_iter()
|
||||||
|
.flat_map(|chunk_idx| {
|
||||||
let base = chunk_idx * chunk_size;
|
let base = chunk_idx * chunk_size;
|
||||||
let mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
|
let end = (base + chunk_size).min(n);
|
||||||
for (j, vec) in chunk.iter().enumerate() {
|
let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
|
||||||
let i = base + j;
|
for i in base..end {
|
||||||
if i < tombstones.len() && tombstones[i] != 0 {
|
if i < tombstones.len() && tombstones[i] != 0 {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let vec_norm = clawhdf5_accel::vector_norm(vec);
|
let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, vec, vec_norm);
|
let score =
|
||||||
|
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
|
||||||
local.push((i, score));
|
local.push((i, score));
|
||||||
}
|
}
|
||||||
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
@@ -209,7 +251,7 @@ pub fn parallel_cosine_batch(
|
|||||||
#[cfg(feature = "parallel")]
|
#[cfg(feature = "parallel")]
|
||||||
pub fn parallel_cosine_batch_prenorm(
|
pub fn parallel_cosine_batch_prenorm(
|
||||||
query: &[f32],
|
query: &[f32],
|
||||||
vectors: &[Vec<f32>],
|
vectors: &(impl VectorSet + Sync + ?Sized),
|
||||||
norms: &[f32],
|
norms: &[f32],
|
||||||
tombstones: &[u8],
|
tombstones: &[u8],
|
||||||
k: usize,
|
k: usize,
|
||||||
@@ -222,23 +264,26 @@ pub fn parallel_cosine_batch_prenorm(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let num_cores = rayon::current_num_threads().max(1);
|
let num_cores = rayon::current_num_threads().max(1);
|
||||||
let chunk_size = vectors.len().div_ceil(num_cores);
|
let chunk_size = vectors.count().div_ceil(num_cores);
|
||||||
if chunk_size == 0 {
|
if chunk_size == 0 {
|
||||||
return Vec::new();
|
return Vec::new();
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut all_results: Vec<(usize, f32)> = vectors
|
// Chunk over index ranges: the corpus may be one flat buffer rather than
|
||||||
.par_chunks(chunk_size)
|
// a slice of rows, so there is nothing to `par_chunks` over.
|
||||||
.enumerate()
|
let n = vectors.count();
|
||||||
.flat_map(|(chunk_idx, chunk)| {
|
let mut all_results: Vec<(usize, f32)> = (0..n.div_ceil(chunk_size))
|
||||||
|
.into_par_iter()
|
||||||
|
.flat_map(|chunk_idx| {
|
||||||
let base = chunk_idx * chunk_size;
|
let base = chunk_idx * chunk_size;
|
||||||
let mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
|
let end = (base + chunk_size).min(n);
|
||||||
for (j, vec) in chunk.iter().enumerate() {
|
let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
|
||||||
let i = base + j;
|
for i in base..end {
|
||||||
if i < tombstones.len() && tombstones[i] != 0 {
|
if i < tombstones.len() && tombstones[i] != 0 {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
let score = crate::cosine_similarity_prenorm(query, query_norm, vec, norms[i]);
|
let score =
|
||||||
|
crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), norms[i]);
|
||||||
local.push((i, score));
|
local.push((i, score));
|
||||||
}
|
}
|
||||||
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||||
|
|||||||
@@ -198,6 +198,57 @@ fn summarize(mut samples: Vec<Duration>) -> Latency {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Counts live heap bytes, so a structure's cost can be measured by
|
||||||
|
/// difference.
|
||||||
|
///
|
||||||
|
/// RSS cannot do this from inside one process: freeing a large structure
|
||||||
|
/// returns its pages to the allocator's pool rather than to the OS, so
|
||||||
|
/// allocating the next one shows no change. Measured that way, a store that
|
||||||
|
/// holds the corpus twice and one that holds it once look identical.
|
||||||
|
struct CountingAllocator;
|
||||||
|
|
||||||
|
static LIVE_BYTES: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0);
|
||||||
|
|
||||||
|
// SAFETY: every method forwards to the system allocator with the same layout
|
||||||
|
// it was given, and only adds bookkeeping around it.
|
||||||
|
unsafe impl std::alloc::GlobalAlloc for CountingAllocator {
|
||||||
|
unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
|
||||||
|
let ptr = unsafe { std::alloc::System.alloc(layout) };
|
||||||
|
if !ptr.is_null() {
|
||||||
|
LIVE_BYTES.fetch_add(layout.size() as i64, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
ptr
|
||||||
|
}
|
||||||
|
|
||||||
|
unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
|
||||||
|
LIVE_BYTES.fetch_sub(layout.size() as i64, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
unsafe { std::alloc::System.dealloc(ptr, layout) }
|
||||||
|
}
|
||||||
|
|
||||||
|
unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
|
||||||
|
let new_ptr = unsafe { std::alloc::System.realloc(ptr, layout, new_size) };
|
||||||
|
if !new_ptr.is_null() {
|
||||||
|
LIVE_BYTES.fetch_add(
|
||||||
|
new_size as i64 - layout.size() as i64,
|
||||||
|
std::sync::atomic::Ordering::Relaxed,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
new_ptr
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[global_allocator]
|
||||||
|
static ALLOCATOR: CountingAllocator = CountingAllocator;
|
||||||
|
|
||||||
|
/// Live heap bytes right now.
|
||||||
|
fn heap_bytes() -> u64 {
|
||||||
|
LIVE_BYTES.load(std::sync::atomic::Ordering::Relaxed).max(0) as u64
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mib(bytes: u64) -> f64 {
|
||||||
|
bytes as f64 / (1 << 20) as f64
|
||||||
|
}
|
||||||
|
|
||||||
fn micros(d: Duration) -> f64 {
|
fn micros(d: Duration) -> f64 {
|
||||||
d.as_secs_f64() * 1e6
|
d.as_secs_f64() * 1e6
|
||||||
}
|
}
|
||||||
@@ -450,6 +501,62 @@ fn fusion_study(n: usize) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// What an in-memory store costs, stage by stage. The vectors are the floor:
|
||||||
|
/// everything above it is bookkeeping that could in principle be shared.
|
||||||
|
fn bench_footprint(n: usize) {
|
||||||
|
let data = make_dataset(n, 0xF007 ^ n as u64);
|
||||||
|
let mut rng = Rng(11);
|
||||||
|
let dir = tempfile::TempDir::new().unwrap();
|
||||||
|
let path = dir.path().join("footprint.h5");
|
||||||
|
|
||||||
|
let base = heap_bytes();
|
||||||
|
let entries: Vec<MemoryEntry> = data
|
||||||
|
.vectors
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(i, v)| MemoryEntry {
|
||||||
|
chunk: text_for(data.cluster_of[i], i, &mut rng),
|
||||||
|
embedding: v.clone(),
|
||||||
|
source_channel: "bench".into(),
|
||||||
|
timestamp: i as f64,
|
||||||
|
session_id: format!("s{}", i % 50),
|
||||||
|
tags: format!("t{i}"),
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let after_entries = heap_bytes();
|
||||||
|
|
||||||
|
let mut mem = HDF5Memory::create(MemoryConfig::new(path, "bench", DIM)).unwrap();
|
||||||
|
mem.save_batch(entries).unwrap();
|
||||||
|
let after_store = heap_bytes();
|
||||||
|
|
||||||
|
// First query builds the vector and keyword indexes.
|
||||||
|
std::hint::black_box(mem.hybrid_search(&data.queries[0], "record", 0.7, 0.3, K));
|
||||||
|
let after_indexes = heap_bytes();
|
||||||
|
|
||||||
|
// Reopening is the figure that matters for a long-lived process, and the
|
||||||
|
// only one RSS reports honestly: memory freed when the ingest buffers went
|
||||||
|
// away stays in the allocator's pool, so the stage deltas above understate
|
||||||
|
// what was given back.
|
||||||
|
let path = mem.config().path.clone();
|
||||||
|
drop(mem);
|
||||||
|
let before_open = heap_bytes();
|
||||||
|
let reopened = HDF5Memory::open(&path).unwrap();
|
||||||
|
let after_open = heap_bytes();
|
||||||
|
let loaded = after_open.saturating_sub(before_open);
|
||||||
|
drop(reopened);
|
||||||
|
|
||||||
|
let raw = (n * DIM * 4) as u64;
|
||||||
|
println!(
|
||||||
|
"| {n} | {:.0} | {:.0} | {:.0} | {:.0} | {:.0} | {:.2}x |",
|
||||||
|
mib(raw),
|
||||||
|
mib(after_entries.saturating_sub(base)),
|
||||||
|
mib(after_store.saturating_sub(after_entries)),
|
||||||
|
mib(after_indexes.saturating_sub(after_store)),
|
||||||
|
mib(loaded),
|
||||||
|
loaded as f64 / raw as f64,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||||
let full = args.iter().any(|a| a == "--full");
|
let full = args.iter().any(|a| a == "--full");
|
||||||
@@ -485,6 +592,18 @@ fn main() {
|
|||||||
|
|
||||||
let mut json = Vec::new();
|
let mut json = Vec::new();
|
||||||
println!("## Search harness");
|
println!("## Search harness");
|
||||||
|
|
||||||
|
if args.iter().any(|a| a == "--footprint") {
|
||||||
|
println!("\n### Resident memory, {DIM}-dim f32\n");
|
||||||
|
println!(
|
||||||
|
"| N | vectors (raw) | entries MiB | store MiB | indexes MiB | reopened MiB | reopened / raw |"
|
||||||
|
);
|
||||||
|
println!("|---:|---:|---:|---:|---:|---:|---:|");
|
||||||
|
for &n in sizes {
|
||||||
|
bench_footprint(n);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
// `--e2e-only` skips the index benchmarks, so the end-to-end section runs
|
// `--e2e-only` skips the index benchmarks, so the end-to-end section runs
|
||||||
// in a process that has not already spun up a thread pool.
|
// in a process that has not already spun up a thread pool.
|
||||||
if !args.iter().any(|a| a == "--e2e-only") {
|
if !args.iter().any(|a| a == "--e2e-only") {
|
||||||
|
|||||||
Reference in New Issue
Block a user