use super::Collection;
use crate::{
CollectionSchema, DataType, Doc, Durability, FieldSchema, FieldValue, HnswQueryParams,
IndexParams, MetricType, SearchQuery,
};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use tempfile::tempdir;
fn path_of(directory: &tempfile::TempDir, name: &str) -> String {
directory
.path()
.join(name)
.to_str()
.expect("temporary path must be UTF-8")
.to_string()
}
fn vector_schema(name: &str, metric: MetricType, hnsw: bool) -> CollectionSchema {
let mut embedding =
FieldSchema::new("embedding", DataType::VectorFp32, false, 4).expect("vector field");
let params = if hnsw {
IndexParams::hnsw(metric, 8, 64).expect("hnsw params")
} else {
IndexParams::flat(metric).expect("flat params")
};
embedding
.set_index_params(¶ms)
.expect("index params must attach");
CollectionSchema::builder(name)
.add_field(embedding)
.build()
.expect("schema")
}
fn vector_doc(id: &str, vector: &[f32]) -> Doc {
let mut doc = Doc::with_pk(id).expect("primary key");
doc.add_vector_f32("embedding", vector)
.expect("vector value");
doc
}
fn ranked(docs: &[Doc]) -> Vec<(String, u32)> {
docs.iter()
.map(|doc| {
(
doc.get_pk()
.expect("query hit must have a primary key")
.to_string(),
doc.get_score().to_bits(),
)
})
.collect()
}
#[test]
fn enterprise_ga_defaults_stay_exact() {
assert_eq!(Durability::default(), Durability::Always);
let ivf = IndexParams::ivf(MetricType::L2, 4, 1, false).expect("ivf params");
assert!(
ivf.params.get("scale_factor").is_none(),
"plain IVF must not gain a default scale_factor"
);
let query = SearchQuery::new("embedding", &[1.0, 0.0, 0.0, 0.0], 3).expect("query");
assert_ne!(
query
.params
.get("is_using_refiner")
.and_then(serde_json::Value::as_bool),
Some(false)
);
}
#[test]
fn enterprise_ga_public_scores_match_exact_oracle_and_radius() {
let temporary = tempdir().expect("temporary directory");
let docs = [
vector_doc("doc-same", &[1.0, 0.0, 0.0, 0.0]),
vector_doc("doc-mid", &[1.0, 1.0, 0.0, 0.0]),
vector_doc("doc-far", &[0.0, 1.0, 0.0, 0.0]),
vector_doc("doc-back", &[-1.0, 0.0, 0.0, 0.0]),
];
let query_vector = [1.0_f32, 0.0, 0.0, 0.0];
let mut flat_hits = None;
let mut hnsw_hits = None;
for (name, hnsw, slot) in [
("flat-oracle", false, &mut flat_hits),
("hnsw-oracle", true, &mut hnsw_hits),
] {
let collection = Collection::create(
&path_of(&temporary, name),
&vector_schema(name, MetricType::Cosine, hnsw),
None,
)
.expect("collection");
let refs: Vec<&Doc> = docs.iter().collect();
collection.insert(&refs).expect("insert");
let query = SearchQuery::new("embedding", &query_vector, 4).expect("query");
*slot = Some(collection.query(&query).expect("query must return"));
}
let flat_hits = flat_hits.expect("flat hits");
let hnsw_hits = hnsw_hits.expect("hnsw hits");
assert_eq!(ranked(&hnsw_hits), ranked(&flat_hits));
assert_eq!(
flat_hits
.iter()
.filter_map(|doc| doc.get_pk())
.collect::<Vec<_>>(),
vec!["doc-same", "doc-mid", "doc-far", "doc-back"]
);
let collection = Collection::open(&path_of(&temporary, "hnsw-oracle"), None).expect("reopen");
let mut excluding = SearchQuery::new("embedding", &query_vector, 4).expect("query");
excluding.set_radius(2.0).expect("radius");
assert!(collection
.query(&excluding)
.expect("high radius query")
.is_empty());
let mut open_radius = SearchQuery::new("embedding", &query_vector, 4).expect("query");
open_radius.set_radius(-1.0).expect("radius");
assert_eq!(
ranked(&collection.query(&open_radius).expect("negative radius")),
ranked(&hnsw_hits)
);
}
#[test]
fn zero_cosine_query_keeps_topk_after_flat_rebuild() {
let temporary = tempdir().expect("temporary directory");
let collection = Collection::create(
&path_of(&temporary, "zero-cosine"),
&vector_schema("zero-cosine", MetricType::Cosine, false),
None,
)
.expect("collection");
for doc in [
vector_doc("c", &[0.0, 0.0, 1.0, 0.0]),
vector_doc("a", &[1.0, 0.0, 0.0, 0.0]),
vector_doc("b", &[0.0, 1.0, 0.0, 0.0]),
] {
collection.insert(&[&doc]).expect("insert");
}
let query = SearchQuery::new("embedding", &[0.0, 0.0, 0.0, 0.0], 2).expect("query");
let before = collection.query(&query).expect("query before rebuild");
collection.rebuild_index("embedding").expect("flat rebuild");
let after = collection.query(&query).expect("query after rebuild");
let expected = vec![
("a".to_string(), 0.0_f32.to_bits()),
("b".to_string(), 0.0_f32.to_bits()),
];
assert_eq!(ranked(&before), expected);
assert_eq!(ranked(&after), expected);
let nonzero = SearchQuery::new("embedding", &[1.0, 0.0, 0.0, 0.0], 1).expect("query");
let nonzero_hits = collection.query(&nonzero).expect("nonzero query");
assert_eq!(nonzero_hits.first().and_then(Doc::get_pk), Some("a"));
assert_eq!(nonzero_hits.first().map(Doc::get_score), Some(1.0));
}
#[test]
fn enterprise_ga_default_hnsw_ef_matches_explicit_64() {
let temporary = tempdir().expect("temporary directory");
let mut embedding =
FieldSchema::new("embedding", DataType::VectorFp32, false, 4).expect("vector field");
embedding
.set_index_params(&IndexParams::hnsw(MetricType::L2, 8, 64).expect("hnsw"))
.expect("attach");
let schema = CollectionSchema::builder("ef-64")
.add_field(embedding)
.build()
.expect("schema");
let collection =
Collection::create(&path_of(&temporary, "ef"), &schema, None).expect("collection");
let docs: Vec<Doc> = (0..96)
.map(|index| {
let id = format!("doc-{index:03}");
let base = f32::from(u16::try_from(index).expect("index fits"));
vector_doc(&id, &[base, base * 0.5, base.sin(), (base + 1.0).cos()])
})
.collect();
let refs: Vec<&Doc> = docs.iter().collect();
collection.insert(&refs).expect("insert");
let query_vector = [3.0_f32, 1.5, 0.1, 0.2];
let implicit = SearchQuery::new("embedding", &query_vector, 5).expect("query");
let mut explicit = SearchQuery::new("embedding", &query_vector, 5).expect("query");
explicit
.set_hnsw_params(HnswQueryParams::new(64, 0.0, false, true))
.expect("ef 64");
assert_eq!(
ranked(&collection.query(&implicit).expect("default ef")),
ranked(&collection.query(&explicit).expect("explicit ef"))
);
}
#[test]
fn enterprise_ga_always_sync_lets_prior_revision_query_return() {
let temporary = tempdir().expect("temporary directory");
let collection = Collection::create(
&path_of(&temporary, "sync"),
&vector_schema("sync", MetricType::Cosine, false),
None,
)
.expect("collection");
let published = vector_doc("published", &[1.0, 0.0, 0.0, 0.0]);
collection.insert(&[&published]).expect("first insert");
let gate = collection.test_arm_wal_sync_stall();
let acked = Arc::new(AtomicBool::new(false));
let pending = vector_doc("pending", &[0.0, 1.0, 0.0, 0.0]);
let writer = collection.clone();
let acked_flag = Arc::clone(&acked);
let insert_thread = thread::spawn(move || {
let result = writer.insert(&[&pending]);
acked_flag.store(true, Ordering::Release);
result
});
assert!(
gate.wait_entered(Duration::from_secs(5)),
"Always sync did not reach the durability stall"
);
assert!(
!acked.load(Ordering::Acquire),
"commit was acknowledged before the durability sync finished"
);
let reader = collection.clone();
let (sender, receiver) = mpsc::channel();
thread::spawn(move || {
let query = SearchQuery::new("embedding", &[1.0, 0.0, 0.0, 0.0], 4).expect("query");
let _ = sender.send(reader.query(&query));
});
let query_result = receiver.recv_timeout(Duration::from_secs(2));
gate.release();
let hits = query_result
.expect("published revision query blocked during durability sync")
.expect("query");
let ids: Vec<_> = hits.iter().filter_map(|doc| doc.get_pk()).collect();
assert_eq!(ids, vec!["published"]);
assert!(collection.fetch(&["pending"]).expect("fetch").is_empty());
insert_thread
.join()
.expect("insert thread")
.expect("stalled insert must succeed after sync");
assert!(acked.load(Ordering::Acquire));
assert_eq!(collection.fetch(&["pending"]).expect("fetch").len(), 1);
let path = path_of(&temporary, "sync");
drop(collection);
let reopened = Collection::open(&path, None).expect("reopen");
assert_eq!(reopened.count().expect("count"), 2);
assert_eq!(reopened.fetch(&["published"]).expect("fetch").len(), 1);
assert_eq!(reopened.fetch(&["pending"]).expect("fetch").len(), 1);
}
#[test]
fn enterprise_ga_diskann_sidecar_failure_keeps_index_cache() {
let temporary = tempdir().expect("temporary directory");
let mut tag = FieldSchema::new("tag", DataType::String, false, 0).expect("tag field");
tag.set_index_params(&IndexParams::invert(false, false).expect("invert"))
.expect("invert attaches");
let mut embedding =
FieldSchema::new("embedding", DataType::VectorFp32, false, 8).expect("vector field");
embedding
.set_index_params(&IndexParams::diskann(MetricType::L2, 8, 16, 1).expect("diskann"))
.expect("diskann attaches");
let schema = CollectionSchema::builder("sidecar")
.add_field(tag)
.add_field(embedding)
.build()
.expect("schema");
let path = path_of(&temporary, "sidecar");
let collection = Collection::create(&path, &schema, None).expect("collection");
let mut docs = Vec::new();
for index in 0..8 {
let id = if index == 0 {
"alpha-doc".to_string()
} else {
format!("doc-{index}")
};
let mut doc = Doc::with_pk(&id).expect("primary key");
let tag_value = if index == 0 { "alpha" } else { "beta" };
doc.add_string("tag", tag_value).expect("tag");
let vector: Vec<f32> = (0..8)
.map(|coordinate| {
f32::from(u16::try_from(index + coordinate).expect("coordinate fits"))
})
.collect();
doc.add_vector_f32("embedding", &vector).expect("vector");
docs.push(doc);
}
let refs: Vec<&Doc> = docs.iter().collect();
collection.insert(&refs).expect("insert");
collection.test_arm_diskann_write_fault();
collection.flush().expect("flush");
assert!(
collection.test_diskann_write_fault_fired(),
"DiskANN sidecar write was not attempted"
);
drop(collection);
let reopened = Collection::open(&path, None).expect("reopen");
assert!(
reopened.stats().expect("stats").index_cache_hit,
"sidecar failure rebuilt the index registry instead of restoring the cache"
);
let mut query =
SearchQuery::new("embedding", &[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], 4).expect("query");
query.set_filter("tag == \"alpha\"").expect("filter");
let hits = reopened.query(&query).expect("filtered query");
assert_eq!(
hits.iter()
.filter_map(|doc| doc.get_pk())
.collect::<Vec<_>>(),
vec!["alpha-doc"]
);
assert_eq!(
reopened
.fetch(&["alpha-doc"])
.expect("fetch")
.first()
.and_then(|doc| doc.field("tag")),
Some(&FieldValue::String("alpha".to_string()))
);
}
#[test]
fn enterprise_ga_stale_diskann_sidecar_keeps_index_cache() {
let temporary = tempdir().expect("temporary directory");
let mut tag = FieldSchema::new("tag", DataType::String, false, 0).expect("tag field");
tag.set_index_params(&IndexParams::invert(false, false).expect("invert"))
.expect("invert attaches");
let mut embedding =
FieldSchema::new("embedding", DataType::VectorFp32, false, 8).expect("vector field");
embedding
.set_index_params(&IndexParams::diskann(MetricType::L2, 8, 16, 1).expect("diskann"))
.expect("diskann attaches");
let schema = CollectionSchema::builder("stale-sidecar")
.add_field(tag)
.add_field(embedding)
.build()
.expect("schema");
let path = path_of(&temporary, "stale-sidecar");
let collection = Collection::create(&path, &schema, None).expect("collection");
let mut docs = Vec::new();
for index in 0..8 {
let id = format!("doc-{index}");
let mut doc = Doc::with_pk(&id).expect("primary key");
doc.add_string("tag", if index == 0 { "alpha" } else { "beta" })
.expect("tag");
let vector: Vec<f32> = (0..8)
.map(|coordinate| {
f32::from(u16::try_from(index + coordinate).expect("coordinate fits"))
})
.collect();
doc.add_vector_f32("embedding", &vector).expect("vector");
docs.push(doc);
}
let refs: Vec<&Doc> = docs.iter().collect();
collection.insert(&refs).expect("insert");
collection.flush().expect("first flush");
let sidecar = std::path::Path::new(&path)
.join("indexes")
.join("diskann-graph.bin");
let stale_bytes = std::fs::read(&sidecar).expect("first sidecar");
assert!(!stale_bytes.is_empty());
let mut extra = Doc::with_pk("doc-extra").expect("primary key");
extra.add_string("tag", "gamma").expect("tag");
extra
.add_vector_f32("embedding", &[8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0])
.expect("vector");
collection.upsert(&[&extra]).expect("revision advances");
collection.test_arm_diskann_write_fault();
collection.flush().expect("second flush");
assert!(collection.test_diskann_write_fault_fired());
assert_eq!(
std::fs::read(&sidecar).expect("stale sidecar remains"),
stale_bytes
);
drop(collection);
let reopened = Collection::open(&path, None).expect("reopen");
assert!(
reopened.stats().expect("stats").index_cache_hit,
"a stale DiskANN sidecar must not discard the restored index cache"
);
let mut query =
SearchQuery::new("embedding", &[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], 4).expect("query");
query.set_filter("tag == \"alpha\"").expect("filter");
let hits = reopened.query(&query).expect("filtered query");
assert_eq!(
hits.iter()
.filter_map(|doc| doc.get_pk())
.collect::<Vec<_>>(),
vec!["doc-0"]
);
assert_eq!(reopened.fetch(&["doc-extra"]).expect("fetch").len(), 1);
}
#[test]
fn enterprise_ga_flush_encodes_the_published_generation_without_body_clones() {
const DOCUMENTS: usize = 2_000;
let temporary = tempdir().expect("temporary directory");
let mut embedding =
FieldSchema::new("embedding", DataType::VectorFp32, false, 4).expect("vector field");
embedding
.set_index_params(&IndexParams::flat(MetricType::Cosine).expect("flat"))
.expect("flat attaches");
let schema = CollectionSchema::builder("flush-generation")
.add_field(embedding)
.build()
.expect("schema");
let path = path_of(&temporary, "flush-generation");
let collection = Collection::create(&path, &schema, None).expect("create");
let mut stored = Vec::with_capacity(DOCUMENTS);
for index in 0..DOCUMENTS {
let mut doc = Doc::with_pk(format!("doc-{index:04}")).expect("primary key");
let index = u16::try_from(index).expect("index fits");
doc.add_vector_f32("embedding", &[f32::from(index % 97), 0.5, 1.5, -0.25])
.expect("vector");
stored.push(doc);
}
let first: Vec<&Doc> = stored[..1_000].iter().collect();
let second: Vec<&Doc> = stored[1_000..].iter().collect();
collection.insert(&first).expect("first batch");
collection.insert(&second).expect("second batch");
crate::doc::reset_doc_body_clones();
collection.flush().expect("flush");
assert_eq!(
crate::doc::doc_body_clones(),
0,
"flush must encode the published generation without cloning every document body"
);
drop(collection);
let full = snapshot_files(std::path::Path::new(&path));
assert_eq!(
full.len(),
1,
"the first content checkpoint is one full file"
);
assert_eq!(full[0].2.first().copied(), Some(0x95));
assert_eq!(full[0].2.get(1).copied(), Some(0x04));
let full_length = full[0].1;
let opened = Collection::open(&path, None).expect("format 4 snapshot must open");
assert_eq!(opened.count().expect("count"), DOCUMENTS);
for doc in &stored {
let id = doc.get_pk().expect("primary key");
let fetched = opened.fetch(&[id]).expect("fetch");
assert_eq!(fetched.len(), 1);
assert_eq!(fetched[0].get_pk(), Some(id));
assert_eq!(fetched[0].vector("embedding"), doc.vector("embedding"));
}
let mut edited = Doc::with_pk(stored[0].get_pk().expect("primary key")).expect("edited key");
edited
.add_vector_f32("embedding", &[9.0, 8.0, 7.0, 6.0])
.expect("edited vector");
opened.upsert(&[&edited]).expect("one document change");
opened.flush().expect("delta checkpoint");
drop(opened);
let files = snapshot_files(std::path::Path::new(&path));
assert!(
files.len() >= 2,
"the base snapshot remains beside the delta"
);
let delta = files.last().expect("delta snapshot");
assert_eq!(delta.2.first().copied(), Some(0x98));
assert_eq!(delta.2.get(1).copied(), Some(0x05));
assert!(
delta.1.saturating_mul(2) < full_length,
"one-document checkpoint length {} is not under half of the full snapshot {full_length}",
delta.1
);
let reopened = Collection::open(&path, None).expect("delta snapshot must open");
let edited_id = edited.get_pk().expect("edited id");
let fetched = reopened.fetch(&[edited_id]).expect("fetch edited");
assert_eq!(fetched[0].vector("embedding"), edited.vector("embedding"));
assert_eq!(reopened.count().expect("count"), DOCUMENTS);
}
fn snapshot_files(root: &std::path::Path) -> Vec<(u64, u64, Vec<u8>)> {
let mut files = Vec::new();
for entry in std::fs::read_dir(root.join("segments")).expect("segments directory") {
let entry = entry.expect("segment entry");
let name = entry.file_name();
let Some(name) = name.to_str() else {
continue;
};
let Some(generation) = name
.strip_prefix("snapshot-")
.and_then(|rest| rest.strip_suffix(".bin"))
.and_then(|rest| rest.parse::<u64>().ok())
else {
continue;
};
let bytes = std::fs::read(entry.path()).expect("snapshot bytes");
let length = u64::try_from(bytes.len()).expect("snapshot length fits");
files.push((generation, length, bytes));
}
files.sort_by_key(|(generation, _, _)| *generation);
files
}