weavatrix-search-vector 0.3.1

Persistent, mutable, bounded vector candidate search for Rust and Weavatrix
Documentation
use std::time::{SystemTime, UNIX_EPOCH};
use weavatrix_search_vector::{
    IndexBundle, IndexConfig, Metadata, MetadataFilter, MetadataValue, MutableVectorIndex,
    MutationOutcome, QuantizationKind, QuantizedIndex, VectorRecord,
};

fn config(dimensions: usize) -> IndexConfig {
    let mut config = IndexConfig::new(dimensions);
    config.connectivity = 2;
    config.expansion_build = 8;
    config.expansion_query = 8;
    config
}

fn metadata(group: &str, rank: u64) -> Metadata {
    let mut metadata = Metadata::new();
    metadata.insert("group", MetadataValue::Text(group.to_owned()));
    metadata.insert("rank", MetadataValue::U64(rank));
    metadata
}

fn bundle_path() -> std::path::PathBuf {
    let nonce = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("system clock should follow Unix epoch")
        .as_nanos();
    std::env::temp_dir().join(format!(
        "weavatrix-vector-bundle-test-{}-{nonce}.wvxb",
        std::process::id()
    ))
}

#[test]
fn mutable_overlay_supports_upsert_delete_filters_and_compaction() {
    let records = vec![
        VectorRecord {
            key: 10,
            vector: vec![1.0, 0.0],
            metadata: metadata("red", 1),
        },
        VectorRecord {
            key: 20,
            vector: vec![0.0, 1.0],
            metadata: metadata("blue", 2),
        },
    ];
    let index = MutableVectorIndex::build(config(2), &records).expect("mutable index should build");
    index
        .insert(30, &[0.9, 0.1], metadata("red", 3))
        .expect("new vector should insert");
    assert_eq!(
        index
            .upsert(20, &[0.8, 0.2], metadata("blue", 4))
            .expect("existing vector should update"),
        MutationOutcome::Updated
    );
    assert!(index.delete(10));
    assert!(!index.delete(10));
    assert_eq!(index.len(), 2);
    assert_eq!(
        index
            .search(&[1.0, 0.0], 2)
            .expect("delta query should succeed")
            .iter()
            .map(|hit| hit.key)
            .collect::<Vec<_>>(),
        vec![30, 20]
    );

    let blue = MetadataFilter::Equal {
        field: "group".to_owned(),
        value: MetadataValue::Text("blue".to_owned()),
    };
    assert_eq!(
        index
            .search_filtered(&[1.0, 0.0], 2, &blue)
            .expect("filtered query should succeed")
            .iter()
            .map(|hit| hit.key)
            .collect::<Vec<_>>(),
        vec![20]
    );
    let ranked = MetadataFilter::U64Range {
        field: "rank".to_owned(),
        minimum: Some(3),
        maximum: Some(4),
    };
    assert_eq!(
        index
            .search_filtered(&[1.0, 0.0], 2, &ranked)
            .expect("range query should succeed")
            .len(),
        2
    );

    index.compact().expect("delta should compact");
    assert_eq!(index.delta_len(), 0);
    assert_eq!(
        index
            .search(&[1.0, 0.0], 2)
            .expect("compacted query should succeed")
            .iter()
            .map(|hit| hit.key)
            .collect::<Vec<_>>(),
        vec![30, 20]
    );
}

