use crate::embedding::single_vector::{
EMBEDDING_DIM, SingleVectorEmbedder, dequantize_int8, l2_normalize, quantize_int8,
};
use crate::search::dense_backend::{DenseBackend, DenseHit, DenseIndexBuilder};
use crate::search::simd::{dot_f32, dot_i8};
use crate::types::ScoredChunkId;
use anyhow::Result;
use instant_distance::{Builder, HnswMap, Point, Search};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
const HNSW_MIN_VECTORS_DEFAULT: usize = 2_000;
const INDEX_FILE: &str = "index.bin";
const VECTORS_FILE: &str = "store.vecs";
const ENCODE_BATCH: usize = 32;
#[derive(Debug, Clone, Copy)]
pub struct HnswParams {
pub m: usize,
pub ef_construction: usize,
pub ef_search: usize,
pub rescore_k: usize,
}
impl Default for HnswParams {
fn default() -> Self {
Self {
m: 16,
ef_construction: 200,
ef_search: 64,
rescore_k: 0,
}
}
}
impl HnswParams {
pub fn preset(name: &str) -> Self {
match name.trim().to_ascii_lowercase().as_str() {
"high_recall" => Self {
m: 32,
ef_construction: 400,
ef_search: 200,
rescore_k: 0,
},
"low_latency" => Self {
m: 16,
ef_construction: 200,
ef_search: 32,
rescore_k: 0,
},
"memory_optimized" => Self {
m: 8,
ef_construction: 100,
ef_search: 48,
rescore_k: 0,
},
_ => Self::default(), }
}
pub fn resolve(preset: &str, ef_search_override: usize, rescore_k_override: usize) -> Self {
let mut p = Self::preset(preset);
if ef_search_override != 0 {
p.ef_search = ef_search_override;
}
p.rescore_k = rescore_k_override;
p
}
}
fn hnsw_min_vectors() -> usize {
crate::config::env_usize("SEMANTEX_HNSW_MIN_VECTORS", HNSW_MIN_VECTORS_DEFAULT)
}
#[derive(Debug, Clone)]
struct HnswPoint(Vec<f32>);
impl Point for HnswPoint {
fn distance(&self, other: &Self) -> f32 {
1.0 - dot_f32(&self.0, &other.0)
}
}
type HnswGraph = HnswMap<HnswPoint, usize>;
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct Int8VectorStore {
pub dim: usize,
pub scales: Vec<f32>,
pub vectors: Vec<Vec<i8>>,
pub chunk_ids: Vec<u64>,
#[serde(default)]
pub tombstoned: Vec<bool>,
}
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
struct StoreIndex {
dim: usize,
chunk_ids: Vec<u64>,
scales: Vec<f32>,
tombstoned: Vec<bool>,
}
fn write_index_atomic(dir: &Path, idx: &StoreIndex) -> Result<()> {
std::fs::create_dir_all(dir)?;
let final_path = dir.join(INDEX_FILE);
let tmp_path = dir.join(format!(".index.{}.tmp", std::process::id()));
std::fs::write(&tmp_path, postcard::to_stdvec(idx)?)?;
std::fs::rename(&tmp_path, &final_path)?;
Ok(())
}
impl Int8VectorStore {
fn len(&self) -> usize {
self.chunk_ids.len()
}
fn push(&mut self, chunk_id: u64, q: Vec<i8>, scale: f32) -> usize {
let idx = self.chunk_ids.len();
self.vectors.push(q);
self.scales.push(scale);
self.chunk_ids.push(chunk_id);
self.tombstoned.push(false);
idx
}
fn is_tombstoned(&self, i: usize) -> bool {
self.tombstoned.get(i).copied().unwrap_or(false)
}
fn live_count(&self) -> usize {
(0..self.chunk_ids.len())
.filter(|&i| !self.is_tombstoned(i))
.count()
}
fn dequant_row(&self, i: usize) -> Vec<f32> {
let mut f = dequantize_int8(&self.vectors[i], self.scales[i]);
l2_normalize(&mut f);
f
}
fn brute_force(&self, query: &[f32], k: usize, rescore_k: usize) -> Vec<DenseHit> {
if self.vectors.is_empty() {
return Vec::new();
}
let shortlist = rescore_k.max(k).max(1);
let (q8, _qs) = quantize_int8(query);
let mut scored: Vec<(usize, f32)> = self
.vectors
.iter()
.enumerate()
.filter(|(i, _)| !self.is_tombstoned(*i))
.map(|(i, v)| (i, dot_i8(&q8, v)))
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
scored.truncate(shortlist);
let positions: Vec<usize> = scored.into_iter().map(|(i, _)| i).collect();
self.fp32_rescore(query, &positions, k)
}
fn fp32_rescore(&self, query: &[f32], positions: &[usize], k: usize) -> Vec<DenseHit> {
let mut hits: Vec<DenseHit> = positions
.iter()
.filter_map(|&i| {
let scale = *self.scales.get(i)?;
let mut f = dequantize_int8(self.vectors.get(i)?, scale);
l2_normalize(&mut f);
let score = dot_f32(query, &f);
Some(ScoredChunkId::new(self.chunk_ids[i], score))
})
.collect();
hits.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
hits.truncate(k.max(1));
hits
}
fn vectors_for(&self, chunk_ids: &[u64]) -> Vec<(u64, Vec<f32>)> {
use std::collections::HashMap;
let pos: HashMap<u64, usize> = self
.chunk_ids
.iter()
.enumerate()
.filter(|(i, _)| !self.is_tombstoned(*i))
.map(|(i, &id)| (id, i))
.collect();
chunk_ids
.iter()
.filter_map(|id| {
let &i = pos.get(id)?;
Some((*id, self.dequant_row(i)))
})
.collect()
}
fn load_index_only(dir: &Path) -> Result<StoreIndex> {
match std::fs::read(dir.join(INDEX_FILE)) {
Ok(bytes) => Ok(postcard::from_bytes(&bytes)?),
Err(_) => Ok(StoreIndex {
dim: EMBEDDING_DIM,
..Default::default()
}),
}
}
fn load_split(dir: &Path) -> Result<Self> {
let idx = Self::load_index_only(dir)?;
let count = idx.chunk_ids.len();
let vectors: Vec<Vec<i8>> = if count == 0 {
Vec::new()
} else {
let bytes = std::fs::read(dir.join(VECTORS_FILE))?;
let need = count * idx.dim;
anyhow::ensure!(
bytes.len() >= need,
"coderank-hnsw {VECTORS_FILE} truncated: have {} bytes, need {need} for \
{count} rows of dim {}",
bytes.len(),
idx.dim,
);
(0..count)
.map(|i| {
let start = i * idx.dim;
bytes[start..start + idx.dim]
.iter()
.map(|&b| b as i8)
.collect()
})
.collect()
};
Ok(Self {
dim: idx.dim,
scales: idx.scales,
vectors,
chunk_ids: idx.chunk_ids,
tombstoned: idx.tombstoned,
})
}
fn to_index(&self) -> StoreIndex {
StoreIndex {
dim: self.dim,
chunk_ids: self.chunk_ids.clone(),
scales: self.scales.clone(),
tombstoned: self.tombstoned.clone(),
}
}
fn save_split(&self, dir: &Path) -> Result<()> {
std::fs::create_dir_all(dir)?;
let mut payload = Vec::with_capacity(self.vectors.len() * self.dim);
for row in &self.vectors {
payload.extend(row.iter().map(|&b| b as u8));
}
std::fs::write(dir.join(VECTORS_FILE), &payload)?;
write_index_atomic(dir, &self.to_index())
}
fn append_split(
dir: &Path,
mut idx: StoreIndex,
new_ids: &[u64],
new_scales: &[f32],
new_rows: &[Vec<i8>],
) -> Result<()> {
use std::io::Write;
std::fs::create_dir_all(dir)?;
if !new_rows.is_empty() {
let mut f = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(dir.join(VECTORS_FILE))?;
for row in new_rows {
let bytes: Vec<u8> = row.iter().map(|&b| b as u8).collect();
f.write_all(&bytes)?;
}
f.sync_data()?;
}
idx.chunk_ids.extend_from_slice(new_ids);
idx.scales.extend_from_slice(new_scales);
idx.tombstoned
.extend(std::iter::repeat_n(false, new_ids.len()));
write_index_atomic(dir, &idx)
}
fn tombstone_split(
dir: &Path,
mut idx: StoreIndex,
remove: &std::collections::HashSet<u64>,
) -> Result<()> {
while idx.tombstoned.len() < idx.chunk_ids.len() {
idx.tombstoned.push(false);
}
for i in 0..idx.chunk_ids.len() {
if remove.contains(&idx.chunk_ids[i]) {
idx.tombstoned[i] = true;
}
}
write_index_atomic(dir, &idx)
}
}
pub struct HnswDenseIndex {
store: Int8VectorStore,
params: HnswParams,
graph: Option<HnswGraph>,
baked_ef_search: usize,
}
impl HnswDenseIndex {
fn baked_ef_search(params: HnswParams) -> usize {
params.ef_search.max(params.rescore_k).max(1)
}
fn build_graph(store: &Int8VectorStore, params: HnswParams) -> Option<HnswGraph> {
if store.live_count() < hnsw_min_vectors() {
return None;
}
let live: Vec<usize> = (0..store.len())
.filter(|&i| !store.is_tombstoned(i))
.collect();
let points: Vec<HnswPoint> = live
.iter()
.map(|&i| HnswPoint(store.dequant_row(i)))
.collect();
let values: Vec<usize> = live;
let graph = Builder::default()
.ef_construction(params.ef_construction.max(1))
.ef_search(Self::baked_ef_search(params))
.build(points, values);
Some(graph)
}
fn load(dir: &Path, params: HnswParams) -> Result<Self> {
let store = Int8VectorStore::load_split(dir)?;
let graph = Self::build_graph(&store, params);
Ok(Self {
store,
params,
graph,
baked_ef_search: Self::baked_ef_search(params),
})
}
fn save(store: &Int8VectorStore, dir: &Path) -> Result<()> {
store.save_split(dir)
}
fn search(&self, query: &[f32], k: usize) -> Vec<DenseHit> {
if self.store.live_count() == 0 {
return Vec::new();
}
let rescore_k = if self.params.rescore_k == 0 {
4 * k
} else {
self.params.rescore_k
}
.max(k);
match &self.graph {
None => self.store.brute_force(query, k, rescore_k),
Some(graph) => {
let effective = rescore_k.min(self.baked_ef_search);
if rescore_k > self.baked_ef_search {
tracing::debug!(
rescore_k,
baked_ef_search = self.baked_ef_search,
"coderank-hnsw rescore window capped by baked ef_search; \
raise SEMANTEX_HNSW_EF_SEARCH or use a higher-recall preset"
);
}
let mut search = Search::default();
let qpoint = HnswPoint(query.to_vec());
let positions: Vec<usize> = graph
.search(&qpoint, &mut search)
.take(effective)
.map(|item| *item.value)
.collect();
self.store.fp32_rescore(query, &positions, k)
}
}
}
}
pub struct CoderankHnswBackend {
index: HnswDenseIndex,
embedder: &'static SingleVectorEmbedder,
}
impl CoderankHnswBackend {
pub const NAME: &'static str = "coderank-hnsw";
pub fn open(dir: &Path, model_dir: &Path, params: HnswParams) -> Result<Self> {
let index = HnswDenseIndex::load(dir, params)?;
let embedder = SingleVectorEmbedder::global(model_dir)?;
Ok(Self { index, embedder })
}
#[cfg(test)]
fn open_for_test(dir: &Path) -> Self {
let index = HnswDenseIndex::load(dir, HnswParams::default()).expect("test store loads");
let tmp = std::env::temp_dir();
let embedder =
SingleVectorEmbedder::global(&tmp).expect("global embedder (lazy; no session built)");
Self { index, embedder }
}
}
impl DenseBackend for CoderankHnswBackend {
fn name(&self) -> &'static str {
Self::NAME
}
fn search(&self, query: &str, k: usize) -> Result<Vec<DenseHit>> {
let q = self.embedder.encode_query(query)?;
Ok(self.index.search(&q, k))
}
fn search_with_subset(&self, query: &str, k: usize, subset: &[u64]) -> Result<Vec<DenseHit>> {
if subset.is_empty() {
return Ok(Vec::new());
}
let q = self.embedder.encode_query(query)?;
let wide = self.index.search(&q, (k * 8).max(64));
let allow: std::collections::HashSet<u64> = subset.iter().copied().collect();
let mut filtered: Vec<DenseHit> = wide
.into_iter()
.filter(|h| allow.contains(&h.chunk_id))
.collect();
filtered.truncate(k);
Ok(filtered)
}
fn positional_chunk_ids(&self) -> Option<&[u64]> {
None
}
fn embed_text_vector(&self, query: &str) -> Option<Vec<f32>> {
self.embedder.encode_query(query).ok()
}
fn embed_doc_vectors(&self, chunk_ids: &[u64]) -> Option<Vec<(u64, Vec<f32>)>> {
Some(self.index.store.vectors_for(chunk_ids))
}
}
pub struct CoderankHnswIndexBuilder {
dir: PathBuf,
#[allow(dead_code)]
params: HnswParams,
models_dir: PathBuf,
dense_context: bool,
context_by_id: std::collections::HashMap<u64, String>,
store: Int8VectorStore,
}
impl CoderankHnswIndexBuilder {
pub fn new(dir: &Path, params: HnswParams) -> Self {
Self {
dir: dir.to_path_buf(),
params,
models_dir: crate::config::SemantexConfig::default().models_dir(),
dense_context: false,
context_by_id: std::collections::HashMap::new(),
store: Int8VectorStore {
dim: EMBEDDING_DIM,
..Default::default()
},
}
}
pub fn with_models_dir(mut self, models_dir: PathBuf) -> Self {
self.models_dir = models_dir;
self
}
pub fn with_context_annotations(
mut self,
enabled: bool,
annotations: std::collections::HashMap<u64, String>,
) -> Self {
self.dense_context = enabled;
self.context_by_id = annotations;
self
}
fn encode_into_store(
&mut self,
embedder: &SingleVectorEmbedder,
chunks: &[(u64, &str)],
) -> Result<()> {
for batch in chunks.chunks(ENCODE_BATCH) {
if let Err(e) = crate::memory::check_rss_or_abort("HNSW encode batch") {
anyhow::bail!("Indexing aborted: {e}");
}
for &(chunk_id, content) in batch {
let emb = if self.dense_context {
let ctx = self.context_by_id.get(&chunk_id).map_or("", String::as_str);
embedder.encode_document_with_context(ctx, content)?
} else {
embedder.encode_document(content)?
};
let (q, scale) = quantize_int8(&emb);
self.store.push(chunk_id, q, scale);
}
}
Ok(())
}
fn encode_streaming_ids<F>(
&mut self,
embedder: &SingleVectorEmbedder,
chunk_ids: &[u64],
mut fetch: F,
) -> Result<()>
where
F: FnMut(&[u64]) -> Result<Vec<(u64, String)>>,
{
for batch_ids in chunk_ids.chunks(ENCODE_BATCH) {
if let Err(e) = crate::memory::check_rss_or_abort("HNSW encode batch") {
anyhow::bail!("Indexing aborted: {e}");
}
let batch = fetch(batch_ids)?;
for (chunk_id, content) in &batch {
let emb = if self.dense_context {
let ctx = self.context_by_id.get(chunk_id).map_or("", String::as_str);
embedder.encode_document_with_context(ctx, content)?
} else {
embedder.encode_document(content)?
};
let (q, scale) = quantize_int8(&emb);
self.store.push(*chunk_id, q, scale);
}
}
Ok(())
}
pub fn build_streaming_ids<F>(&mut self, chunk_ids: &[u64], mut fetch: F) -> Result<()>
where
F: FnMut(&[u64]) -> Result<Vec<(u64, String)>>,
{
if chunk_ids.is_empty() {
tracing::info!("No chunks to encode for coderank-hnsw index");
return Ok(());
}
let embedder = self.make_embedder()?;
self.store = Int8VectorStore {
dim: EMBEDDING_DIM,
..Default::default()
};
self.encode_streaming_ids(&embedder, chunk_ids, &mut fetch)?;
crate::memory::purge_allocator();
let dir = self.dir.clone();
HnswDenseIndex::save(&self.store, &dir)?;
tracing::info!("coderank-hnsw index built ({} vectors)", self.store.len());
Ok(())
}
pub fn insert_streaming_ids<F>(&mut self, chunk_ids: &[u64], mut fetch: F) -> Result<()>
where
F: FnMut(&[u64]) -> Result<Vec<(u64, String)>>,
{
if chunk_ids.is_empty() {
return Ok(());
}
let embedder = self.make_embedder()?;
self.store = Int8VectorStore {
dim: EMBEDDING_DIM,
..Default::default()
};
self.encode_streaming_ids(&embedder, chunk_ids, &mut fetch)?;
let dir = self.dir.clone();
let idx = Int8VectorStore::load_index_only(&dir)?;
Int8VectorStore::append_split(
&dir,
idx,
&self.store.chunk_ids,
&self.store.scales,
&self.store.vectors,
)?;
Ok(())
}
fn make_embedder(&self) -> Result<SingleVectorEmbedder> {
let model_dir =
crate::embedding::single_vector_model::ensure_coderank_model(&self.models_dir)?;
SingleVectorEmbedder::for_indexing(&model_dir)
}
}
impl DenseIndexBuilder for CoderankHnswIndexBuilder {
fn name(&self) -> &'static str {
CoderankHnswBackend::NAME
}
fn build(&mut self, chunks: &[(u64, &str)]) -> Result<()> {
if chunks.is_empty() {
tracing::info!("No chunks to encode for coderank-hnsw index");
return Ok(());
}
let embedder = self.make_embedder()?;
self.store = Int8VectorStore {
dim: EMBEDDING_DIM,
..Default::default()
};
self.encode_into_store(&embedder, chunks)?;
crate::memory::purge_allocator();
let dir = self.dir.clone();
HnswDenseIndex::save(&self.store, &dir)?;
tracing::info!("coderank-hnsw index built ({} vectors)", self.store.len());
Ok(())
}
fn insert(&mut self, chunks: &[(u64, &str)]) -> Result<()> {
if chunks.is_empty() {
return Ok(());
}
let embedder = self.make_embedder()?;
self.store = Int8VectorStore {
dim: EMBEDDING_DIM,
..Default::default()
};
self.encode_into_store(&embedder, chunks)?;
let dir = self.dir.clone();
let idx = Int8VectorStore::load_index_only(&dir)?;
Int8VectorStore::append_split(
&dir,
idx,
&self.store.chunk_ids,
&self.store.scales,
&self.store.vectors,
)?;
Ok(())
}
fn delete(&mut self, chunk_ids: &[u64]) -> Result<()> {
if chunk_ids.is_empty() {
return Ok(());
}
let remove: std::collections::HashSet<u64> = chunk_ids.iter().copied().collect();
let dir = self.dir.clone();
let idx = Int8VectorStore::load_index_only(&dir)?;
Int8VectorStore::tombstone_split(&dir, idx, &remove)
}
fn persist(&self, _dir: &Path) -> Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn save_split_then_load_split_round_trips() {
let tmp = tempfile::TempDir::new().unwrap();
let mut store = Int8VectorStore {
dim: 3,
..Default::default()
};
let (q1, s1) = quantize_int8(&[0.6, 0.8, 0.0]);
let (q2, s2) = quantize_int8(&[0.0, 0.6, 0.8]);
store.push(10, q1.clone(), s1);
store.push(20, q2.clone(), s2);
store.save_split(tmp.path()).unwrap();
assert!(tmp.path().join(INDEX_FILE).exists());
assert!(tmp.path().join(VECTORS_FILE).exists());
let back = Int8VectorStore::load_split(tmp.path()).unwrap();
assert_eq!(back.chunk_ids, vec![10, 20]);
assert_eq!(back.dim, 3);
assert_eq!(back.vectors, vec![q1.clone(), q2.clone()]);
assert_eq!(back.scales, vec![s1, s2]);
assert_eq!(back.tombstoned, vec![false, false]);
}
#[test]
fn append_split_adds_rows_without_touching_existing_ones() {
let tmp = tempfile::TempDir::new().unwrap();
let mut store = Int8VectorStore {
dim: 2,
..Default::default()
};
let (q1, s1) = quantize_int8(&[1.0, 0.0]);
store.push(1, q1.clone(), s1);
store.save_split(tmp.path()).unwrap();
let idx = Int8VectorStore::load_index_only(tmp.path()).unwrap();
let (q2, s2) = quantize_int8(&[0.0, 1.0]);
Int8VectorStore::append_split(tmp.path(), idx, &[2], &[s2], std::slice::from_ref(&q2))
.unwrap();
let back = Int8VectorStore::load_split(tmp.path()).unwrap();
assert_eq!(back.chunk_ids, vec![1, 2]);
assert_eq!(back.vectors, vec![q1, q2]);
assert_eq!(back.tombstoned, vec![false, false]);
}
#[test]
fn tombstone_split_marks_flag_without_touching_vectors_file() {
let tmp = tempfile::TempDir::new().unwrap();
let mut store = Int8VectorStore {
dim: 2,
..Default::default()
};
let (q1, s1) = quantize_int8(&[1.0, 0.0]);
let (q2, s2) = quantize_int8(&[0.0, 1.0]);
store.push(1, q1, s1);
store.push(2, q2, s2);
store.save_split(tmp.path()).unwrap();
let vectors_bytes_before = std::fs::read(tmp.path().join(VECTORS_FILE)).unwrap();
let idx = Int8VectorStore::load_index_only(tmp.path()).unwrap();
let remove: std::collections::HashSet<u64> = [1].into_iter().collect();
Int8VectorStore::tombstone_split(tmp.path(), idx, &remove).unwrap();
let vectors_bytes_after = std::fs::read(tmp.path().join(VECTORS_FILE)).unwrap();
assert_eq!(
vectors_bytes_before, vectors_bytes_after,
"tombstoning must never rewrite the vectors payload file"
);
let back = Int8VectorStore::load_split(tmp.path()).unwrap();
assert!(back.is_tombstoned(0));
assert!(!back.is_tombstoned(1));
assert_eq!(back.live_count(), 1);
}
#[test]
fn load_split_on_empty_store_does_not_require_vectors_file() {
let tmp = tempfile::TempDir::new().unwrap();
let store = Int8VectorStore {
dim: 4,
..Default::default()
};
store.save_split(tmp.path()).unwrap();
std::fs::remove_file(tmp.path().join(VECTORS_FILE)).ok();
let back = Int8VectorStore::load_split(tmp.path()).unwrap();
assert_eq!(back.chunk_ids.len(), 0);
}
#[test]
fn builder_insert_appends_without_rewriting_existing_payload_bytes() {
let tmp = tempfile::TempDir::new().unwrap();
let mut store = Int8VectorStore {
dim: 2,
..Default::default()
};
let (q1, s1) = quantize_int8(&[1.0, 0.0]);
store.push(1, q1, s1);
store.save_split(tmp.path()).unwrap();
let existing_bytes = std::fs::read(tmp.path().join(VECTORS_FILE)).unwrap();
let idx = Int8VectorStore::load_index_only(tmp.path()).unwrap();
let (q2, s2) = quantize_int8(&[0.0, 1.0]);
Int8VectorStore::append_split(tmp.path(), idx, &[2], &[s2], std::slice::from_ref(&q2))
.unwrap();
let new_bytes = std::fs::read(tmp.path().join(VECTORS_FILE)).unwrap();
assert!(
new_bytes.starts_with(&existing_bytes),
"insert must append after existing bytes, never rewrite them"
);
assert_eq!(new_bytes.len(), existing_bytes.len() + 2 );
}
#[test]
fn dense_sentinel_file_is_index_bin_for_coderank_hnsw() {
use crate::search::dense_backend::{DenseBackendKind, dense_sentinel_file};
assert_eq!(
dense_sentinel_file(DenseBackendKind::CoderankHnsw),
"index.bin"
);
}
#[test]
fn preset_names_resolve() {
assert_eq!(HnswParams::preset("default").m, 16);
assert_eq!(HnswParams::preset("high_recall").ef_search, 200);
assert_eq!(HnswParams::preset("low_latency").ef_search, 32);
assert_eq!(HnswParams::preset("memory_optimized").m, 8);
assert_eq!(HnswParams::preset("garbage").m, 16); }
#[test]
fn resolve_keeps_preset_ef_search_when_override_is_zero() {
let p = HnswParams::resolve("high_recall", 0, 0);
assert_eq!(
p.ef_search, 200,
"preset ef_search must survive a 0 override"
);
assert_eq!(p.rescore_k, 0);
let p2 = HnswParams::resolve("high_recall", 96, 0);
assert_eq!(p2.ef_search, 96, "explicit ef_search override must win");
assert_eq!(HnswParams::resolve("default", 0, 0).ef_search, 64);
assert_eq!(HnswParams::resolve("default", 0, 256).rescore_k, 256);
}
#[test]
fn baked_ef_search_covers_explicit_rescore_k() {
let params = HnswParams {
ef_search: 64,
rescore_k: 300,
..HnswParams::default()
};
assert_eq!(HnswDenseIndex::baked_ef_search(params), 300);
let params2 = HnswParams {
ef_search: 200,
rescore_k: 50,
..HnswParams::default()
};
assert_eq!(HnswDenseIndex::baked_ef_search(params2), 200);
}
#[test]
fn brute_force_ranks_by_cosine_descending() {
let store = Int8VectorStore {
dim: 2,
scales: vec![1.0, 1.0, 1.0],
vectors: vec![
quantize_int8(&[1.0, 0.0]).0, quantize_int8(&[
std::f32::consts::FRAC_1_SQRT_2,
std::f32::consts::FRAC_1_SQRT_2,
])
.0, quantize_int8(&[0.0, 1.0]).0, ],
chunk_ids: vec![10, 20, 30],
tombstoned: vec![false, false, false],
};
let q = vec![1.0f32, 0.0];
let hits = store.brute_force(&q, 3, 0);
assert_eq!(
hits.iter().map(|h| h.chunk_id).collect::<Vec<_>>(),
vec![10, 20, 30]
);
assert!(hits[0].score >= hits[1].score && hits[1].score >= hits[2].score);
}
#[test]
fn empty_store_returns_no_hits() {
let store = Int8VectorStore {
dim: 4,
scales: vec![],
vectors: vec![],
chunk_ids: vec![],
tombstoned: vec![],
};
assert!(store.brute_force(&[0.0, 0.0, 0.0, 0.0], 5, 0).is_empty());
}
#[test]
fn rescore_reorders_int8_candidates_by_fp32() {
let store = Int8VectorStore {
dim: 3,
scales: vec![1.0, 1.0],
vectors: vec![
quantize_int8(&[0.9, 0.1, 0.0]).0,
quantize_int8(&[0.8, 0.2, 0.0]).0,
],
chunk_ids: vec![1, 2],
tombstoned: vec![false, false],
};
let q = vec![1.0f32, 0.0, 0.0];
let hits = store.fp32_rescore(&q, &[0, 1], 2);
assert_eq!(hits[0].chunk_id, 1);
}
#[test]
fn brute_force_wide_shortlist_recovers_int8_misranked_topk() {
let mut v1 = [0.952f32, 0.218, 0.505];
let mut v2 = [0.605f32, 0.394, 0.096];
l2_normalize(&mut v1);
l2_normalize(&mut v2);
let (q1, s1) = quantize_int8(&v1);
let (q2, s2) = quantize_int8(&v2);
let store = Int8VectorStore {
dim: 3,
scales: vec![s1, s2],
vectors: vec![q1, q2],
chunk_ids: vec![1, 2],
tombstoned: vec![false, false],
};
let mut query = vec![1.0f32, 0.028, 0.0];
l2_normalize(&mut query);
let narrow = store.brute_force(&query, 1, 1);
assert_eq!(
narrow[0].chunk_id, 2,
"with shortlist == k the int8 inversion is unrecoverable (returns the WRONG id) \
— this is exactly the bug FIX 4 addresses"
);
let wide = store.brute_force(&query, 1, 4);
assert_eq!(
wide[0].chunk_id, 1,
"wide int8 shortlist + fp32 rescore must recover the true fp32 top-1"
);
}
#[test]
fn builder_name_and_empty_build_is_ok() {
let tmp = tempfile::TempDir::new().unwrap();
let mut b = CoderankHnswIndexBuilder::new(tmp.path(), HnswParams::default());
assert_eq!(DenseIndexBuilder::name(&b), "coderank-hnsw");
b.build(&[]).unwrap();
assert!(!tmp.path().join(INDEX_FILE).exists());
assert!(!tmp.path().join(VECTORS_FILE).exists());
}
#[test]
fn store_persist_and_reload_round_trips() {
let tmp = tempfile::TempDir::new().unwrap();
let mut store = Int8VectorStore {
dim: 2,
..Default::default()
};
let (q, s) = quantize_int8(&[0.6, 0.8]);
store.push(42, q, s);
let path = tmp.path().join("vectors.bin");
std::fs::write(&path, postcard::to_stdvec(&store).unwrap()).unwrap();
let back: Int8VectorStore = postcard::from_bytes(&std::fs::read(&path).unwrap()).unwrap();
assert_eq!(back.chunk_ids, vec![42]);
assert_eq!(back.dim, 2);
}
#[test]
fn backend_positional_chunk_ids_is_none() {
let tmp = tempfile::TempDir::new().unwrap();
let store = Int8VectorStore {
dim: 4,
..Default::default()
};
store.save_split(tmp.path()).unwrap();
let backend = CoderankHnswBackend::open_for_test(tmp.path());
assert!(backend.positional_chunk_ids().is_none());
}
#[test]
fn vectors_for_dequantizes_stored_ids_and_skips_missing() {
let mut store = Int8VectorStore {
dim: 3,
..Default::default()
};
let (q10, s10) = quantize_int8(&[0.6, 0.8, 0.0]);
let (q20, s20) = quantize_int8(&[0.0, 0.6, 0.8]);
store.push(10, q10, s10);
store.push(20, q20, s20);
let got = store.vectors_for(&[20, 99, 10]);
assert_eq!(
got.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
vec![20, 10]
);
for (id, v) in &got {
assert_eq!(v.len(), store.dim);
let row = store.chunk_ids.iter().position(|c| c == id).unwrap();
let expected = store.dequant_row(row);
assert_eq!(
v, &expected,
"vectors_for must return the re-normalized stored int8 row"
);
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
assert!(
(norm - 1.0).abs() < 1e-5,
"embed_doc_vectors row must be unit-norm (FIX 5), got {norm}"
);
}
}
#[test]
#[ignore = "needs the CodeRankEmbed model download"]
fn vector_accessor_methods_return_some_on_built_index() {
let tmp = tempfile::TempDir::new().unwrap();
let dir = tmp.path().join("dense");
let models = tmp.path().join("models");
let model_dir = crate::embedding::single_vector_model::ensure_coderank_model(&models)
.expect("model provisions");
let mut b = CoderankHnswIndexBuilder::new(&dir, HnswParams::default())
.with_models_dir(models.clone());
b.build(&[
(7, "fn parse_int(s:&str)->i64 { s.parse().unwrap() }"),
(9, "def add(a,b): return a+b"),
])
.unwrap();
let backend = CoderankHnswBackend::open(&dir, &model_dir, HnswParams::default()).unwrap();
let qv = backend
.embed_text_vector("parse an integer")
.expect("Some on coderank-hnsw");
assert_eq!(qv.len(), SingleVectorEmbedder::embedding_dim());
let norm: f32 = qv.iter().map(|x| x * x).sum::<f32>().sqrt();
assert!(
(norm - 1.0).abs() < 1e-3,
"query vector must be L2-normalized, got {norm}"
);
let dv = backend
.embed_doc_vectors(&[9, 7, 404])
.expect("Some on coderank-hnsw");
assert_eq!(dv.iter().map(|(id, _)| *id).collect::<Vec<_>>(), vec![9, 7]);
assert!(
dv.iter()
.all(|(_, v)| v.len() == SingleVectorEmbedder::embedding_dim())
);
}
#[test]
fn hnsw_graph_and_brute_force_agree_on_tiny_corpus() {
unsafe {
std::env::set_var("SEMANTEX_HNSW_MIN_VECTORS", "4");
}
let mut store = Int8VectorStore {
dim: 4,
..Default::default()
};
for i in 0..12u64 {
let mut v = vec![
((i * 7 % 11) as f32) - 5.0,
((i * 13 % 11) as f32) - 5.0,
((i * 17 % 11) as f32) - 5.0,
((i * 19 % 11) as f32) - 5.0,
];
l2_normalize(&mut v);
let (q, s) = quantize_int8(&v);
store.push(100 + i, q, s);
}
let params = HnswParams::default();
let graph = HnswDenseIndex::build_graph(&store, params);
assert!(
graph.is_some(),
"graph must build above the lowered threshold"
);
let idx = HnswDenseIndex {
store: store.clone(),
params,
graph,
baked_ef_search: HnswDenseIndex::baked_ef_search(params),
};
let mut query = vec![1.0f32, 0.5, -0.3, 0.2];
l2_normalize(&mut query);
let hnsw_hits: Vec<u64> = idx.search(&query, 5).iter().map(|h| h.chunk_id).collect();
let bf_hits: Vec<u64> = store
.brute_force(&query, 5, 4 * 5)
.iter()
.map(|h| h.chunk_id)
.collect();
assert_eq!(
hnsw_hits, bf_hits,
"HNSW (with fp32 rescore) must agree with brute-force on a tiny corpus"
);
unsafe {
std::env::remove_var("SEMANTEX_HNSW_MIN_VECTORS");
}
}
#[test]
fn tombstoned_row_is_excluded_from_brute_force() {
let mut store = Int8VectorStore {
dim: 2,
..Default::default()
};
let (q1, s1) = quantize_int8(&[1.0, 0.0]);
let (q2, s2) = quantize_int8(&[0.9, 0.1]);
store.push(1, q1, s1);
store.push(2, q2, s2);
assert_eq!(store.live_count(), 2);
store.tombstoned[0] = true;
assert_eq!(store.live_count(), 1);
assert!(store.is_tombstoned(0));
assert!(!store.is_tombstoned(1));
let hits = store.brute_force(&[1.0, 0.0], 5, 5);
assert_eq!(
hits.iter().map(|h| h.chunk_id).collect::<Vec<_>>(),
vec![2],
"tombstoned row (chunk_id=1) must never be returned"
);
}
#[test]
fn tombstoned_row_is_excluded_from_vectors_for() {
let mut store = Int8VectorStore {
dim: 2,
..Default::default()
};
let (q1, s1) = quantize_int8(&[1.0, 0.0]);
let (q2, s2) = quantize_int8(&[0.0, 1.0]);
store.push(10, q1, s1);
store.push(20, q2, s2);
store.tombstoned[0] = true;
let got = store.vectors_for(&[10, 20]);
assert_eq!(
got.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
vec![20],
"tombstoned chunk_id=10 must be skipped even when explicitly requested"
);
}
#[test]
fn build_graph_skips_tombstoned_rows_and_uses_live_count_for_threshold() {
unsafe {
std::env::set_var("SEMANTEX_HNSW_MIN_VECTORS", "4");
}
let mut store = Int8VectorStore {
dim: 4,
..Default::default()
};
for i in 0..6u64 {
let mut v = vec![
((i * 7 % 11) as f32) - 5.0,
((i * 13 % 11) as f32) - 5.0,
((i * 17 % 11) as f32) - 5.0,
((i * 19 % 11) as f32) - 5.0,
];
l2_normalize(&mut v);
let (q, s) = quantize_int8(&v);
store.push(100 + i, q, s);
}
store.tombstoned[0] = true;
store.tombstoned[1] = true;
store.tombstoned[2] = true;
assert_eq!(store.live_count(), 3);
let graph = HnswDenseIndex::build_graph(&store, HnswParams::default());
assert!(
graph.is_none(),
"live_count (3) below the lowered threshold (4) must skip the graph, \
even though raw row count (6) is above it"
);
unsafe {
std::env::remove_var("SEMANTEX_HNSW_MIN_VECTORS");
}
}
}