velesdb-core 7.0.0

//! Tests for `DirectVectorWriter`.

use super::direct_writer::DirectVectorWriter;
use super::upsert::UpsertResult;
use super::HnswIndex;
use crate::distance::DistanceMetric;

/// Creates a test `HnswIndex` with the given dimension and its exact-distance
/// features on.
///
/// Euclidean on purpose: its arena stores vectors as given, so the placement
/// tests below can compare them exactly. A cosine arena normalizes them — see
/// `a_direct_write_into_a_cosine_arena_is_unit_norm`.
fn make_index(dim: usize) -> HnswIndex {
    HnswIndex::new(dim, DistanceMetric::Euclidean).expect("test index creation")
}

/// Creates a test `HnswIndex` with its exact-distance features off
/// (`new_fast_insert`): graph inserts still store every vector, the direct
/// writer places none.
fn make_fast_insert_index(dim: usize) -> HnswIndex {
    HnswIndex::new_fast_insert(dim, DistanceMetric::Cosine).expect("test index creation")
}

/// Reads a vector back from the graph's `ContiguousVectors`.
fn contiguous_get(index: &HnswIndex, idx: usize) -> Option<Vec<f32>> {
    index
        .inner
        .read()
        .with_contiguous_vectors_read(|cv| cv.get(idx).map(<[f32]>::to_vec))
}

/// `write_batch_direct` into an index with its exact-distance features on,
/// which places every vector.
fn placed(writer: &DirectVectorWriter<'_>, batch: &[(u64, &[f32])]) -> Vec<UpsertResult> {
    writer
        .write_batch_direct(batch)
        .expect("test: direct write")
        .expect("test: an index with exact-distance features on places")
}

#[test]
fn test_write_batch_direct_empty() {
    let index = make_index(4);
    let writer = DirectVectorWriter::new(&index);
    let results = placed(&writer, &[]);
    assert!(results.is_empty());
}

#[test]
fn test_write_batch_direct_single_vector() {
    let index = make_index(4);
    let writer = DirectVectorWriter::new(&index);
    let vec = [1.0_f32, 2.0, 3.0, 4.0];
    let results = placed(&writer, &[(1, &vec)]);

    assert_eq!(results.len(), 1);
    assert!(results[0].old_idx.is_none());

    // Verify mapping exists
    assert!(index.mappings.get_idx(1).is_some());

    // Verify vector is in ContiguousVectors (single vector store)
    assert_eq!(contiguous_get(&index, results[0].idx), Some(vec.to_vec()));
}

#[test]
fn test_write_batch_direct_multiple_vectors() {
    let index = make_index(3);
    let writer = DirectVectorWriter::new(&index);
    let v1 = [1.0_f32, 0.0, 0.0];
    let v2 = [0.0_f32, 1.0, 0.0];
    let v3 = [0.0_f32, 0.0, 1.0];
    let batch: Vec<(u64, &[f32])> = vec![(10, &v1), (20, &v2), (30, &v3)];

    let results = placed(&writer, &batch);
    assert_eq!(results.len(), 3);

    // All mappings registered
    assert!(index.mappings.get_idx(10).is_some());
    assert!(index.mappings.get_idx(20).is_some());
    assert!(index.mappings.get_idx(30).is_some());

    // Every vector landed in ContiguousVectors at its assigned slot.
    assert_eq!(contiguous_get(&index, results[0].idx), Some(v1.to_vec()));
    assert_eq!(contiguous_get(&index, results[1].idx), Some(v2.to_vec()));
    assert_eq!(contiguous_get(&index, results[2].idx), Some(v3.to_vec()));
}

