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

crates/clawhdf5-agent/src/cache.rs

@@ -2,16 +2,143 @@
 
 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.
 #[derive(Debug, Clone)]
 pub struct MemoryCache {
     pub chunks: Vec<String>,
-    pub embeddings: Vec<Vec<f32>>,
-    /// `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 embeddings: Embeddings,
     pub source_channels: Vec<String>,
     pub timestamps: Vec<f64>,
     pub session_ids: Vec<String>,
@@ -28,8 +155,7 @@ impl MemoryCache {
     pub fn new(embedding_dim: usize) -> Self {
         Self {
             chunks: Vec::new(),
-            embeddings: Vec::new(),
-            embeddings_flat: Vec::new(),
+            embeddings: Embeddings::new(embedding_dim),
             source_channels: Vec::new(),
             timestamps: Vec::new(),
             session_ids: Vec::new(),
@@ -41,15 +167,14 @@ impl MemoryCache {
         }
     }
 
-    /// Rebuild `embeddings_flat` from `embeddings` from scratch. Callers that
-    /// populate `embeddings` directly (bulk loads) must call this afterward.
-    pub fn rebuild_flat(&mut self) {
-        self.embeddings_flat.clear();
-        self.embeddings_flat
-            .reserve(self.embeddings.len() * self.embedding_dim);
-        for emb in &self.embeddings {
-            self.embeddings_flat.extend_from_slice(emb);
-        }
+    /// Kept for callers that used to have to re-flatten after a bulk load.
+    /// The buffer is always flat now, so there is nothing to rebuild.
+    #[deprecated(note = "embeddings are stored flat; this is a no-op")]
+    pub fn rebuild_flat(&mut self) {}
+
+    /// The embeddings as one contiguous `[N x dim]` buffer.
+    pub fn flat_embeddings(&self) -> &[f32] {
+        self.embeddings.as_flat()
     }
 
     /// Total number of entries (including tombstoned).
@@ -79,8 +204,7 @@ impl MemoryCache {
         let idx = self.chunks.len();
         let norm = vector_search::compute_norm(&embedding);
         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.timestamps.push(timestamp);
         self.session_ids.push(session_id);
@@ -118,20 +242,7 @@ impl MemoryCache {
         if idx < self.chunks.len() {
             let norm = vector_search::compute_norm(&embedding);
             self.chunks[idx] = chunk;
-            let dim = self.embedding_dim;
-            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.embeddings.set(idx, &embedding);
             self.source_channels[idx] = source_channel;
             self.timestamps[idx] = timestamp;
             self.session_ids[idx] = session_id;
@@ -183,7 +294,7 @@ impl MemoryCache {
                 new_idx += 1;
                 let norm = vector_search::compute_norm(&self.embeddings[i]);
                 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_timestamps.push(self.timestamps[i]);
                 new_session_ids.push(self.session_ids[i].clone());
@@ -196,7 +307,8 @@ impl MemoryCache {
 
         let removed = old_len - new_chunks.len();
         self.chunks = new_chunks;
-        self.embeddings = new_embeddings;
+        self.embeddings
+            .reset_from(self.embedding_dim, new_embeddings);
         self.source_channels = new_source_channels;
         self.timestamps = new_timestamps;
         self.session_ids = new_session_ids;
@@ -204,16 +316,14 @@ impl MemoryCache {
         self.tombstones = new_tombstones;
         self.norms = new_norms;
         self.activation_weights = new_activation_weights;
-        self.rebuild_flat();
 
         (removed, index_map)
     }
 
-    /// Flatten all embeddings into a single Vec<f32> for HDF5 storage.
-    /// `embeddings_flat` is already maintained incrementally, so this just
-    /// clones it — kept as a method for callers that want an owned copy.
-    pub fn flat_embeddings(&self) -> Vec<f32> {
-        self.embeddings_flat.clone()
+    /// All embeddings as one owned `[N x dim]` buffer, for HDF5 storage.
+    /// Prefer [`MemoryCache::flat_embeddings`] where a borrow will do.
+    pub fn flat_embeddings_owned(&self) -> Vec<f32> {
+        self.embeddings.as_flat().to_vec()
     }
 }
 
@@ -224,7 +334,7 @@ mod tests {
     /// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
     fn assert_flat_in_sync(cache: &MemoryCache) {
         let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
-        assert_eq!(cache.embeddings_flat, expected);
+        assert_eq!(cache.embeddings.as_flat(), expected);
     }
 
     #[test]
@@ -247,7 +357,10 @@ mod tests {
             String::new(),
         );
         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]
@@ -279,7 +392,7 @@ mod tests {
         );
         assert_flat_in_sync(&cache);
         assert_eq!(
-            cache.embeddings_flat,
+            cache.embeddings.as_flat(),
             vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
             "update must overwrite the correct flat slice, not just append"
         );
@@ -315,14 +428,15 @@ mod tests {
         cache.mark_deleted(1);
         cache.compact();
         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]
     fn rebuild_flat_matches_manual_flatten() {
         let mut cache = MemoryCache::new(2);
-        cache.embeddings = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
-        cache.rebuild_flat();
-        assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0]);
+        cache
+            .embeddings
+            .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]);
     }
 }

crates/clawhdf5-agent/src/hybrid.rs

@@ -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));

crates/clawhdf5-agent/src/lib.rs

@@ -565,7 +565,9 @@ impl HDF5Memory {
         if u64::from_le_bytes(stamp.try_into().ok()?) != generation {
             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()?;
         if index.dimension() != cache.embedding_dim {
             return None;
@@ -573,7 +575,7 @@ impl HDF5Memory {
         // Records appended since (replayed from the WAL) join incrementally.
         for id in n_checkpoint..cache.embeddings.len() {
             if cache.embeddings[id].len() != index.dimension()
-                || index.insert(cache.embeddings[id].clone()) != id
+                || index.insert(cache.embeddings[id].to_vec()) != id
             {
                 return None;
             }
@@ -816,8 +818,11 @@ impl HDF5Memory {
         if self.cache.embeddings.iter().any(|e| e.len() != dim) {
             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(
-            &self.cache.embeddings,
+            &rows,
             HNSW_M,
             HNSW_EF_CONSTRUCTION,
             DistanceMetric::Cosine,
@@ -846,7 +851,7 @@ impl HDF5Memory {
             let dim = index.dimension();
             let appended = (self.hnsw_synced_len..n).all(|id| {
                 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 {
                 for id in self.hnsw_synced_len..n {
@@ -877,7 +882,7 @@ impl HDF5Memory {
         let emb_len = self.cache.embeddings[idx].len();
         match self.hnsw.as_mut() {
             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 {
                     self.hnsw_synced_len = self.cache.embeddings.len();
                 } else {

crates/clawhdf5-agent/src/openclaw.rs

@@ -466,7 +466,7 @@ impl ClawhdfBackend {
             let record = MemoryRecord {
                 id: i as u64,
                 chunk: cache.chunks[i].clone(),
-                embedding: cache.embeddings[i].clone(),
+                embedding: cache.embeddings[i].to_vec(),
                 tier: MemoryTier::Working,
                 importance: cache.activation_weights[i],
                 access_count: 0,
@@ -717,11 +717,13 @@ impl MemoryBackend for ClawhdfBackend {
 
         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
-            .embeddings
+            .norms
             .iter()
             .enumerate()
-            .filter(|(i, emb)| cache.tombstones[*i] == 0 && !emb.is_empty())
+            .filter(|(i, norm)| cache.tombstones[*i] == 0 && **norm > 0.0)
             .count();
 
         let file_size_bytes = std::fs::metadata(&self.hdf5_path)

crates/clawhdf5-agent/src/schema.rs

@@ -153,7 +153,7 @@ fn build_memory_group(
     {
         let ds = group
             .create_dataset("embeddings")
-            .with_f32_data(&flat)
+            .with_f32_data(flat)
             .with_shape(&[n, d]);
 
         // Chunk size tuning: target ~256KB per chunk for optimal I/O
@@ -563,12 +563,7 @@ fn load_memory_group(
             .collect(),
     };
 
-    // Unflatten embeddings
-    let embeddings: Vec<Vec<f32>> = flat_embeddings
-        .chunks(embedding_dim)
-        .map(|c| c.to_vec())
-        .collect();
-
+    // No unflattening: the cache stores the buffer as it is on disk.
     // Read activation_weights if present, default to vec![1.0; N] for backward compat
     let activation_weights = match read_f32_dataset(&group, "activation_weights") {
         Ok(w) if w.len() == n => w,
@@ -576,7 +571,7 @@ fn load_memory_group(
     };
 
     cache.chunks = chunks;
-    cache.embeddings = embeddings;
+    cache.embeddings.set_flat(embedding_dim, flat_embeddings);
     cache.source_channels = source_channels;
     cache.timestamps = timestamps;
     cache.session_ids = session_ids;
@@ -584,7 +579,6 @@ fn load_memory_group(
     cache.tombstones = tombstones;
     cache.norms = norms;
     cache.activation_weights = activation_weights;
-    cache.rebuild_flat();
 
     Ok(cache)
 }

crates/clawhdf5-agent/src/vector_search.rs

@@ -4,6 +4,44 @@
 //! `clawhdf5_accel`, with optional float16 support via the `half` crate.
 //! 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.
 ///
 /// 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.
 pub fn cosine_similarity_batch(
     query: &[f32],
-    vectors: &[Vec<f32>],
+    vectors: &(impl VectorSet + ?Sized),
     tombstones: &[u8],
 ) -> Vec<(usize, f32)> {
     let query_norm = clawhdf5_accel::vector_norm(query);
@@ -30,7 +68,7 @@ pub fn cosine_similarity_batch(
         return Vec::new();
     }
 
-    let n = vectors.len();
+    let n = vectors.count();
     let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
 
     // Process 4 vectors at a time where possible
@@ -42,8 +80,9 @@ pub fn cosine_similarity_batch(
             if i < tombstones.len() && tombstones[i] != 0 {
                 continue;
             }
-            let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
-            let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
+            let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
+            let score =
+                crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
             results.push((i, score));
         }
     }
@@ -53,8 +92,8 @@ pub fn cosine_similarity_batch(
         if i < tombstones.len() && tombstones[i] != 0 {
             continue;
         }
-        let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]);
-        let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm);
+        let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
+        let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
         results.push((i, score));
     }
 
@@ -68,7 +107,7 @@ pub fn cosine_similarity_batch(
 /// collections. Uses `score = dot(query, vec) / (query_norm * stored_norm)`.
 pub fn cosine_similarity_batch_prenorm(
     query: &[f32],
-    vectors: &[Vec<f32>],
+    vectors: &(impl VectorSet + ?Sized),
     norms: &[f32],
     tombstones: &[u8],
 ) -> Vec<(usize, f32)> {
@@ -77,7 +116,7 @@ pub fn cosine_similarity_batch_prenorm(
         return Vec::new();
     }
 
-    let n = vectors.len();
+    let n = vectors.count();
     let mut results: Vec<(usize, f32)> = Vec::with_capacity(n);
 
     for i in 0..n {
@@ -85,7 +124,7 @@ pub fn cosine_similarity_batch_prenorm(
             continue;
         }
         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));
     }
 
@@ -162,7 +201,7 @@ pub fn cosine_similarity_f16(
 #[cfg(feature = "parallel")]
 pub fn parallel_cosine_batch(
     query: &[f32],
-    vectors: &[Vec<f32>],
+    vectors: &(impl VectorSet + Sync + ?Sized),
     tombstones: &[u8],
     k: usize,
 ) -> Vec<(usize, f32)> {
@@ -174,24 +213,27 @@ pub fn parallel_cosine_batch(
     }
 
     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 {
         return Vec::new();
     }
 
-    let mut all_results: Vec<(usize, f32)> = vectors
-        .par_chunks(chunk_size)
-        .enumerate()
-        .flat_map(|(chunk_idx, chunk)| {
+    // Chunk over index ranges: the corpus may be one flat buffer rather than
+    // a slice of rows, so there is nothing to `par_chunks` over.
+    let n = vectors.count();
+    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 mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
-            for (j, vec) in chunk.iter().enumerate() {
-                let i = base + j;
+            let end = (base + chunk_size).min(n);
+            let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
+            for i in base..end {
                 if i < tombstones.len() && tombstones[i] != 0 {
                     continue;
                 }
-                let vec_norm = clawhdf5_accel::vector_norm(vec);
-                let score = crate::cosine_similarity_prenorm(query, query_norm, vec, vec_norm);
+                let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i));
+                let score =
+                    crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm);
                 local.push((i, score));
             }
             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")]
 pub fn parallel_cosine_batch_prenorm(
     query: &[f32],
-    vectors: &[Vec<f32>],
+    vectors: &(impl VectorSet + Sync + ?Sized),
     norms: &[f32],
     tombstones: &[u8],
     k: usize,
@@ -222,23 +264,26 @@ pub fn parallel_cosine_batch_prenorm(
     }
 
     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 {
         return Vec::new();
     }
 
-    let mut all_results: Vec<(usize, f32)> = vectors
-        .par_chunks(chunk_size)
-        .enumerate()
-        .flat_map(|(chunk_idx, chunk)| {
+    // Chunk over index ranges: the corpus may be one flat buffer rather than
+    // a slice of rows, so there is nothing to `par_chunks` over.
+    let n = vectors.count();
+    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 mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len());
-            for (j, vec) in chunk.iter().enumerate() {
-                let i = base + j;
+            let end = (base + chunk_size).min(n);
+            let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base);
+            for i in base..end {
                 if i < tombstones.len() && tombstones[i] != 0 {
                     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.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));