use selene_core::{
CancellationChecker, CancellationToken, CoreError, GraphId, LabelSet, NodeId, PropertyMap,
Value, VectorMetric, VectorValue, db_string,
};
use crate::{
GraphError, SharedGraph, VectorCandidateSet, VectorNeighborDirection,
VectorNeighborSearchOptions, VectorSearchError,
};
#[path = "score_tests/candidate_batch.rs"]
mod candidate_batch;
#[path = "score_tests/candidate_set.rs"]
mod candidate_set;
fn vector(components: &[f32]) -> VectorValue {
VectorValue::new(components.to_vec()).expect("test vector is valid")
}
fn props(key: &selene_core::DbString, value: Value) -> PropertyMap {
PropertyMap::from_pairs([(key.clone(), value)]).expect("test property map is valid")
}
#[test]
fn score_vector_nodes_ranks_unique_live_vector_candidates() {
let shared = SharedGraph::new(GraphId::new(974));
let label = db_string("vector.score.doc").unwrap();
let embedding = db_string("embedding").unwrap();
let other = db_string("other").unwrap();
let ids = {
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
let mut ids = Vec::new();
for value in 0..5 {
ids.push(
mutator
.create_node(
LabelSet::single(label.clone()),
props(&embedding, Value::Vector(vector(&[value as f32, 0.0]))),
)
.unwrap(),
);
}
ids.push(
mutator
.create_node(
LabelSet::single(label.clone()),
props(&other, Value::String(db_string("not-a-vector").unwrap())),
)
.unwrap(),
);
txn.commit().unwrap();
ids
};
{
let mut txn = shared.begin_write();
txn.mutator().delete_node(ids[4]).unwrap();
txn.commit().unwrap();
}
let hits = shared
.score_vector_nodes_checked(
&embedding,
&vector(&[2.1, 0.0]),
&[
ids[3],
ids[2],
ids[2],
ids[0],
ids[4],
ids[5],
NodeId::new(999),
],
VectorMetric::SquaredEuclidean,
3,
CancellationChecker::disabled(),
)
.unwrap();
assert_eq!(hits.len(), 3);
assert_eq!(hits[0].node_id, ids[2]);
assert_eq!(hits[1].node_id, ids[3]);
assert_eq!(hits[2].node_id, ids[0]);
}
#[test]
fn score_vector_nodes_zero_k_does_not_bind_query() {
let shared = SharedGraph::new(GraphId::new(975));
let embedding = db_string("embedding").unwrap();
let hits = shared
.score_vector_nodes_checked(
&embedding,
&vector(&[0.0, 0.0]),
&[NodeId::new(1)],
VectorMetric::Cosine,
0,
CancellationChecker::disabled(),
)
.unwrap();
assert!(hits.is_empty());
}
#[test]
fn score_vector_nodes_surfaces_candidate_dimension_errors() {
let shared = SharedGraph::new(GraphId::new(976));
let label = db_string("vector.score.dim.doc").unwrap();
let embedding = db_string("embedding").unwrap();
let node = {
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
let node = mutator
.create_node(
LabelSet::single(label),
props(&embedding, Value::Vector(vector(&[1.0, 2.0, 3.0]))),
)
.unwrap();
txn.commit().unwrap();
node
};
let err = shared
.score_vector_nodes_checked(
&embedding,
&vector(&[1.0, 2.0]),
&[node],
VectorMetric::SquaredEuclidean,
10,
CancellationChecker::disabled(),
)
.expect_err("dimension mismatch must error");
assert!(matches!(
err,
VectorSearchError::Graph(GraphError::Core(CoreError::VectorDimensionMismatch {
lhs: 2,
rhs: 3
}))
));
}
#[test]
fn score_vector_nodes_batch_matches_single_queries() {
let shared = SharedGraph::new(GraphId::new(977));
let label = db_string("vector.score.batch.doc").unwrap();
let embedding = db_string("embedding").unwrap();
let ids = {
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
let mut ids = Vec::new();
for value in 0..8 {
ids.push(
mutator
.create_node(
LabelSet::single(label.clone()),
props(&embedding, Value::Vector(vector(&[value as f32, 0.0]))),
)
.unwrap(),
);
}
txn.commit().unwrap();
ids
};
let queries = vec![vector(&[2.2, 0.0]), vector(&[5.1, 0.0])];
let candidate_sets = vec![
vec![ids[4], ids[2], ids[2], ids[0], NodeId::new(999)],
vec![ids[7], ids[5], ids[1], ids[5]],
];
let batched = shared
.score_vector_nodes_batch_checked(
&embedding,
&queries,
&candidate_sets,
VectorMetric::SquaredEuclidean,
2,
CancellationChecker::disabled(),
)
.unwrap();
let singles: Vec<_> = queries
.iter()
.zip(&candidate_sets)
.map(|(query, candidates)| {
shared
.score_vector_nodes_checked(
&embedding,
query,
candidates,
VectorMetric::SquaredEuclidean,
2,
CancellationChecker::disabled(),
)
.unwrap()
})
.collect();
assert_eq!(batched, singles);
assert_eq!(batched[0][0].node_id, ids[2]);
assert_eq!(batched[1][0].node_id, ids[5]);
}
#[test]
fn score_vector_nodes_batch_empty_and_zero_k_do_not_bind_queries() {
let shared = SharedGraph::new(GraphId::new(978));
let embedding = db_string("embedding").unwrap();
let empty = shared
.score_vector_nodes_batch_checked::<Vec<NodeId>>(
&embedding,
&[],
&[],
VectorMetric::Cosine,
10,
CancellationChecker::disabled(),
)
.unwrap();
assert!(empty.is_empty());
let hits = shared
.score_vector_nodes_batch_checked(
&embedding,
&[vector(&[0.0, 0.0]), vector(&[0.0, 0.0])],
&[vec![NodeId::new(1)], vec![NodeId::new(2)]],
VectorMetric::Cosine,
0,
CancellationChecker::disabled(),
)
.unwrap();
assert_eq!(hits, vec![Vec::new(), Vec::new()]);
}
#[test]
fn score_vector_nodes_batch_rejects_invalid_batch_shape() {
let shared = SharedGraph::new(GraphId::new(979));
let embedding = db_string("embedding").unwrap();
let err = shared
.score_vector_nodes_batch_checked(
&embedding,
&[vector(&[0.0, 0.0]), vector(&[1.0, 0.0])],
&[vec![NodeId::new(1)]],
VectorMetric::SquaredEuclidean,
1,
CancellationChecker::disabled(),
)
.expect_err("query/candidate batch arity mismatch must error");
assert!(matches!(
err,
VectorSearchError::BatchLengthMismatch {
queries: 2,
candidate_sets: 1
}
));
let err = shared
.score_vector_nodes_batch_checked(
&embedding,
&[vector(&[0.0, 0.0]), vector(&[1.0, 0.0, 0.0])],
&[vec![NodeId::new(1)], vec![NodeId::new(2)]],
VectorMetric::SquaredEuclidean,
1,
CancellationChecker::disabled(),
)
.expect_err("mixed query dimensions must error");
assert!(matches!(
err,
VectorSearchError::Graph(GraphError::Core(CoreError::VectorDimensionMismatch {
lhs: 2,
rhs: 3
}))
));
}
#[test]
fn score_vector_neighbors_filters_direction_and_ranks_unique_live_vectors() {
let shared = SharedGraph::new(GraphId::new(980));
let anchor_label = db_string("vector.neighbor.anchor").unwrap();
let doc_label = db_string("vector.neighbor.doc").unwrap();
let embedding = db_string("embedding").unwrap();
let link = db_string("DEPENDS_ON").unwrap();
let other_link = db_string("MENTIONS").unwrap();
let other = db_string("other").unwrap();
let (anchor, out_near, out_far, incoming, deleted, non_vector) = {
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
let anchor = mutator
.create_node(LabelSet::single(anchor_label), PropertyMap::new())
.unwrap();
let out_near = mutator
.create_node(
LabelSet::single(doc_label.clone()),
props(&embedding, Value::Vector(vector(&[2.0, 0.0]))),
)
.unwrap();
let out_far = mutator
.create_node(
LabelSet::single(doc_label.clone()),
props(&embedding, Value::Vector(vector(&[8.0, 0.0]))),
)
.unwrap();
let incoming = mutator
.create_node(
LabelSet::single(doc_label.clone()),
props(&embedding, Value::Vector(vector(&[1.0, 0.0]))),
)
.unwrap();
let deleted = mutator
.create_node(
LabelSet::single(doc_label.clone()),
props(&embedding, Value::Vector(vector(&[0.0, 0.0]))),
)
.unwrap();
let non_vector = mutator
.create_node(
LabelSet::single(doc_label),
props(&other, Value::String(db_string("skip").unwrap())),
)
.unwrap();
mutator
.create_edge(link.clone(), anchor, out_far, PropertyMap::new())
.unwrap();
mutator
.create_edge(link.clone(), anchor, out_near, PropertyMap::new())
.unwrap();
mutator
.create_edge(link.clone(), anchor, out_near, PropertyMap::new())
.unwrap();
mutator
.create_edge(link.clone(), anchor, deleted, PropertyMap::new())
.unwrap();
mutator
.create_edge(link.clone(), anchor, non_vector, PropertyMap::new())
.unwrap();
mutator
.create_edge(other_link.clone(), anchor, incoming, PropertyMap::new())
.unwrap();
mutator
.create_edge(link.clone(), incoming, anchor, PropertyMap::new())
.unwrap();
txn.commit().unwrap();
(anchor, out_near, out_far, incoming, deleted, non_vector)
};
{
let mut txn = shared.begin_write();
txn.mutator().delete_node(deleted).unwrap();
txn.commit().unwrap();
}
let outgoing = shared
.score_vector_neighbors_checked(
&embedding,
&vector(&[2.2, 0.0]),
anchor,
VectorNeighborSearchOptions::new(
&link,
VectorNeighborDirection::Outgoing,
VectorMetric::SquaredEuclidean,
10,
),
CancellationChecker::disabled(),
)
.unwrap();
assert_eq!(
outgoing.iter().map(|hit| hit.node_id).collect::<Vec<_>>(),
vec![out_near, out_far]
);
let incoming_hits = shared
.score_vector_neighbors_checked(
&embedding,
&vector(&[2.2, 0.0]),
anchor,
VectorNeighborSearchOptions::new(
&link,
VectorNeighborDirection::Incoming,
VectorMetric::SquaredEuclidean,
10,
),
CancellationChecker::disabled(),
)
.unwrap();
assert_eq!(
incoming_hits
.iter()
.map(|hit| hit.node_id)
.collect::<Vec<_>>(),
vec![incoming]
);
let both = shared
.score_vector_neighbors_checked(
&embedding,
&vector(&[2.2, 0.0]),
anchor,
VectorNeighborSearchOptions::new(
&link,
VectorNeighborDirection::Both,
VectorMetric::SquaredEuclidean,
10,
),
CancellationChecker::disabled(),
)
.unwrap();
assert_eq!(
both.iter().map(|hit| hit.node_id).collect::<Vec<_>>(),
vec![out_near, incoming, out_far]
);
let missing_anchor = shared
.score_vector_neighbors_checked(
&embedding,
&vector(&[2.2, 0.0]),
non_vector,
VectorNeighborSearchOptions::new(
&other_link,
VectorNeighborDirection::Outgoing,
VectorMetric::SquaredEuclidean,
10,
),
CancellationChecker::disabled(),
)
.unwrap();
assert!(missing_anchor.is_empty());
}
#[test]
fn expand_vector_candidate_set_preserves_roots_and_walks_labeled_edges() {
let shared = SharedGraph::new(GraphId::new(982));
let root_label = db_string("vector.expand.root").unwrap();
let doc_label = db_string("vector.expand.doc").unwrap();
let support = db_string("SUPPORTS").unwrap();
let other = db_string("MENTIONS").unwrap();
let (root_a, root_b, out_a, out_b, incoming, wrong_label) = {
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
let root_a = mutator
.create_node(LabelSet::single(root_label.clone()), PropertyMap::new())
.unwrap();
let root_b = mutator
.create_node(LabelSet::single(root_label), PropertyMap::new())
.unwrap();
let out_a = mutator
.create_node(LabelSet::single(doc_label.clone()), PropertyMap::new())
.unwrap();
let out_b = mutator
.create_node(LabelSet::single(doc_label.clone()), PropertyMap::new())
.unwrap();
let incoming = mutator
.create_node(LabelSet::single(doc_label.clone()), PropertyMap::new())
.unwrap();
let wrong_label = mutator
.create_node(LabelSet::single(doc_label), PropertyMap::new())
.unwrap();
mutator
.create_edge(support.clone(), root_a, out_b, PropertyMap::new())
.unwrap();
mutator
.create_edge(support.clone(), root_a, out_b, PropertyMap::new())
.unwrap();
mutator
.create_edge(support.clone(), root_b, out_a, PropertyMap::new())
.unwrap();
mutator
.create_edge(support.clone(), incoming, root_a, PropertyMap::new())
.unwrap();
mutator
.create_edge(other.clone(), root_a, wrong_label, PropertyMap::new())
.unwrap();
txn.commit().unwrap();
(root_a, root_b, out_a, out_b, incoming, wrong_label)
};
let roots = VectorCandidateSet::from_nodes([root_b, root_a, root_a]);
let outgoing =
shared.expand_vector_candidate_set(&roots, &support, VectorNeighborDirection::Outgoing);
assert_eq!(outgoing.as_nodes(), &[root_a, root_b, out_a, out_b]);
let incoming_set =
shared.expand_vector_candidate_set(&roots, &support, VectorNeighborDirection::Incoming);
assert_eq!(incoming_set.as_nodes(), &[root_a, root_b, incoming]);
let both = shared.expand_vector_candidate_set(&roots, &support, VectorNeighborDirection::Both);
assert_eq!(both.as_nodes(), &[root_a, root_b, out_a, out_b, incoming]);
let missing =
shared.expand_vector_candidate_set(&roots, &other, VectorNeighborDirection::Incoming);
assert_eq!(missing.as_nodes(), &[root_a, root_b]);
assert!(!both.as_nodes().contains(&wrong_label));
}
#[test]
fn expand_vector_candidate_set_handles_empty_and_cancelled_inputs() {
let shared = SharedGraph::new(GraphId::new(983));
let link = db_string("EXPANDS_TO").unwrap();
let empty = VectorCandidateSet::default();
assert!(
shared
.expand_vector_candidate_set(&empty, &link, VectorNeighborDirection::Both)
.is_empty()
);
let token = CancellationToken::new();
token.cancel();
let err = shared
.expand_vector_candidate_set_checked(
&empty,
&link,
VectorNeighborDirection::Both,
CancellationChecker::new(Some(&token), None),
)
.expect_err("cancelled expansion must report cancellation");
assert!(matches!(err, VectorSearchError::Cancelled));
}
#[test]
fn score_vector_neighbors_batch_matches_single_anchor_queries() {
let shared = SharedGraph::new(GraphId::new(981));
let anchor_label = db_string("vector.neighbor.batch.anchor").unwrap();
let doc_label = db_string("vector.neighbor.batch.doc").unwrap();
let embedding = db_string("embedding").unwrap();
let link = db_string("NEAR").unwrap();
let (anchor_a, anchor_b, ids) = {
let mut txn = shared.begin_write();
let mut mutator = txn.mutator();
let anchor_a = mutator
.create_node(LabelSet::single(anchor_label.clone()), PropertyMap::new())
.unwrap();
let anchor_b = mutator
.create_node(LabelSet::single(anchor_label), PropertyMap::new())
.unwrap();
let mut ids = Vec::new();
for value in 0..8 {
ids.push(
mutator
.create_node(
LabelSet::single(doc_label.clone()),
props(&embedding, Value::Vector(vector(&[value as f32, 0.0]))),
)
.unwrap(),
);
}
for &node in &[ids[1], ids[2], ids[5]] {
mutator
.create_edge(link.clone(), anchor_a, node, PropertyMap::new())
.unwrap();
}
for &node in &[ids[4], ids[6], ids[7]] {
mutator
.create_edge(link.clone(), anchor_b, node, PropertyMap::new())
.unwrap();
}
txn.commit().unwrap();
(anchor_a, anchor_b, ids)
};
let queries = vec![vector(&[2.1, 0.0]), vector(&[6.2, 0.0])];
let anchors = vec![anchor_a, anchor_b];
let batched = shared
.score_vector_neighbors_batch_checked(
&embedding,
&queries,
&anchors,
VectorNeighborSearchOptions::new(
&link,
VectorNeighborDirection::Outgoing,
VectorMetric::SquaredEuclidean,
2,
),
CancellationChecker::disabled(),
)
.unwrap();
let singles: Vec<_> = queries
.iter()
.zip(&anchors)
.map(|(query, anchor)| {
shared
.score_vector_neighbors_checked(
&embedding,
query,
*anchor,
VectorNeighborSearchOptions::new(
&link,
VectorNeighborDirection::Outgoing,
VectorMetric::SquaredEuclidean,
2,
),
CancellationChecker::disabled(),
)
.unwrap()
})
.collect();
assert_eq!(batched, singles);
assert_eq!(
batched[0].iter().map(|hit| hit.node_id).collect::<Vec<_>>(),
vec![ids[2], ids[1]]
);
assert_eq!(
batched[1].iter().map(|hit| hit.node_id).collect::<Vec<_>>(),
vec![ids[6], ids[7]]
);
let repeated_anchor_queries = vec![vector(&[1.2, 0.0]), vector(&[5.1, 0.0])];
let repeated_anchor_anchors = vec![anchor_a, anchor_a];
let repeated_anchor_batched = shared
.score_vector_neighbors_batch_checked(
&embedding,
&repeated_anchor_queries,
&repeated_anchor_anchors,
VectorNeighborSearchOptions::new(
&link,
VectorNeighborDirection::Outgoing,
VectorMetric::SquaredEuclidean,
2,
),
CancellationChecker::disabled(),
)
.unwrap();
let repeated_anchor_singles: Vec<_> = repeated_anchor_queries
.iter()
.map(|query| {
shared
.score_vector_neighbors_checked(
&embedding,
query,
anchor_a,
VectorNeighborSearchOptions::new(
&link,
VectorNeighborDirection::Outgoing,
VectorMetric::SquaredEuclidean,
2,
),
CancellationChecker::disabled(),
)
.unwrap()
})
.collect();
assert_eq!(repeated_anchor_batched, repeated_anchor_singles);
let err = shared
.score_vector_neighbors_batch_checked(
&embedding,
&queries,
&[anchor_a],
VectorNeighborSearchOptions::new(
&link,
VectorNeighborDirection::Outgoing,
VectorMetric::SquaredEuclidean,
2,
),
CancellationChecker::disabled(),
)
.expect_err("query/anchor count mismatch must error");
assert!(matches!(
err,
VectorSearchError::BatchLengthMismatch {
queries: 2,
candidate_sets: 1
}
));
}