/// Vectors written by `write_batch_direct` are immediately visible to
/// brute-force search (the deferred-HNSW window must not hide them).
#[test]
fn test_write_batch_direct_visible_to_brute_force() {
    let index = make_index(3);
    let writer = DirectVectorWriter::new(&index);
    let v1 = [1.0_f32, 2.0, 3.0];
    let v2 = [4.0_f32, 5.0, 6.0];
    let batch: Vec<(u64, &[f32])> = vec![(1, &v1), (2, &v2)];

    placed(&writer, &batch);

    let results = index.brute_force_search_parallel(&v1, 2).unwrap();
    assert_eq!(results.len(), 2);
    assert_eq!(results[0].id, 1);
}

#[test]
fn test_upsert_deduplication() {
    let index = make_index(2);
    let writer = DirectVectorWriter::new(&index);
    let v1 = [1.0_f32, 0.0];
    let v2 = [0.0_f32, 1.0];

    // Insert ID=1 twice — second should replace first
    let r1 = placed(&writer, &[(1, &v1)]);
    let r2 = placed(&writer, &[(1, &v2)]);

    assert!(r1[0].old_idx.is_none());
    assert!(r2[0].old_idx.is_some());
    assert_eq!(r2[0].old_idx, Some(r1[0].idx));

    // Only one mapping for ID=1
    assert_eq!(index.mappings.len(), 1);
    let current_idx = index.mappings.get_idx(1).unwrap();
    assert_eq!(current_idx, r2[0].idx);
}

#[test]
fn test_dimension_mismatch_returns_error() {
    let index = make_index(4);
    let writer = DirectVectorWriter::new(&index);
    let wrong_dim = [1.0_f32, 2.0]; // dim=2, expected=4

    let result = writer.write_batch_direct(&[(1, &wrong_dim)]);
    assert!(result.is_err());

    // State unchanged
    assert!(index.mappings.is_empty());
}

/// With exact-distance features off the direct writer places no vector, so
/// it maps nothing: mapping `id` to a slot it never filled would leave that
/// slot for another writer's vector to land in (#2246). The deferred graph
/// insert places the vector and maps it instead.
#[test]
fn test_direct_write_skipped_with_exact_distance_features_off() {
    let index = make_fast_insert_index(3);
    let writer = DirectVectorWriter::new(&index);
    let v = [1.0_f32, 2.0, 3.0];

    let results = writer.write_batch_direct(&[(1, &v)]).unwrap();

    assert!(
        results.is_none(),
        "the writer reports that it placed nothing"
    );
    assert_eq!(index.mappings.get_idx(1), None);
    assert_eq!(
        contiguous_get(&index, 0),
        None,
        "the writer placed a vector it does not map"
    );
}

/// A wrong dimension is refused even with exact-distance features off: the
/// `# Errors` contract holds whatever the index's features.
#[test]
fn test_a_wrong_dimension_is_refused_with_exact_distance_features_off() {
    let index = make_fast_insert_index(3);
    let writer = DirectVectorWriter::new(&index);
    let wrong = [1.0_f32, 2.0, 3.0, 4.0];
    let err = writer
        .write_batch_direct(&[(1, &wrong[..])])
        .expect_err("test: a 4-d vector in a 3-d index");
    assert!(
        matches!(err, crate::error::Error::DimensionMismatch { .. }),
        "{err:?}"
    );
}

/// A pre-normalized engine keeps its cosine arena unit-norm, and `.vectors`
/// declares it; a direct write must keep it so, or that declaration is false
/// for every slot written here (#2246).
#[test]
fn a_direct_write_into_a_cosine_arena_is_unit_norm() {
    let index = HnswIndex::new(4, DistanceMetric::Cosine).expect("test index creation");
    let writer = DirectVectorWriter::new(&index);
    let raw = [3.0_f32, 0.0, 0.0, 4.0];
    let results = placed(&writer, &[(1, &raw)]);

    let stored = contiguous_get(&index, results[0].idx).expect("stored");
    for (got, want) in stored.iter().zip([0.6_f32, 0.0, 0.0, 0.8]) {
        assert!(
            (got - want).abs() < 1e-6,
            "stored {stored:?}, want the unit vector"
        );
    }
}