use selene_core::{CancellationChecker, GraphId, LabelSet, NodeId, Value, VectorMetric, db_string};
use super::super::{ApproximateVectorSearchOptions, VectorCandidateSet, VectorSearchError};
use super::{props, vector};
use crate::VectorIndexKind;
use crate::shared::SharedGraph;
#[test]
fn turbo_quant_search_reranks_primary_vectors_without_shadow_storage() {
let shared = SharedGraph::new(GraphId::new(973));
let doc = db_string("vector.ann.turbo.doc").unwrap();
let embedding = db_string("embedding").unwrap();
{
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
for components in [
[1.0, 0.0, 0.0],
[0.9, 0.1, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 1.0],
] {
mutator
.create_node(
LabelSet::single(doc.clone()),
props(&embedding, Value::Vector(vector(&components))),
)
.unwrap();
}
txn.commit().unwrap();
}
shared
.create_vector_index(
doc.clone(),
embedding.clone(),
VectorIndexKind::TurboQuantCosine,
3,
)
.unwrap();
let usage = shared
.read()
.vector_index_for(&doc, &embedding)
.unwrap()
.memory_usage();
assert_eq!(usage.turbo_quant_entries, 4);
assert_eq!(usage.turbo_quant_referenced_vector_bytes, 0);
assert_eq!(usage.estimated_reachable_bytes, usage.estimated_index_bytes);
let query = vector(&[1.0, 0.0, 0.0]);
let exact = shared
.exact_vector_search_nodes(&doc, &embedding, &query, VectorMetric::Cosine, 2)
.unwrap();
let approx = shared
.approximate_vector_search_nodes_checked(
&doc,
&embedding,
&query,
ApproximateVectorSearchOptions::new(VectorMetric::Cosine, 2, 4),
CancellationChecker::disabled(),
)
.unwrap();
assert_eq!(approx, exact);
}
#[test]
fn turbo_quant_batch_search_matches_single_queries() {
let shared = SharedGraph::new(GraphId::new(974));
let doc = db_string("vector.ann.turbo.batch.doc").unwrap();
let embedding = db_string("embedding").unwrap();
{
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
for value in 0..32 {
let angle = value as f32 * 0.05;
mutator
.create_node(
LabelSet::single(doc.clone()),
props(
&embedding,
Value::Vector(vector(&[angle.cos(), angle.sin()])),
),
)
.unwrap();
}
txn.commit().unwrap();
}
shared
.create_vector_index(
doc.clone(),
embedding.clone(),
VectorIndexKind::TurboQuantCosine,
2,
)
.unwrap();
let queries = vec![
vector(&[1.0, 0.0]),
vector(&[0.8, 0.6]),
vector(&[0.25, 1.0]),
];
let options = ApproximateVectorSearchOptions::new(VectorMetric::Cosine, 4, 16);
let batched = shared
.approximate_vector_search_nodes_batch_checked(
&doc,
&embedding,
&queries,
options,
CancellationChecker::disabled(),
)
.unwrap();
let singles: Vec<_> = queries
.iter()
.map(|query| {
shared
.approximate_vector_search_nodes_checked(
&doc,
&embedding,
query,
options,
CancellationChecker::disabled(),
)
.unwrap()
})
.collect();
assert_eq!(batched, singles);
let covered_options = ApproximateVectorSearchOptions::new(VectorMetric::Cosine, 4, 32);
let covered = shared
.approximate_vector_search_nodes_batch_checked(
&doc,
&embedding,
&queries,
covered_options,
CancellationChecker::disabled(),
)
.unwrap();
let exact = shared
.exact_vector_search_nodes_batch_checked(
&doc,
&embedding,
&queries,
VectorMetric::Cosine,
4,
CancellationChecker::disabled(),
)
.unwrap();
assert_eq!(covered, exact);
}
#[test]
fn turbo_quant_candidate_set_search_filters_index_rows_and_reranks() {
let shared = SharedGraph::new(GraphId::new(975));
let doc = db_string("vector.ann.turbo.candidates.doc").unwrap();
let other = db_string("vector.ann.turbo.candidates.other").unwrap();
let embedding = db_string("embedding").unwrap();
let (doc_ids, wrong_label) = {
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
let mut doc_ids = Vec::new();
for components in [
[1.0, 0.0, 0.0],
[0.95, 0.1, 0.0],
[0.0, 1.0, 0.0],
[0.8, 0.2, 0.0],
[-1.0, 0.0, 0.0],
] {
doc_ids.push(
mutator
.create_node(
LabelSet::single(doc.clone()),
props(&embedding, Value::Vector(vector(&components))),
)
.unwrap(),
);
}
let wrong_label = mutator
.create_node(
LabelSet::single(other),
props(&embedding, Value::Vector(vector(&[1.0, 0.0, 0.0]))),
)
.unwrap();
txn.commit().unwrap();
(doc_ids, wrong_label)
};
shared
.create_vector_index(
doc.clone(),
embedding.clone(),
VectorIndexKind::TurboQuantCosine,
3,
)
.unwrap();
let candidate_set = VectorCandidateSet::from_nodes([
doc_ids[4],
doc_ids[3],
doc_ids[1],
wrong_label,
NodeId::new(99_999),
]);
let hits = shared
.approximate_vector_search_candidate_set_checked(
&doc,
&embedding,
&vector(&[1.0, 0.0, 0.0]),
&candidate_set,
ApproximateVectorSearchOptions::new(VectorMetric::Cosine, 2, 8),
CancellationChecker::disabled(),
)
.unwrap();
assert_eq!(
hits.iter().map(|hit| hit.node_id).collect::<Vec<_>>(),
vec![doc_ids[1], doc_ids[3]]
);
assert!(!hits.iter().any(|hit| hit.node_id == wrong_label));
}
#[test]
fn turbo_quant_candidate_set_batch_search_matches_single_queries() {
let shared = SharedGraph::new(GraphId::new(976));
let doc = db_string("vector.ann.turbo.candidate.batch.doc").unwrap();
let embedding = db_string("embedding").unwrap();
let doc_ids = {
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
let mut doc_ids = Vec::new();
for value in 0..32 {
let angle = value as f32 * 0.05;
doc_ids.push(
mutator
.create_node(
LabelSet::single(doc.clone()),
props(
&embedding,
Value::Vector(vector(&[angle.cos(), angle.sin()])),
),
)
.unwrap(),
);
}
txn.commit().unwrap();
doc_ids
};
shared
.create_vector_index(
doc.clone(),
embedding.clone(),
VectorIndexKind::TurboQuantCosine,
2,
)
.unwrap();
let queries = vec![
vector(&[1.0, 0.0]),
vector(&[0.8, 0.6]),
vector(&[0.25, 1.0]),
];
let candidate_sets = vec![
VectorCandidateSet::from_nodes(doc_ids[0..16].iter().copied()),
VectorCandidateSet::from_nodes(doc_ids[8..24].iter().copied()),
VectorCandidateSet::from_nodes(doc_ids[16..32].iter().copied()),
];
let options = ApproximateVectorSearchOptions::new(VectorMetric::Cosine, 4, 16);
let batch = shared
.approximate_vector_search_candidate_sets_batch_checked(
&doc,
&embedding,
&queries,
&candidate_sets,
options,
CancellationChecker::disabled(),
)
.unwrap();
let singles = queries
.iter()
.zip(&candidate_sets)
.map(|(query, candidates)| {
shared
.approximate_vector_search_candidate_set_checked(
&doc,
&embedding,
query,
candidates,
options,
CancellationChecker::disabled(),
)
.unwrap()
})
.collect::<Vec<_>>();
assert_eq!(batch, singles);
let exact = shared
.score_vector_candidate_sets_batch_checked(
&embedding,
&queries,
&candidate_sets,
VectorMetric::Cosine,
4,
CancellationChecker::disabled(),
)
.unwrap();
assert_eq!(batch, exact);
}
#[test]
fn turbo_quant_candidate_set_batch_rejects_shape_mismatch() {
let shared = SharedGraph::new(GraphId::new(977));
let doc = db_string("vector.ann.turbo.candidate.batch.err.doc").unwrap();
let embedding = db_string("embedding").unwrap();
let error = shared
.approximate_vector_search_candidate_sets_batch_checked(
&doc,
&embedding,
&[vector(&[1.0, 0.0])],
&[],
ApproximateVectorSearchOptions::new(VectorMetric::Cosine, 4, 16),
CancellationChecker::disabled(),
)
.unwrap_err();
assert!(matches!(
error,
VectorSearchError::BatchLengthMismatch {
queries: 1,
candidate_sets: 0
}
));
}