#[test]
fn mutable_batches_rename_and_sealed_delta_avoid_base_compaction() {
    let records = vec![
        VectorRecord {
            key: 10,
            vector: vec![1.0, 0.0],
            metadata: metadata("base", 1),
        },
        VectorRecord {
            key: 20,
            vector: vec![0.0, 1.0],
            metadata: metadata("base", 2),
        },
    ];
    let index = MutableVectorIndex::build(config(2), &records).expect("mutable index should build");
    let batch = vec![
        VectorRecord {
            key: 30,
            vector: vec![0.9, 0.1],
            metadata: metadata("delta", 3),
        },
        VectorRecord {
            key: 40,
            vector: vec![0.8, 0.2],
            metadata: metadata("delta", 4),
        },
    ];
    index
        .insert_batch(&batch)
        .expect("batch should insert atomically");
    assert_eq!(index.staged_len(), 2);
    index.seal_delta().expect("delta should seal");
    assert_eq!(index.staged_len(), 0);
    assert_eq!(index.sealed_len(), 2);

    assert!(index.rename(30, 300).expect("rename should succeed"));
    assert_eq!(
        index
            .metadata(300)
            .and_then(|value| value.get("group").cloned()),
        Some(MetadataValue::Text("delta".to_owned()))
    );
    let outcomes = index
        .upsert_batch(&[
            VectorRecord::new(20, vec![0.7, 0.3]),
            VectorRecord::new(50, vec![0.6, 0.4]),
        ])
        .expect("upsert batch should apply");
    assert_eq!(
        outcomes,
        vec![MutationOutcome::Updated, MutationOutcome::Inserted]
    );
    assert_eq!(index.delete_batch(&[10, 40, 999]), 2);
    assert_eq!(index.len(), 3);
    assert_eq!(
        index
            .search(&[1.0, 0.0], 3)
            .expect("sealed query should succeed")
            .iter()
            .map(|hit| hit.key)
            .collect::<Vec<_>>(),
        vec![300, 20, 50]
    );

    index.seal_delta().expect("new writes should reseal");
    assert_eq!(index.staged_len(), 0);
    assert_eq!(index.len(), 3);
    index
        .compact()
        .expect("explicit base compaction should work");
    assert_eq!(index.sealed_len(), 0);
    assert_eq!(index.delta_len(), 0);
}

#[test]
fn mutable_bundle_roundtrip_preserves_metadata_tombstones_and_both_deltas() {
    let path = bundle_path();
    let records = vec![
        VectorRecord {
            key: 10,
            vector: vec![1.0, 0.0],
            metadata: metadata("base", 1),
        },
        VectorRecord {
            key: 20,
            vector: vec![0.0, 1.0],
            metadata: metadata("base", 2),
        },
    ];
    let index = MutableVectorIndex::build(config(2), &records).expect("mutable index should build");
    index
        .insert(30, &[0.9, 0.1], metadata("sealed", 3))
        .expect("sealed record should insert");
    index.seal_delta().expect("delta should seal");
    index
        .insert(40, &[0.8, 0.2], metadata("staged", 4))
        .expect("staged record should insert");
    assert!(index.delete(20));
    let active = [(10, [1.0, 0.0]), (30, [0.9, 0.1]), (40, [0.8, 0.2])];
    let active = active
        .iter()
        .map(|(key, vector)| (*key, vector.as_slice()))
        .collect::<Vec<_>>();
    let quantized = QuantizedIndex::build(config(2), QuantizationKind::Float16, &active)
        .expect("bundle quantized index should build");
    let bundle = IndexBundle::new(index, Some(quantized)).expect("complete bundle should validate");
    bundle.save(&path).expect("complete bundle should save");

    let loaded = IndexBundle::load(&path).expect("complete bundle should load");
    let mutable = loaded.mutable();
    assert_eq!(mutable.len(), 3);
    assert_eq!(mutable.sealed_len(), 1);
    assert_eq!(mutable.staged_len(), 1);
    assert_eq!(
        mutable
            .search(&[1.0, 0.0], 4)
            .expect("loaded bundle should search")
            .iter()
            .map(|hit| hit.key)
            .collect::<Vec<_>>(),
        vec![10, 30, 40]
    );
    assert_eq!(
        mutable
            .metadata(40)
            .and_then(|value| value.get("group").cloned()),
        Some(MetadataValue::Text("staged".to_owned()))
    );
    assert!(mutable.metadata(20).is_none());
    assert_eq!(
        loaded
            .quantized()
            .expect("quantized state should load")
            .search(&[1.0, 0.0], 3)
            .expect("loaded quantized state should search")
            .iter()
            .map(|hit| hit.key)
            .collect::<Vec<_>>(),
        vec![10, 30, 40]
    );
    std::fs::remove_file(path).expect("bundle should be removable");
}