Merge feat/ann-quantisation: optional int8 vector index; interop suites run again

Cuts a loaded store's memory from 2.72x to 1.74x the raw vectors at 100k
x 384 via MemoryConfig::quantized_index, with recall held at the f32
index's level by re-scoring candidates against the exact embeddings the
store already holds. Off by default: it trades ~13% of QPS for the
memory.

Also restores the Python interop suites, which had been skipping
silently on this machine because no interpreter has h5py and PEP 668
blocks installing it into the system one.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-19 20:44:59 -07:00
co-authored by Claude Opus 5
21 changed files with 801 additions and 94 deletions
+1
View File
@@ -4,3 +4,4 @@ benchmarks/longmemeval/*.json
# Local model weights (MiniLM etc.) — large, not committed
weights/
.venv
+46
View File
@@ -57,6 +57,52 @@ 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.
### Quantising the index copy (`quantized_index`)
`MemoryConfig::quantized_index` stores the index's copy as `i8` instead of
`f32`. Same harness, same binary, `--footprint --full` with and without
`--int8`:
| N | vectors (raw) | indexes, f32 | indexes, int8 | reopened, f32 | reopened, int8 |
|---:|---:|---:|---:|---:|---:|
| 1 000 | 1 MiB | 2 MiB | 1 MiB | 4 MiB (2.40x) | 2 MiB (1.64x) |
| 10 000 | 15 MiB | 32 MiB | 14 MiB | 44 MiB (3.03x) | 27 MiB (1.81x) |
| 100 000 | 146 MiB | 266 MiB | **123 MiB** | 399 MiB (2.72x) | **256 MiB (1.74x)** |
The scale is **per row**, not global. A unit-length row in `d` dimensions has
components around `1/sqrt(d)`, so a fixed `[-1, 1]` scale spends fewer than 12
of the 255 levels on a 128-dimensional vector: measured against an exact
ranking that gives 0.35 top-10 overlap — unusable. Scaling each row by its own
largest component brings the same measurement to 0.99.
Quantised distances still cost recall on their own, and **`ef` does not buy it
back**, because the loss is in the distances rather than in the graph
(`--ann-only --full`, N = 100 000):
| ef | recall@10, f32 | recall@10, int8 | recall@10, int8 + re-score |
|---:|---:|---:|---:|
| 32 | 0.9775 | 0.9415 | 0.9785 |
| 64 | 0.9945 | 0.9625 | 0.9940 |
| 128 | 0.9995 | 0.9670 | 0.9990 |
| 256 | 0.9995 | 0.9670 (ceiling) | 0.9990 |
Re-scoring closes the gap: the store already holds the exact embeddings, so
the query path re-scores the candidate pool against them before fusion. That
is done automatically whenever the index is quantised. What it costs is
throughput — about 13% of QPS and 16% of build time at 100 000 x 384. So the
setting trades ~13% of query speed for ~36% of the process's memory at equal
recall. It is **off by default**: the right side of that trade depends on
whether the deployment is short of memory or short of CPU.
A measurement trap worth recording: the synthetic `clustered` generator in the
`clawhdf5-ann` tests draws clusters far tighter than any real embedding, so
neighbours there sit closer together than the quantisation error and top-10
*identity* is noise. Scored on that fixture int8 looks catastrophic (0.57
overlap) — a fact about the fixture, not the storage. The tests use random
vectors, and recall is measured against brute-force ground truth rather than
against the f32 index, whose own approximation errors a re-scored search is
entitled to get right.
## Read harness
Produced by `cargo run --release -p clawhdf5-bench --bin read_harness`: a 4096 x
+29
View File
@@ -2,6 +2,35 @@
## Unreleased
### Testing
- The Python interop suites honour **`CLAWHDF5_PYTHON`**, and `ci-test.sh`
picks up a `.venv/bin/python` automatically. On a PEP 668 "externally
managed" system h5py cannot be installed into the system interpreter at all,
so every interop suite — the h5py writer round-trips, the facade, netCDF4
and the reference files — was skipping silently. A silent skip here is
exactly how the v5 compound-datatype bug reached a release.
`CLAWHDF5_REQUIRE_INTEROP=1` still turns a skip into a failure.
### Memory
- `clawhdf5-agent`: **`MemoryConfig::quantized_index`** stores the vector
index's own copy of the embeddings as `i8` rather than `f32`, which at 100k
384-dim entries takes the index from 266 to 123 MiB and the whole reopened
store from 399 to 256 MiB (2.72x -> **1.74x** the raw vectors). Quantised
distances are approximate and `ef` cannot compensate — recall@10 tops out at
0.967 against f32's 0.9995 — so the query path re-scores the candidate pool
against the exact embeddings the store already holds, which restores recall
(0.9940 vs 0.9945 at ef=64) for about 13% of QPS. **Off by default**: it
trades query speed for memory, and which side is worth more depends on the
deployment. The setting is persisted, so a reopened store does not silently
revert to four times the index memory.
- `clawhdf5-ann`: `Storage::Int8` and the `build_with` / `new_with` /
`from_graph_bytes_with` constructors that select it. The scale is per row,
not global — a fixed `[-1, 1]` scale spends fewer than 12 of the 255 levels
on a unit-length 128-dim vector and is unusable (0.35 top-10 overlap against
an exact ranking, versus 0.99 per row). `compact()` keeps the storage it was
given; serialized indexes still carry f32 vectors, so a quantised index is
rebuilt rather than loaded.
### 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
+7 -1
View File
@@ -39,7 +39,13 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F
(plain closest-M capped recall on clustered data: 0.31 recall@10 at 100K). Its
graph is saved to `<store>.h5.ann` at each checkpoint and reloaded by `open()`
(tied to the checkpoint by a generation id; stale/damaged sidecars are
ignored and the index rebuilt). `hybrid_search` keeps one incremental BM25
ignored and the index rebuilt). `MemoryConfig::quantized_index` (off by
default, persisted) stores the index's own copy of the embeddings as `i8`,
which roughly halves a loaded store's memory (2.72x -> 1.74x the raw vectors
at 100K); because quantised distances are approximate and `ef` cannot
compensate, the query path then re-scores the candidate pool against the
exact embeddings, which holds recall at the f32 index's level and costs
~13% of QPS. `hybrid_search` keeps one incremental BM25
index for the life of the store and never writes the store: Hebbian
activation boosts are persisted by the next checkpoint (or on drop), not per
query. Measure any search-path change with
+7
View File
@@ -432,6 +432,13 @@ ClawhDF5's agent memory design draws from 15+ recent papers:
| `agent` | no | Full agent memory layer |
| `float16` | **yes** | Half-precision embedding storage (2× compression) |
| `hnsw` | **yes** | HNSW approximate vector index for `hybrid_search` (via `clawhdf5-ann`); disable for an exact linear scan |
`MemoryConfig::quantized_index` (off by default) stores the HNSW index's own
copy of the embeddings as `i8`, roughly halving a loaded store's memory
(2.72x -> 1.74x the raw vectors at 100k x 384). Quantised distances are
approximate, so the query path re-scores the candidate pool against the exact
embeddings the store already holds — recall matches the `f32` index, at about
13% fewer queries per second. See `BENCHMARKS.md`, "Quantising the index copy".
| `parallel` | no | Rayon parallel search |
| `fast-math` | no | BLAS matrix-vector multiply |
| `accelerate` | no | Apple Accelerate / AMX (macOS) |
+1
View File
@@ -118,6 +118,7 @@ mod tests {
created_at: "2025-01-01T00:00:00Z".to_string(),
wal_enabled: false,
wal_max_entries: 500,
quantized_index: false,
}
}
+38 -4
View File
@@ -62,7 +62,7 @@ use std::path::{Path, PathBuf};
use cache::MemoryCache;
#[cfg(feature = "hnsw")]
use clawhdf5_ann::{DistanceMetric, HnswIndex};
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
use ephemeral::{EphemeralConfig, EphemeralStore};
/// HNSW construction parameters used for the agent's vector index. Cosine is the
@@ -139,6 +139,17 @@ pub struct MemoryConfig {
pub created_at: String,
pub wal_enabled: bool,
pub wal_max_entries: usize,
/// Store the vector index's own copy of the embeddings as int8 rather than
/// f32, a quarter of the memory.
///
/// The index's copy is the single largest part of a loaded store's
/// footprint. Quantised distances are approximate, so the candidate pool
/// is re-scored against the cache's exact embeddings before fusion, which
/// restores recall; what it costs is throughput — roughly 13% of queries
/// per second and 16% of build time at 100K x 384. See `BENCHMARKS.md`.
///
/// Has no effect without the `hnsw` feature.
pub quantized_index: bool,
}
impl MemoryConfig {
@@ -160,6 +171,7 @@ impl MemoryConfig {
created_at,
wal_enabled: true,
wal_max_entries: 500,
quantized_index: false,
}
}
}
@@ -444,7 +456,17 @@ impl HDF5Memory {
#[cfg(feature = "hnsw")]
let loaded_index = if replay_only_appended {
Self::load_vector_index(path, checkpoint.ann_generation, &cache, n_checkpoint)
Self::load_vector_index(
path,
checkpoint.ann_generation,
&cache,
n_checkpoint,
if config.quantized_index {
Storage::Int8
} else {
Storage::Float32
},
)
} else {
None
};
@@ -558,6 +580,7 @@ impl HDF5Memory {
generation: Option<u64>,
cache: &MemoryCache,
n_checkpoint: usize,
storage: Storage,
) -> Option<HnswIndex> {
let generation = generation?;
let bytes = std::fs::read(Self::vector_index_path(store)).ok()?;
@@ -568,7 +591,7 @@ impl HDF5Memory {
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_with(graph, vectors, storage).ok()?;
if index.dimension() != cache.embedding_dim {
return None;
}
@@ -803,6 +826,16 @@ impl HDF5Memory {
// the index length drifts from the cache length (covering any mutation path
// that doesn't call a hook, e.g. consolidation pushes).
/// How the index should store its copy of the vectors, per the config.
#[cfg(feature = "hnsw")]
fn index_storage(&self) -> Storage {
if self.config.quantized_index {
Storage::Int8
} else {
Storage::Float32
}
}
/// Build an HNSW index over the entire cache, re-applying tombstones as
/// soft-deletions so node ids stay aligned with cache indices.
///
@@ -821,11 +854,12 @@ impl HDF5Memory {
// 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(
&rows,
HNSW_M,
HNSW_EF_CONSTRUCTION,
DistanceMetric::Cosine,
self.index_storage(),
);
for (i, &t) in self.cache.tombstones.iter().enumerate() {
if t != 0 {
+5
View File
@@ -104,6 +104,10 @@ pub fn build_hdf5_file_with_meta(
"wal_max_entries",
AttrValue::I64(config.wal_max_entries as i64),
);
meta.set_attr(
"quantized_index",
AttrValue::I64(config.quantized_index.into()),
);
meta.set_attr(
"edgehdf5_version",
AttrValue::String(ZEROCLAW_VERSION.into()),
@@ -484,6 +488,7 @@ pub fn validate_and_load(
wal_max_entries: optional_i64_attr(&attrs, "wal_max_entries")
.and_then(|v| usize::try_from(v).ok())
.unwrap_or(500),
quantized_index: optional_bool_attr(&attrs, "quantized_index", false),
};
// Load /memory group
+19 -3
View File
@@ -29,10 +29,26 @@ impl HDF5Memory {
// Over-fetch so the merge sees a useful vector pool; cosine
// distance from the index converts back to similarity (1 - d).
let pool = (k * 8).max(64);
let vec_scores: Vec<(usize, f32)> = index
.search(query_embedding, pool, pool)
let candidates = index.search(query_embedding, pool, pool);
// A quantised index returns approximate distances, and no
// amount of `ef` fixes that — the loss is in the distances,
// not the graph. Re-score the pool against the cache's exact
// embeddings, which cost nothing extra to keep: recall then
// matches an f32 index. See `BENCHMARKS.md`.
let exact = index.storage() == clawhdf5_ann::Storage::Int8;
let vec_scores: Vec<(usize, f32)> = candidates
.into_iter()
.map(|(id, dist)| (id, 1.0 - dist))
.map(|(id, dist)| {
let score = if exact {
crate::vector_search::cosine_similarity(
query_embedding,
&self.cache.embeddings[id],
)
} else {
1.0 - dist
};
(id, score)
})
.collect();
// Fusion normalises over every keyword match, so it needs all
// the scores — but not ranked.
@@ -165,3 +165,72 @@ fn save_batch_then_search_is_consistent() {
);
}
}
#[test]
fn quantized_index_matches_the_f32_index_after_re_scoring() {
// A quantised index holds approximate vectors, but the store still has the
// exact ones, so the query path re-scores the candidate pool before
// fusion. The results a caller sees should therefore be the same.
let dim = 64;
let n = 400;
let mut seed = 0x5EED_1234_5678_9ABC;
let vectors: Vec<Vec<f32>> = (0..n).map(|_| make_vector(&mut seed, dim)).collect();
let queries: Vec<Vec<f32>> = (0..20).map(|_| make_vector(&mut seed, dim)).collect();
let build = |dir: &TempDir, quantized: bool| {
let mut config = MemoryConfig::new(dir.path().join("mem.h5"), "agent", dim);
config.quantized_index = quantized;
let mut mem = HDF5Memory::create(config).unwrap();
for (i, v) in vectors.iter().enumerate() {
mem.save(entry(&format!("chunk {i}"), v.clone(), &format!("k{i}")))
.unwrap();
}
mem
};
let exact_dir = TempDir::new().unwrap();
let quant_dir = TempDir::new().unwrap();
let mut exact = build(&exact_dir, false);
let mut quantized = build(&quant_dir, true);
let k = 10;
let mut agree = 0;
for q in &queries {
let want: Vec<usize> = exact
.hybrid_search(q, "", 1.0, 0.0, k)
.iter()
.map(|r| r.index)
.collect();
agree += quantized
.hybrid_search(q, "", 1.0, 0.0, k)
.iter()
.filter(|r| want.contains(&r.index))
.count();
}
let overlap = agree as f64 / (k * queries.len()) as f64;
assert!(
overlap >= 0.95,
"quantised store should match the f32 one: {overlap}"
);
}
#[test]
fn quantized_index_setting_survives_a_reopen() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("mem.h5");
let mut config = MemoryConfig::new(path.clone(), "agent", 8);
config.quantized_index = true;
let mut mem = HDF5Memory::create(config).unwrap();
let mut seed = 7;
for i in 0..30 {
mem.save(entry(&format!("c{i}"), make_vector(&mut seed, 8), "t"))
.unwrap();
}
mem.flush_wal().unwrap();
drop(mem);
// Reopening must not silently quadruple the index's memory, so the flag
// is part of the stored config rather than a per-session choice.
let reopened = HDF5Memory::open(&path).unwrap();
assert!(reopened.config().quantized_index);
}
+431 -62
View File
@@ -154,6 +154,237 @@ impl Ord for FarCandidate {
}
}
/// How the index keeps its copy of the vectors.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Storage {
/// Exactly as given: `dim * 4` bytes per vector.
#[default]
Float32,
/// Each component scaled to an `i8`: `dim` bytes per vector, a quarter of
/// the space, at some cost in precision.
///
/// Only meaningful for [`DistanceMetric::Cosine`]: rows are stored
/// unit-length, so a quantised dot product reconstructs the similarity
/// directly. Requesting it for `L2` keeps `Float32`, because an L2
/// distance cannot be recovered from a dot product alone.
Int8,
}
/// The index's copy of the vectors, flat and row-major.
#[derive(Debug, Clone)]
enum Vectors {
F32 {
dim: usize,
flat: Vec<f32>,
},
/// `flat[i * dim + j]` is component `j` of vector `i` divided by
/// `scales[i]`; multiplying back recovers it.
///
/// The scale is per row rather than global. A unit-length row in `d`
/// dimensions has components around `1/sqrt(d)`, so a fixed `[-1, 1]`
/// scale spends fewer than 12 of the 255 levels on a 128-dimensional
/// vector and the reconstruction error swamps the gaps between near
/// neighbours — measured at 0.35 top-10 overlap with the exact ranking.
/// Scaling each row by its own largest component uses the full range.
Int8 {
dim: usize,
flat: Vec<i8>,
scales: Vec<f32>,
},
}
/// Levels either side of zero. 127, not 128, so the range is symmetric.
const INT8_LEVELS: f32 = 127.0;
/// Quantise one row, returning the codes and the scale that inverts them.
fn quantise_row(v: &[f32], out: &mut Vec<i8>) -> f32 {
let max_abs = v.iter().fold(0.0f32, |m, x| m.max(x.abs()));
if max_abs <= f32::MIN_POSITIVE {
out.extend(core::iter::repeat_n(0i8, v.len()));
return 0.0;
}
let inv = INT8_LEVELS / max_abs;
out.extend(
v.iter()
.map(|x| (x * inv).round().clamp(-INT8_LEVELS, INT8_LEVELS) as i8),
);
max_abs / INT8_LEVELS
}
impl Vectors {
fn new(dim: usize, storage: Storage, metric: DistanceMetric) -> Self {
match storage {
Storage::Int8 if metric == DistanceMetric::Cosine => Vectors::Int8 {
dim,
flat: Vec::new(),
scales: Vec::new(),
},
_ => Vectors::F32 {
dim,
flat: Vec::new(),
},
}
}
fn dim(&self) -> usize {
match self {
Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim,
}
}
fn storage(&self) -> Storage {
match self {
Vectors::F32 { .. } => Storage::Float32,
Vectors::Int8 { .. } => Storage::Int8,
}
}
fn len(&self) -> usize {
let dim = self.dim();
if dim == 0 {
return 0;
}
match self {
Vectors::F32 { flat, .. } => flat.len() / dim,
Vectors::Int8 { flat, .. } => flat.len() / dim,
}
}
/// Set the row width, for a store seeded empty by `new`.
fn set_dim(&mut self, new_dim: usize) {
match self {
Vectors::F32 { dim, .. } | Vectors::Int8 { dim, .. } => *dim = new_dim,
}
}
fn push(&mut self, vector: &[f32]) {
match self {
Vectors::F32 { flat, .. } => flat.extend_from_slice(vector),
Vectors::Int8 { flat, scales, .. } => scales.push(quantise_row(vector, flat)),
}
}
/// Row `i` as `f32`, for callers that need the values back (serialization,
/// and the f32 fast paths). Quantised rows are reconstructed, so this is
/// lossy in exactly the way the storage is.
fn row(&self, i: usize) -> Vec<f32> {
let dim = self.dim();
let start = i * dim;
match self {
Vectors::F32 { flat, .. } => flat[start..start + dim].to_vec(),
Vectors::Int8 { flat, scales, .. } => flat[start..start + dim]
.iter()
.map(|&q| f32::from(q) * scales[i])
.collect(),
}
}
/// Distance between two stored vectors.
fn dist(&self, a: usize, b: usize, metric: DistanceMetric) -> f32 {
let dim = self.dim();
match self {
Vectors::F32 { flat, .. } => {
let (x, y) = (a * dim, b * dim);
compute_distance(&flat[x..x + dim], &flat[y..y + dim], metric)
}
Vectors::Int8 { flat, scales, .. } => {
let (x, y) = (a * dim, b * dim);
let dot = dot_i8(&flat[x..x + dim], &flat[y..y + dim]);
1.0 - dot as f32 * scales[a] * scales[b]
}
}
}
/// Distance from a prepared query to stored vector `i`.
fn dist_query(&self, query: &Query, i: usize, metric: DistanceMetric) -> f32 {
let dim = self.dim();
let start = i * dim;
match (self, query) {
(Vectors::F32 { flat, .. }, Query::F32(q)) => {
compute_distance(q, &flat[start..start + dim], metric)
}
(Vectors::Int8 { flat, scales, .. }, Query::Int8(q, q_scale)) => {
let dot = dot_i8(q, &flat[start..start + dim]);
1.0 - dot as f32 * q_scale * scales[i]
}
// Mixed forms cannot occur: `Query` is built from the same storage.
_ => f32::MAX,
}
}
/// Build a store from prepared rows.
fn from_rows(rows: &[Vec<f32>], storage: Storage, metric: DistanceMetric) -> Self {
let dim = rows.first().map_or(0, Vec::len);
let mut out = Vectors::new(dim, storage, metric);
for row in rows {
out.push(row);
}
out
}
/// Prepare `query` for comparison against this store.
fn query(&self, query: Vec<f32>) -> Query {
match self {
Vectors::F32 { .. } => Query::F32(query),
Vectors::Int8 { .. } => {
let mut codes = Vec::with_capacity(query.len());
let scale = quantise_row(&query, &mut codes);
Query::Int8(codes, scale)
}
}
}
}
/// What a layer search is measuring distance *to*: an incoming query, or a
/// node already in the index (which is what insertion compares against).
enum Target<'a> {
Query(&'a Query),
Node(usize),
}
impl Vectors {
fn dist_to(&self, target: &Target<'_>, i: usize, metric: DistanceMetric) -> f32 {
match target {
Target::Query(q) => self.dist_query(q, i, metric),
Target::Node(n) => self.dist(*n, i, metric),
}
}
}
/// A search query in whichever form the store compares against.
enum Query {
F32(Vec<f32>),
/// Codes and the scale that inverts them, as in [`Vectors::Int8`].
Int8(Vec<i8>, f32),
}
/// Sum of products, widened so it cannot overflow: `dim` terms of at most
/// `127 * 127`, so `i32` suffices for any realistic dimension.
fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
// Four independent accumulators over 32-lane blocks: the widening product
// has to sit in a fixed-length chunk for the vectoriser to see it, and the
// separate accumulators keep it off one dependency chain.
const LANE: usize = 8;
let (a_blocks, a_tail) = a.as_chunks::<{ LANE * 4 }>();
let (b_blocks, b_tail) = b.as_chunks::<{ LANE * 4 }>();
let mut acc = [0i32; 4];
for (x, y) in a_blocks.iter().zip(b_blocks) {
for (lane, slot) in acc.iter_mut().enumerate() {
let mut sum = 0i32;
for k in 0..LANE {
sum += i32::from(x[lane * LANE + k]) * i32::from(y[lane * LANE + k]);
}
*slot += sum;
}
}
let tail: i32 = a_tail
.iter()
.zip(b_tail)
.map(|(&x, &y)| i32::from(x) * i32::from(y))
.sum();
acc[0] + acc[1] + acc[2] + acc[3] + tail
}
/// Magic for [`HnswIndex::graph_to_bytes`].
const GRAPH_MAGIC: &[u8; 4] = b"CHG1";
@@ -174,8 +405,8 @@ pub const HNSW_FORMAT_VERSION: i64 = 2;
/// HDF5 format.
#[derive(Debug, Clone)]
pub struct HnswIndex {
/// All vectors in the index.
vectors: Vec<Vec<f32>>,
/// All vectors in the index, flat and row-major.
vectors: Vectors,
/// Adjacency lists per layer. `graph[layer][node]` = list of neighbor IDs.
graph: Vec<Vec<Vec<usize>>>,
/// Soft-deletion flags, one per node. Deleted nodes remain in the graph for
@@ -214,6 +445,20 @@ impl HnswIndex {
m: usize,
ef_construction: usize,
metric: DistanceMetric,
) -> Self {
Self::build_with(vectors, m, ef_construction, metric, Storage::default())
}
/// Build an index, choosing how the vectors are stored.
///
/// [`Storage::Int8`] keeps them at a quarter of the size; see its docs for
/// what that costs and when it applies.
pub fn build_with(
vectors: &[Vec<f32>],
m: usize,
ef_construction: usize,
metric: DistanceMetric,
storage: Storage,
) -> Self {
assert!(!vectors.is_empty(), "cannot build index from empty vectors");
assert!(m >= 2, "m must be at least 2");
@@ -224,8 +469,11 @@ impl HnswIndex {
let m_max0 = m * 2;
let n = vectors.len();
let prepared: Vec<Vec<f32>> = vectors.iter().map(|v| prepare(v.clone(), metric)).collect();
let vectors: &[Vec<f32>] = &prepared;
let mut prepared = Vectors::new(dim, storage, metric);
for v in vectors {
prepared.push(&prepare(v.clone(), metric));
}
let vectors = &prepared;
// Assign levels to all nodes
let mut node_levels = Vec::with_capacity(n);
@@ -322,9 +570,20 @@ impl HnswIndex {
/// point for incremental [`HnswIndex::insert`] and as the result of
/// [`HnswIndex::compact`] when every vector has been deleted.
pub fn new(m: usize, ef_construction: usize, metric: DistanceMetric) -> Self {
Self::new_with(m, ef_construction, metric, Storage::default())
}
/// [`HnswIndex::new`], choosing how the vectors are stored.
pub fn new_with(
m: usize,
ef_construction: usize,
metric: DistanceMetric,
storage: Storage,
) -> Self {
assert!(m >= 2, "m must be at least 2");
Self {
vectors: Vec::new(),
// The dimension is set by the first insert.
vectors: Vectors::new(0, storage, metric),
graph: Vec::new(),
deleted: Vec::new(),
entry_point: 0,
@@ -351,7 +610,8 @@ impl HnswIndex {
// Seed an empty index.
if id == 0 {
let node_level = assign_level(0, self.m);
self.vectors.push(vector);
self.vectors.set_dim(vector.len());
self.vectors.push(&vector);
self.deleted.push(false);
self.node_levels.push(node_level);
self.graph = (0..=node_level).map(|_| vec![Vec::new(); 1]).collect();
@@ -361,12 +621,12 @@ impl HnswIndex {
assert_eq!(
vector.len(),
self.vectors[0].len(),
self.vectors.dim(),
"insert dimension mismatch"
);
let node_level = assign_level(id, self.m);
self.vectors.push(vector);
self.vectors.push(&vector);
self.deleted.push(false);
self.node_levels.push(node_level);
@@ -387,7 +647,7 @@ impl HnswIndex {
ep = greedy_closest(
&self.vectors,
&self.graph[layer],
&self.vectors[id],
&Target::Node(id),
ep,
self.metric,
);
@@ -400,7 +660,7 @@ impl HnswIndex {
let neighbors = search_layer(
&self.vectors,
&self.graph[layer],
&self.vectors[id],
&Target::Node(id),
ep,
self.ef_construction,
self.metric,
@@ -466,16 +726,25 @@ impl HnswIndex {
pub fn compact(&mut self) -> Vec<Option<usize>> {
let mut mapping = vec![None; self.vectors.len()];
let mut surviving: Vec<Vec<f32>> = Vec::with_capacity(self.active_len());
for (old, v) in self.vectors.iter().enumerate() {
for (old, slot) in mapping.iter_mut().enumerate() {
if !self.deleted[old] {
mapping[old] = Some(surviving.len());
surviving.push(v.clone());
*slot = Some(surviving.len());
surviving.push(self.vectors.row(old));
}
}
// Rebuilding must keep the storage the caller chose; a compaction is
// not the place to silently quadruple the index's memory.
let storage = self.vectors.storage();
*self = if surviving.is_empty() {
Self::new(self.m, self.ef_construction, self.metric)
Self::new_with(self.m, self.ef_construction, self.metric, storage)
} else {
Self::build_with_metric(&surviving, self.m, self.ef_construction, self.metric)
Self::build_with(
&surviving,
self.m,
self.ef_construction,
self.metric,
storage,
)
};
mapping
}
@@ -490,24 +759,22 @@ impl HnswIndex {
/// # Returns
/// A vector of `(id, distance)` pairs sorted by distance (closest first).
pub fn search(&self, query: &[f32], k: usize, ef: usize) -> Vec<(usize, f32)> {
if self.vectors.is_empty() {
if self.vectors.len() == 0 {
return Vec::new();
}
assert_eq!(
query.len(),
self.vectors[0].len(),
"query dimension mismatch"
);
assert_eq!(query.len(), self.vectors.dim(), "query dimension mismatch");
let ef = ef.max(k);
let prepared_query = prepare(query.to_vec(), self.metric);
let query = prepared_query.as_slice();
// Prepared and, for a quantised store, quantised once per search
// rather than once per comparison.
let prepared = self.vectors.query(prepare(query.to_vec(), self.metric));
let target = Target::Query(&prepared);
let mut ep = self.entry_point;
let top_layer = self.graph.len().saturating_sub(1);
// Greedy search from top layer down to layer 1
for layer in (1..=top_layer).rev() {
ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric);
ep = greedy_closest(&self.vectors, &self.graph[layer], &target, ep, self.metric);
}
// Search layer 0 for the ef nearest *live* nodes. Deleted nodes are
@@ -516,7 +783,7 @@ impl HnswIndex {
let candidates = search_layer(
&self.vectors,
&self.graph[0],
query,
&target,
ep,
ef,
self.metric,
@@ -543,14 +810,13 @@ impl HnswIndex {
pub fn to_hdf5_bytes(&self) -> Result<Vec<u8>, FormatError> {
let mut fw = FmtWriter::new();
let n = self.vectors.len();
let dim = if n > 0 { self.vectors[0].len() } else { 0 };
let dim = self.vectors.dim();
// Flatten vectors into a 1D array for storage
let flat_vectors: Vec<f32> = self
.vectors
.iter()
.flat_map(|v| v.iter().copied())
.collect();
let mut flat_vectors: Vec<f32> = Vec::with_capacity(n * dim);
for i in 0..n {
flat_vectors.extend_from_slice(&self.vectors.row(i));
}
let mut group = fw.create_group("ann");
@@ -709,7 +975,9 @@ impl HnswIndex {
};
Ok(Self {
vectors,
// Serialized files carry f32 vectors and no storage tag: a
// quantised index is rebuilt, not loaded.
vectors: Vectors::from_rows(&vectors, Storage::Float32, metric),
graph,
deleted,
entry_point,
@@ -773,6 +1041,16 @@ impl HnswIndex {
/// `bytes` is validated — a corrupt or mismatched graph is an error, never
/// an index that panics or walks out of bounds during a search.
pub fn from_graph_bytes(bytes: &[u8], vectors: Vec<Vec<f32>>) -> Result<Self, FormatError> {
Self::from_graph_bytes_with(bytes, vectors, Storage::default())
}
/// As [`from_graph_bytes`](Self::from_graph_bytes), choosing how the
/// rehydrated vectors are stored.
pub fn from_graph_bytes_with(
bytes: &[u8],
vectors: Vec<Vec<f32>>,
storage: Storage,
) -> Result<Self, FormatError> {
let bad = |what: &str| FormatError::SerializationError(format!("HNSW graph: {what}"));
let body_len = bytes
.len()
@@ -863,7 +1141,14 @@ impl HnswIndex {
}
Ok(Self {
vectors: vectors.into_iter().map(|v| prepare(v, metric)).collect(),
vectors: Vectors::from_rows(
&vectors
.into_iter()
.map(|v| prepare(v, metric))
.collect::<Vec<_>>(),
storage,
metric,
),
graph,
deleted,
entry_point,
@@ -882,16 +1167,17 @@ impl HnswIndex {
/// Returns true if the index is empty.
pub fn is_empty(&self) -> bool {
self.vectors.is_empty()
self.vectors.len() == 0
}
/// How this index stores its copy of the vectors.
pub fn storage(&self) -> Storage {
self.vectors.storage()
}
/// Returns the dimension of vectors in the index.
pub fn dimension(&self) -> usize {
if self.vectors.is_empty() {
0
} else {
self.vectors[0].len()
}
self.vectors.dim()
}
/// Returns the number of layers in the graph.
@@ -916,17 +1202,17 @@ impl HnswIndex {
/// Greedy search: find the single closest node to `query` starting from `ep`.
fn greedy_closest(
vectors: &[Vec<f32>],
vectors: &Vectors,
layer: &[Vec<usize>],
query: &[f32],
target: &Target<'_>,
mut ep: usize,
metric: DistanceMetric,
) -> usize {
let mut best_dist = compute_distance(query, &vectors[ep], metric);
let mut best_dist = vectors.dist_to(target, ep, metric);
loop {
let mut changed = false;
for &neighbor in &layer[ep] {
let d = compute_distance(query, &vectors[neighbor], metric);
let d = vectors.dist_to(target, neighbor, metric);
if d < best_dist {
best_dist = d;
ep = neighbor;
@@ -950,15 +1236,15 @@ fn greedy_closest(
/// instead meant a query whose neighbourhood had been deleted got back fewer
/// than `k` results, or none, however many live records were nearby.
fn search_layer(
vectors: &[Vec<f32>],
vectors: &Vectors,
layer: &[Vec<usize>],
query: &[f32],
target: &Target<'_>,
ep: usize,
ef: usize,
metric: DistanceMetric,
skip: Option<&[bool]>,
) -> Vec<Candidate> {
let ep_dist = compute_distance(query, &vectors[ep], metric);
let ep_dist = vectors.dist_to(target, ep, metric);
// Min-heap of candidates to explore
let mut candidates = BinaryHeap::new();
@@ -980,7 +1266,7 @@ fn search_layer(
visited.begin(vectors.len());
visited.insert(ep);
search_layer_visit(
vectors, layer, query, ef, metric, skip, visited, candidates, results,
vectors, layer, target, ef, metric, skip, visited, candidates, results,
)
})
}
@@ -1023,9 +1309,9 @@ thread_local! {
#[allow(clippy::too_many_arguments)]
fn search_layer_visit(
vectors: &[Vec<f32>],
vectors: &Vectors,
layer: &[Vec<usize>],
query: &[f32],
target: &Target<'_>,
ef: usize,
metric: DistanceMetric,
skip: Option<&[bool]>,
@@ -1044,7 +1330,7 @@ fn search_layer_visit(
continue;
}
let d = compute_distance(query, &vectors[neighbor], metric);
let d = vectors.dist_to(target, neighbor, metric);
let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance);
if d < furthest_dist || results.len() < ef {
@@ -1095,7 +1381,7 @@ fn search_layer_visit(
/// remaining slots are then filled with the closest rejected candidates, so a
/// node is never left under-connected.
fn select_neighbors(
vectors: &[Vec<f32>],
vectors: &Vectors,
candidates: &[(usize, f32)],
max_conn: usize,
metric: DistanceMetric,
@@ -1111,7 +1397,7 @@ fn select_neighbors(
}
let diverse = selected
.iter()
.all(|&s| compute_distance(&vectors[id], &vectors[s], metric) > dist_to_node);
.all(|&s| vectors.dist(id, s, metric) > dist_to_node);
if diverse {
selected.push(id);
} else {
@@ -1136,7 +1422,7 @@ fn batch_len(linked: usize) -> usize {
/// layers, found by searching the graph as it currently stands.
#[allow(clippy::too_many_arguments)]
fn plan_batch(
vectors: &[Vec<f32>],
vectors: &Vectors,
graph: &[Vec<Vec<usize>>],
node_levels: &[usize],
batch: std::ops::Range<usize>,
@@ -1150,7 +1436,7 @@ fn plan_batch(
let mut ep = entry_point;
// Phase 1: greedy descent from the top layer down to node_level + 1.
for layer in (node_level + 1..=ep_level).rev() {
ep = greedy_closest(vectors, &graph[layer], &vectors[i], ep, metric);
ep = greedy_closest(vectors, &graph[layer], &Target::Node(i), ep, metric);
}
// Phase 2: search and select on every layer the node lives on.
let mut plan = Vec::with_capacity(node_level.min(ep_level) + 1);
@@ -1159,7 +1445,7 @@ fn plan_batch(
let neighbors = search_layer(
vectors,
&graph[layer],
&vectors[i],
&Target::Node(i),
ep,
ef_construction,
metric,
@@ -1186,7 +1472,7 @@ fn plan_batch(
/// Prune every `(layer, node)` neighbour list in `overflowed` back to its
/// limit. Each list belongs to a different node, so they are independent.
fn prune_overflowed(
vectors: &[Vec<f32>],
vectors: &Vectors,
graph: &mut [Vec<Vec<usize>>],
overflowed: Vec<(usize, usize)>,
(m, m_max0): (usize, usize),
@@ -1226,7 +1512,7 @@ const PARALLEL_MIN: usize = 8;
/// prunes is too fine-grained to parallelise profitably — measured 1.45x on 16
/// cores; bulk builds batch their pruning instead, see `prune_overflowed`.)
fn link_back(
vectors: &[Vec<f32>],
vectors: &Vectors,
layer: &mut [Vec<usize>],
new_id: usize,
selected: &[usize],
@@ -1248,7 +1534,7 @@ fn link_back(
/// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`].
fn prune_connections(
vectors: &[Vec<f32>],
vectors: &Vectors,
neighbors: &mut Vec<usize>,
node: usize,
max_conn: usize,
@@ -1259,7 +1545,7 @@ fn prune_connections(
}
let mut scored: Vec<(usize, f32)> = neighbors
.iter()
.map(|&n| (n, compute_distance(&vectors[node], &vectors[n], metric)))
.map(|&n| (n, vectors.dist(node, n, metric)))
.collect();
scored.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0)));
*neighbors = select_neighbors(vectors, &scored, max_conn, metric);
@@ -1519,6 +1805,88 @@ mod tests {
assert!(recall >= 0.95, "incremental recall@10 = {recall}");
}
#[test]
fn int8_storage_needs_an_exact_re_score_to_match_f32() {
// Cosine only: rows are unit-length, so a quantised dot product
// reconstructs the similarity directly.
//
// Not the `clustered` generator: its clusters are far tighter than any
// real embedding, so neighbours sit closer together than the
// quantisation error and top-10 identity there is noise — that would
// measure the fixture, not the storage.
let mut vectors = make_random_vectors(3060, 128, 5);
let queries = vectors.split_off(3000);
let f32_index =
HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Float32);
let quantised =
HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::Cosine, Storage::Int8);
assert_eq!(quantised.storage(), Storage::Int8);
// Ground truth, not the f32 index's answers: re-scoring can beat that
// index, and measuring against it would score being right as drift.
let truth: Vec<Vec<usize>> = queries
.iter()
.map(|q| {
let mut d: Vec<(usize, f32)> = vectors
.iter()
.enumerate()
.map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::Cosine)))
.collect();
d.sort_by(|a, b| a.1.total_cmp(&b.1));
d[..10].iter().map(|x| x.0).collect()
})
.collect();
let recall = |got: &dyn Fn(&[f32]) -> Vec<usize>| -> f64 {
let mut hits = 0;
for (q, want) in queries.iter().zip(&truth) {
hits += got(q).iter().filter(|id| want.contains(id)).count();
}
hits as f64 / (10 * queries.len()) as f64
};
let exact_recall = recall(&|q| f32_index.search(q, 10, 64).iter().map(|r| r.0).collect());
let raw_recall = recall(&|q| quantised.search(q, 10, 64).iter().map(|r| r.0).collect());
// Quantised distances alone cost recall, and `ef` cannot buy it back:
// the loss is in the distances, not in the graph.
assert!(
raw_recall < exact_recall,
"int8 alone should cost recall: {raw_recall} vs {exact_recall}"
);
// Re-scoring a wider candidate pool against the exact vectors — what a
// caller holding them (the agent's embedding cache) does — puts it
// back, because only the *ordering* was approximate.
let rescored_recall = recall(&|q| {
let mut pool: Vec<(usize, f32)> = quantised
.search(q, 40, 64)
.into_iter()
.map(|(id, _)| {
(
id,
compute_distance(q, &vectors[id], DistanceMetric::Cosine),
)
})
.collect();
pool.sort_by(|a, b| a.1.total_cmp(&b.1));
pool.truncate(10);
pool.into_iter().map(|p| p.0).collect()
});
assert!(
rescored_recall >= exact_recall - 0.01,
"int8 + exact re-score should match f32: {rescored_recall} vs {exact_recall} (raw {raw_recall})"
);
}
#[test]
fn int8_storage_falls_back_to_f32_for_non_cosine_metrics() {
// L2 distance is not recoverable from a quantised dot product, so the
// store silently stays f32 rather than returning wrong distances.
let vectors = clustered(100, 8, 5, 3);
let index = HnswIndex::build_with(&vectors, 8, 40, DistanceMetric::L2, Storage::Int8);
assert_eq!(index.storage(), Storage::Float32);
}
#[test]
fn deletions_near_the_query_do_not_shrink_or_degrade_results() {
let mut vectors = clustered(2040, 16, 20, 11);
@@ -1650,6 +2018,7 @@ mod tests {
vec![1.2, 0.0], // 3
vec![-2.0, 0.0], // 4
];
let store = Vectors::from_rows(&vectors, Storage::Float32, DistanceMetric::L2);
let scored: Vec<(usize, f32)> = (1..5)
.map(|i| {
(
@@ -1659,12 +2028,12 @@ mod tests {
})
.collect();
assert_eq!(
select_neighbors(&vectors, &scored, 2, DistanceMetric::L2),
select_neighbors(&store, &scored, 2, DistanceMetric::L2),
[1, 4]
);
// Spare capacity is filled with the closest rejected candidates.
assert_eq!(
select_neighbors(&vectors, &scored, 3, DistanceMetric::L2),
select_neighbors(&store, &scored, 3, DistanceMetric::L2),
[1, 4, 2]
);
}
@@ -1790,7 +2159,7 @@ mod tests {
// Verify vectors match
for i in 0..loaded.len() {
assert_eq!(loaded.vectors[i], index.vectors[i]);
assert_eq!(loaded.vectors.row(i), index.vectors.row(i));
}
}
+1 -1
View File
@@ -5,4 +5,4 @@
mod hnsw;
pub use hnsw::{DistanceMetric, HnswIndex};
pub use hnsw::{DistanceMetric, HnswIndex, Storage};
@@ -24,7 +24,7 @@
use std::time::{Duration, Instant};
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
use clawhdf5_ann::{DistanceMetric, HnswIndex};
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
const DIM: usize = 384;
const K: usize = 10;
@@ -84,6 +84,22 @@ struct Dataset {
/// that appears only on clustered data points at graph connectivity.
static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
/// `--int8`: build the HNSW index over int8-quantised vectors (a quarter of
/// the memory) instead of f32, to price the recall it costs.
static INT8: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
/// `--rerank`: re-score the candidate pool against the exact vectors before
/// taking the top K.
static RERANK: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
fn storage() -> Storage {
if INT8.load(std::sync::atomic::Ordering::Relaxed) {
Storage::Int8
} else {
Storage::Float32
}
}
fn make_dataset(n: usize, seed: u64) -> Dataset {
let mut rng = Rng(seed);
if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) {
@@ -169,6 +185,11 @@ fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String {
// Measurement helpers
// ---------------------------------------------------------------------------
/// Exact cosine distance between unit-length vectors.
fn exact_dist(a: &[f32], b: &[f32]) -> f32 {
1.0 - a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>()
}
fn exact_top_k(vectors: &[Vec<f32>], query: &[f32], k: usize) -> Vec<usize> {
// Vectors are unit length, so cosine order == dot-product order.
let mut scored: Vec<(usize, f32)> = vectors
@@ -270,11 +291,12 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
.collect();
let started = Instant::now();
let index = HnswIndex::build_with_metric(
let index = HnswIndex::build_with(
&data.vectors,
HNSW_M,
HNSW_EF_CONSTRUCTION,
DistanceMetric::Cosine,
storage(),
);
let build = started.elapsed();
@@ -291,7 +313,8 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
);
println!(
"\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}\n"
"\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}, storage = {:?}\n",
index.storage()
);
println!(
"build: {:.1} ms ({:.0} vectors/s) · exact scan: {:.0} QPS, p50 {:.0} µs\n",
@@ -302,12 +325,26 @@ fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
);
println!("| ef | recall@{K} | QPS | p50 µs | p99 µs |");
println!("|---:|---:|---:|---:|---:|");
// With a quantised index the distances it returns are approximate, so
// the candidates are re-scored against the exact vectors the caller
// already holds (in the agent, the embedding cache) before taking the
// top K. `--rerank` prices that: it costs one exact distance per
// candidate and is what decides whether int8 is usable.
let rerank = RERANK.load(std::sync::atomic::Ordering::Relaxed);
let pool = if rerank { K * 4 } else { K };
for ef in EF_VALUES {
let mut hits = 0usize;
let mut samples = Vec::with_capacity(data.queries.len());
for (q, want) in data.queries.iter().zip(&truth) {
let t = Instant::now();
let got = index.search(q, K, ef);
let mut got = index.search(q, pool, ef.max(pool));
if rerank {
for cand in &mut got {
cand.1 = exact_dist(&data.vectors[cand.0], q);
}
got.select_nth_unstable_by(K - 1, |a, b| a.1.total_cmp(&b.1));
got.truncate(K);
}
samples.push(t.elapsed());
hits += got.iter().filter(|(id, _)| want.contains(id)).count();
}
@@ -445,11 +482,12 @@ fn fusion_study(n: usize) {
.map(|(i, c)| text_for(*c, i, &mut rng))
.collect();
let bm25 = BM25Index::build(&texts, &vec![0u8; n]);
let index = HnswIndex::build_with_metric(
let index = HnswIndex::build_with(
&data.vectors,
HNSW_M,
HNSW_EF_CONSTRUCTION,
DistanceMetric::Cosine,
storage(),
);
let vec_pool = (K * 8).max(64);
@@ -525,7 +563,9 @@ fn bench_footprint(n: usize) {
.collect();
let after_entries = heap_bytes();
let mut mem = HDF5Memory::create(MemoryConfig::new(path, "bench", DIM)).unwrap();
let mut config = MemoryConfig::new(path, "bench", DIM);
config.quantized_index = INT8.load(std::sync::atomic::Ordering::Relaxed);
let mut mem = HDF5Memory::create(config).unwrap();
mem.save_batch(entries).unwrap();
let after_store = heap_bytes();
@@ -571,6 +611,14 @@ fn main() {
}
return;
}
if args.iter().any(|a| a == "--int8") {
INT8.store(true, std::sync::atomic::Ordering::Relaxed);
println!("(int8-quantised index vectors)");
}
if args.iter().any(|a| a == "--rerank") {
RERANK.store(true, std::sync::atomic::Ordering::Relaxed);
println!("(candidates re-scored against exact vectors)");
}
if args.iter().any(|a| a == "--uniform") {
UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed);
println!("(uniform random data)");
+12 -1
View File
@@ -28,6 +28,10 @@ enum Commands {
/// Enable write-ahead log
#[arg(long)]
wal: bool,
/// Store the vector index's copy of the embeddings as int8, roughly
/// halving a loaded store's memory at about 13% fewer queries/second
#[arg(long)]
quantized_index: bool,
},
/// Save a memory entry (reads JSON from stdin or --json)
Save {
@@ -88,9 +92,15 @@ fn main() {
fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
match cli.command {
Commands::Create { agent_id, dim, wal } => {
Commands::Create {
agent_id,
dim,
wal,
quantized_index,
} => {
let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim);
config.wal_enabled = wal;
config.quantized_index = quantized_index;
let mem = HDF5Memory::create(config)?;
let j = serde_json::json!({
"status": "created",
@@ -98,6 +108,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
"agent_id": agent_id,
"embedding_dim": dim,
"wal_enabled": wal,
"quantized_index": quantized_index,
"count": mem.count(),
});
println!("{}", serde_json::to_string_pretty(&j)?);
+15 -4
View File
@@ -15,7 +15,6 @@ use crate::filter_pipeline::{
FilterDescription, FilterPipeline,
};
use crate::filters::compress_chunk;
/// Round a file offset up to the next cache-line boundary.
///
/// This ensures chunk data starts at an address that is a multiple of the
@@ -928,6 +927,7 @@ pub fn write_selection_to_buffer(
#[cfg(test)]
mod tests {
use super::*;
use crate::chunked_read::read_chunked_data;
use crate::data_layout::DataLayout;
@@ -1512,9 +1512,20 @@ mod tests {
// ---- h5py round-trip tests for chunked writes ----
/// The Python interpreter to drive interop checks with.
///
/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py,
/// which on a PEP 668 "externally managed" system is the only place it
/// can be installed. Without it the suite silently skips, and a silent
/// skip here is how a datatype bug once reached a release.
#[cfg(feature = "std")]
fn python() -> String {
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
}
#[cfg(feature = "std")]
fn h5py_available() -> bool {
std::process::Command::new("python3")
std::process::Command::new(python())
.args(["-c", "import h5py"])
.output()
.map(|o| o.status.success())
@@ -1526,10 +1537,10 @@ mod tests {
if !h5py_available() {
panic!("h5py not installed — skipping interop test");
}
let o = std::process::Command::new("python3")
let o = std::process::Command::new(python())
.args(["-c", script])
.output()
.expect("python3");
.expect("python interpreter");
if !o.status.success() {
panic!("h5py: {}", String::from_utf8_lossy(&o.stderr));
}
@@ -2,6 +2,15 @@
use clawhdf5_format::data_read::{read_object_references, read_region_references};
use clawhdf5_format::datatype::{Datatype, ReferenceType};
/// The Python interpreter to drive interop checks with.
///
/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, which
/// on a PEP 668 "externally managed" system is the only place it can be
/// installed. Without it the suite silently skips, and a silent skip here is
/// how a datatype bug once reached a release.
fn python() -> String {
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
}
#[test]
fn object_ref_single_valid() {
@@ -173,7 +182,7 @@ print('ok')
"#,
path.display()
);
let output = std::process::Command::new("python3")
let output = std::process::Command::new(python())
.args(["-c", &script])
.output();
@@ -4,9 +4,18 @@
//! (and vice versa). They require python3 + h5py to be installed.
use clawhdf5_format::file_writer::{AttrValue, CompoundTypeBuilder, EnumTypeBuilder, FileWriter};
/// The Python interpreter to drive interop checks with.
///
/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, which
/// on a PEP 668 "externally managed" system is the only place it can be
/// installed. Without it the suite silently skips, and a silent skip here is
/// how a datatype bug once reached a release.
fn python() -> String {
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
}
fn h5py_available() -> bool {
std::process::Command::new("python3")
std::process::Command::new(python())
.args(["-c", "import h5py"])
.output()
.map(|o| o.status.success())
@@ -17,10 +26,10 @@ fn h5py_read(_path: &std::path::Path, script: &str) -> String {
if !h5py_available() {
panic!("h5py not installed — skipping interop test");
}
let o = std::process::Command::new("python3")
let o = std::process::Command::new(python())
.args(["-c", script])
.output()
.expect("python3");
.expect("python interpreter");
if !o.status.success() {
panic!("h5py: {}", String::from_utf8_lossy(&o.stderr));
}
+12 -3
View File
@@ -9,6 +9,15 @@ use clawhdf5_netcdf4::{AttrValue, NetCDF4File};
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/// The Python interpreter to drive interop checks with.
///
/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, which
/// on a PEP 668 "externally managed" system is the only place it can be
/// installed. Without it the suite silently skips, and a silent skip here is
/// how a datatype bug once reached a release.
fn python() -> String {
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
}
/// When `CLAWHDF5_REQUIRE_INTEROP=1` (set in CI), a missing Python dependency
/// is a test failure instead of a silent skip.
@@ -17,7 +26,7 @@ fn interop_required() -> bool {
}
fn netcdf4_python_available() -> bool {
Command::new("python3")
Command::new(python())
.args(["-c", "import netCDF4; print(netCDF4.__version__)"])
.output()
.map(|o| o.status.success())
@@ -25,7 +34,7 @@ fn netcdf4_python_available() -> bool {
}
fn xarray_available() -> bool {
Command::new("python3")
Command::new(python())
.args(["-c", "import xarray; print(xarray.__version__)"])
.output()
.map(|o| o.status.success())
@@ -59,7 +68,7 @@ macro_rules! skip_if_no_xarray {
}
fn run_python(script: &str) {
let output = Command::new("python3")
let output = Command::new(python())
.args(["-c", script])
.output()
.expect("failed to run python3");
+12 -3
View File
@@ -9,6 +9,15 @@ use clawhdf5::{AttrValue, CompoundTypeBuilder, DType, File, FileBuilder};
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/// The Python interpreter to drive interop checks with.
///
/// `CLAWHDF5_PYTHON` lets these run against a virtualenv holding h5py, which
/// on a PEP 668 "externally managed" system is the only place it can be
/// installed. Without it the suite silently skips, and a silent skip here is
/// how a datatype bug once reached a release.
fn python() -> String {
std::env::var("CLAWHDF5_PYTHON").unwrap_or_else(|_| "python3".to_string())
}
/// When `CLAWHDF5_REQUIRE_INTEROP=1` (set in CI), a missing Python dependency
/// is a test failure instead of a silent skip.
@@ -17,7 +26,7 @@ fn interop_required() -> bool {
}
fn python_available() -> bool {
Command::new("python3")
Command::new(python())
.args(["-c", "import h5py; print(h5py.__version__)"])
.output()
.map(|o| o.status.success())
@@ -39,7 +48,7 @@ macro_rules! skip_if_no_python {
/// Run a Python script and panic if it fails.
fn run_python(script: &str) {
let output = Command::new("python3")
let output = Command::new(python())
.args(["-c", script])
.output()
.expect("failed to run python3");
@@ -52,7 +61,7 @@ fn run_python(script: &str) {
/// Run a Python script and return stdout as a trimmed string.
fn run_python_output(script: &str) -> String {
let output = Command::new("python3")
let output = Command::new(python())
.args(["-c", script])
.output()
.expect("failed to run python3");
+5
View File
@@ -364,6 +364,11 @@ cargo install --path crates/clawhdf5-cli
clawhdf5 --path agent.h5 create --agent-id my-agent --dim 384 --wal
```
Add `--quantized-index` to store the vector index's copy of the embeddings as
int8. That roughly halves a loaded store's memory at about 13% fewer queries
per second, with recall unchanged — the query path re-scores candidates
against the exact embeddings. The setting is recorded in the file.
Output:
```json
{
+15 -2
View File
@@ -20,6 +20,13 @@
set -uo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
# Interop suites drive a Python interpreter. On a PEP 668 "externally managed"
# system h5py can only live in a virtualenv, so pick one up here — before any
# test step, since the non-ignored interop suites read the same variable.
if [ -z "${CLAWHDF5_PYTHON:-}" ] && [ -x "$SCRIPT_DIR/../.venv/bin/python" ]; then
export CLAWHDF5_PYTHON="$SCRIPT_DIR/../.venv/bin/python"
fi
PASS=0
FAIL=0
STEPS=()
@@ -85,12 +92,18 @@ run_step "cargo test (ann parallel)" cargo test \
# 5. Python interop suites. The h5py writer tests are #[ignore]d so a plain
# `cargo test` stays hermetic; run them explicitly here.
if python3 -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then
# On a PEP 668 "externally managed" system h5py can only live in a
# virtualenv, so honour CLAWHDF5_PYTHON (and a local .venv) rather than
# skipping — the tests read the same variable.
PYTHON="${CLAWHDF5_PYTHON:-python3}"
if "$PYTHON" -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then
run_step "h5py interop (format, ignored tests)" cargo test \
-p clawhdf5-format --test writer_h5py_tests -- --include-ignored
else
echo ""
echo "==> [h5py interop] SKIPPED: python3 with h5py not available"
echo "==> [h5py interop] SKIPPED: no h5py in $PYTHON"
echo " (set CLAWHDF5_PYTHON=/path/to/venv/bin/python, or create .venv;"
echo " CLAWHDF5_REQUIRE_INTEROP=1 makes this a failure instead)"
STEPS+=("SKIP: h5py interop (format, ignored tests)")
fi