use std::fs;
use selene_core::{
Change, HnswIndexConfig, IvfIndexConfig, LabelSet, NodeId, PropertyValueType, SchemaChange,
SchemaVectorIndexKind, Value, VectorValue, db_string,
};
use super::*;
use crate::{VectorIndexConfig, VectorIndexKind};
fn vector_value() -> Value {
Value::Vector(VectorValue::new(vec![0.25, 0.5, 0.75]).unwrap())
}
#[test]
fn recover_snapshot_preserves_vector_property() {
let dir = temp_dir("snapshot-vector");
let shared = SharedGraph::builder(GraphId::new(37)).build().unwrap();
let mut txn = shared.begin_write();
{
let mut mutator = txn.mutator();
mutator
.create_node(
LabelSet::single(db_string("recover.vector.node").unwrap()),
prop("recover.vector", vector_value()),
)
.unwrap();
}
txn.commit().unwrap();
write_snapshot(&dir, &shared, 1);
let recovered = SharedGraph::recover(&dir, GraphId::new(37)).unwrap();
let snapshot = recovered.read();
assert_eq!(snapshot.node_count(), 1);
expect_prop(
snapshot.node_properties(NodeId::new(1)).unwrap(),
"recover.vector",
&vector_value(),
);
let _ = fs::remove_dir_all(dir);
}
#[test]
fn recover_snapshot_preserves_vector_index_registration() {
let dir = temp_dir("snapshot-vector-index");
let label = db_string("recover.vector.index.node").unwrap();
let property = db_string("recover.vector.index.embedding").unwrap();
let shared = SharedGraph::builder(GraphId::new(40)).build().unwrap();
{
let mut txn = shared.begin_write();
txn.mutator()
.create_node(
LabelSet::single(label.clone()),
prop(
"recover.vector.index.embedding",
Value::Vector(VectorValue::new(vec![1.0, 2.0, 3.0]).unwrap()),
),
)
.unwrap();
txn.commit().unwrap();
}
shared
.create_vector_index(label.clone(), property.clone(), VectorIndexKind::Flat, 3)
.unwrap();
write_snapshot(&dir, &shared, 1);
let recovered = SharedGraph::recover(&dir, GraphId::new(40)).unwrap();
let snapshot = recovered.read();
let index = snapshot.vector_index_for(&label, &property).unwrap();
assert_eq!(index.kind(), VectorIndexKind::Flat);
assert_eq!(index.dimension(), 3);
assert_eq!(index.rows().iter().collect::<Vec<_>>(), vec![0]);
let _ = fs::remove_dir_all(dir);
}
#[test]
fn recover_snapshot_preserves_hnsw_vector_index_registration() {
let dir = temp_dir("snapshot-hnsw-vector-index");
let label = db_string("recover.hnsw.vector.index.node").unwrap();
let property = db_string("recover.hnsw.vector.index.embedding").unwrap();
let config = HnswIndexConfig::new(24, 128);
let shared = SharedGraph::builder(GraphId::new(42)).build().unwrap();
{
let mut txn = shared.begin_write();
txn.mutator()
.create_node(
LabelSet::single(label.clone()),
prop(
"recover.hnsw.vector.index.embedding",
Value::Vector(VectorValue::new(vec![1.0, 0.0, 0.0]).unwrap()),
),
)
.unwrap();
txn.commit().unwrap();
}
shared
.create_vector_index_named_with_config(
label.clone(),
property.clone(),
VectorIndexKind::HnswCosine,
3,
None,
Some(config),
)
.unwrap();
write_snapshot(&dir, &shared, 1);
let recovered = SharedGraph::recover(&dir, GraphId::new(42)).unwrap();
let snapshot = recovered.read();
let index = snapshot.vector_index_for(&label, &property).unwrap();
assert_eq!(index.kind(), VectorIndexKind::HnswCosine);
assert_eq!(index.dimension(), 3);
assert_eq!(index.hnsw_config(), Some(config));
assert_eq!(index.rows().iter().collect::<Vec<_>>(), vec![0]);
let _ = fs::remove_dir_all(dir);
}
#[test]
fn recover_snapshot_preserves_ivf_vector_index_registration() {
let dir = temp_dir("snapshot-ivf-vector-index");
let label = db_string("recover.ivf.vector.index.node").unwrap();
let property = db_string("recover.ivf.vector.index.embedding").unwrap();
let config = IvfIndexConfig::new(4);
let shared = SharedGraph::builder(GraphId::new(43)).build().unwrap();
{
let mut txn = shared.begin_write();
txn.mutator()
.create_node(
LabelSet::single(label.clone()),
prop(
"recover.ivf.vector.index.embedding",
Value::Vector(VectorValue::new(vec![1.0, 0.0, 0.0]).unwrap()),
),
)
.unwrap();
txn.commit().unwrap();
}
shared
.create_vector_index_named_with_configs(
label.clone(),
property.clone(),
VectorIndexKind::IvfCosine,
3,
None,
VectorIndexConfig::ivf(config),
)
.unwrap();
write_snapshot(&dir, &shared, 1);
let recovered = SharedGraph::recover(&dir, GraphId::new(43)).unwrap();
let snapshot = recovered.read();
let index = snapshot.vector_index_for(&label, &property).unwrap();
assert_eq!(index.kind(), VectorIndexKind::IvfCosine);
assert_eq!(index.dimension(), 3);
assert_eq!(index.hnsw_config(), None);
assert_eq!(index.ivf_config(), Some(config));
assert_eq!(index.rows().iter().collect::<Vec<_>>(), vec![0]);
assert_eq!(index.memory_usage().ivf_centroids, 1);
assert_eq!(index.memory_usage().ivf_live_entries, 1);
let _ = fs::remove_dir_all(dir);
}
#[test]
fn recover_snapshot_preserves_turbo_quant_vector_index_registration() {
let dir = temp_dir("snapshot-turbo-quant-vector-index");
let label = db_string("recover.turbo.vector.index.node").unwrap();
let property = db_string("recover.turbo.vector.index.embedding").unwrap();
let shared = SharedGraph::builder(GraphId::new(44)).build().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]] {
mutator
.create_node(
LabelSet::single(label.clone()),
prop(
"recover.turbo.vector.index.embedding",
Value::Vector(VectorValue::new(components.to_vec()).unwrap()),
),
)
.unwrap();
}
txn.commit().unwrap();
}
shared
.create_vector_index(
label.clone(),
property.clone(),
VectorIndexKind::TurboQuantCosine,
3,
)
.unwrap();
write_snapshot(&dir, &shared, 1);
let recovered = SharedGraph::recover(&dir, GraphId::new(44)).unwrap();
let snapshot = recovered.read();
let index = snapshot.vector_index_for(&label, &property).unwrap();
assert_eq!(index.kind(), VectorIndexKind::TurboQuantCosine);
assert_eq!(index.dimension(), 3);
assert_eq!(index.rows().iter().collect::<Vec<_>>(), vec![0, 1, 2]);
assert_eq!(index.memory_usage().turbo_quant_entries, 3);
assert_eq!(index.memory_usage().turbo_quant_live_entries, 3);
assert!(index.memory_usage().turbo_quant_code_bytes > 0);
let _ = fs::remove_dir_all(dir);
}
#[test]
fn recover_wal_only_replays_vector_property() {
let dir = temp_dir("wal-vector");
append_wal(
&dir,
0,
&[Change::NodeCreated {
id: NodeId::new(1),
labels: LabelSet::single(db_string("recover.wal.vector.node").unwrap()),
properties: prop("recover.wal.vector", vector_value()),
}],
);
let recovered = SharedGraph::recover(&dir, GraphId::new(38)).unwrap();
let snapshot = recovered.read();
assert_eq!(snapshot.node_count(), 1);
expect_prop(
snapshot.node_properties(NodeId::new(1)).unwrap(),
"recover.wal.vector",
&vector_value(),
);
let _ = fs::remove_dir_all(dir);
}
#[test]
fn recover_wal_only_replays_vector_index_registration() {
let dir = temp_dir("wal-vector-index");
let label = db_string("recover.wal.vector.index.node").unwrap();
let property = db_string("recover.wal.vector.index.embedding").unwrap();
let config = HnswIndexConfig::new(24, 128);
append_wal(
&dir,
0,
&[
Change::NodeCreated {
id: NodeId::new(1),
labels: LabelSet::single(label.clone()),
properties: prop(
"recover.wal.vector.index.embedding",
Value::Vector(VectorValue::new(vec![0.25, 0.5, 0.75]).unwrap()),
),
},
Change::SchemaChanged {
graph: GraphId::new(41),
change: SchemaChange::VectorIndexCreated {
label: label.clone(),
property: property.clone(),
kind: SchemaVectorIndexKind::HnswSquaredEuclidean,
dimension: 3,
name: None,
hnsw_config: Some(config),
ivf_config: None,
},
},
],
);
let recovered = SharedGraph::recover(&dir, GraphId::new(41)).unwrap();
let snapshot = recovered.read();
let index = snapshot.vector_index_for(&label, &property).unwrap();
assert_eq!(index.kind(), VectorIndexKind::HnswSquaredEuclidean);
assert_eq!(index.dimension(), 3);
assert_eq!(index.hnsw_config(), Some(config));
assert_eq!(index.rows().iter().collect::<Vec<_>>(), vec![0]);
let _ = fs::remove_dir_all(dir);
}
#[test]
fn recover_wal_only_replays_ivf_vector_index_registration() {
let dir = temp_dir("wal-ivf-vector-index");
let label = db_string("recover.wal.ivf.vector.index.node").unwrap();
let property = db_string("recover.wal.ivf.vector.index.embedding").unwrap();
let config = IvfIndexConfig::new(4);
append_wal(
&dir,
0,
&[
Change::NodeCreated {
id: NodeId::new(1),
labels: LabelSet::single(label.clone()),
properties: prop(
"recover.wal.ivf.vector.index.embedding",
Value::Vector(VectorValue::new(vec![0.25, 0.5, 0.75]).unwrap()),
),
},
Change::SchemaChanged {
graph: GraphId::new(44),
change: SchemaChange::VectorIndexCreated {
label: label.clone(),
property: property.clone(),
kind: SchemaVectorIndexKind::IvfSquaredEuclidean,
dimension: 3,
name: None,
hnsw_config: None,
ivf_config: Some(config),
},
},
],
);
let recovered = SharedGraph::recover(&dir, GraphId::new(44)).unwrap();
let snapshot = recovered.read();
let index = snapshot.vector_index_for(&label, &property).unwrap();
assert_eq!(index.kind(), VectorIndexKind::IvfSquaredEuclidean);
assert_eq!(index.dimension(), 3);
assert_eq!(index.hnsw_config(), None);
assert_eq!(index.ivf_config(), Some(config));
assert_eq!(index.rows().iter().collect::<Vec<_>>(), vec![0]);
assert_eq!(index.memory_usage().ivf_live_entries, 1);
let _ = fs::remove_dir_all(dir);
}
#[test]
fn recover_closed_wal_only_preserves_vector_property_type() {
let dir = temp_dir("closed-schema-vector-wal-only");
let graph_id = GraphId::new(39);
let base = empty_closed_graph_type();
let shared = SharedGraph::builder(graph_id)
.bound_to(base.clone())
.unwrap()
.build()
.unwrap();
let sensor = db_string("VectorSensor").unwrap();
let embedding = db_string("embedding").unwrap();
let changes = {
let mut txn = shared.begin_write();
txn.mutator()
.create_node_type(
sensor.clone(),
LabelSet::single(sensor),
vec![PropertyTypeDef {
name: embedding.clone(),
value_type: PropertyValueType::Vector,
list_element_type: None,
required: true,
default: None,
immutable: false,
unique: false,
decimal_type: None,
character_string_type: None,
byte_string_type: None,
record_field_types: None,
}],
ValidationMode::Strict,
)
.unwrap();
txn.commit().unwrap().changes
};
append_wal(&dir, 0, &changes);
let recovered = SharedGraph::recover_closed(&dir, graph_id, base).unwrap();
let graph_type = recovered.graph_type().unwrap();
let property = &graph_type.node_types[0].properties[0];
assert_eq!(property.name, embedding);
assert_eq!(property.value_type, PropertyValueType::Vector);
assert_eq!(property.list_element_type, None);
assert!(property.required);
let _ = fs::remove_dir_all(dir);
}