haematite-cli 0.21.7

//! V5: exact search over an embedded field's vectors in the companion.

use std::error::Error;
use std::num::{NonZeroU32, NonZeroUsize};

use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::value::RawValue;

use super::{Found, Searched, State, search};
use crate::serve::access::{Access, Standalone, TestSource, request};
use crate::serve::protocol::{
    Change, CollectionWrite, DocumentImport, DocumentWrite, Fork, ImportedDocument, Metric,
    Mutation, Selection, Version,
};
use crate::serve::storage::embed_job::{Made, Page, Reach, field_key};
use crate::serve::storage::embedded::{
    self, Declaration, ModelName, ModelTag, Position, Provenance, Toward, VectorKind,
};
use crate::serve::storage::{
    Captured, DocumentRevision, Failure, Storage, branch_name, document_key, reads, revision_key,
    version,
};

type TestResult = Result<(), Box<dyn Error>>;

const FIELD: &str = "body_vec";

/// The compared bound a search is handed where no test sets one: a stand-in
/// for the setting's value.
const BOUND: u64 = 1_000;

fn create() -> Result<(tempfile::TempDir, Storage), Box<dyn Error>> {
    let directory = tempfile::tempdir()?;
    let mut storage = Storage::create(
        DatabaseConfig {
            data_dir: directory.path().join("database"),
            shard_count: 2,
            distributed: None,
            executor_threads: Some(1),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::reclaim(0),
        },
        NodeCacheBudget::Unlimited,
        fixture_long_wait_interval(),
    )?;
    mutate(
        &mut storage,
        Change::Collection(CollectionWrite {
            collection: "notes".to_owned(),
            schema: None,
        }),
    )?;
    Ok((directory, storage))
}

fn head(storage: &Storage) -> Result<Captured, Failure> {
    let record = storage
        .refs
        .get(&branch_name("main"))
        .ok_or_else(|| Failure::new("branch_missing", "main"))?;
    Ok(Captured {
        branch: "main".to_owned(),
        version: version(record),
        shards: record.shards.clone(),
    })
}

fn mutate(storage: &mut Storage, change: Change) -> Result<Version, Failure> {
    let at = head(storage)?.version;
    let receipt = storage.mutate(&Mutation {
        request_id: format!("search-{}-{}", at.generation, at.sequence),
        branch: "main".to_owned(),
        expected: at,
        change,
    })?;
    Ok(receipt.version)
}

fn put(storage: &mut Storage, id: &str, body: &str) -> Result<Version, Box<dyn Error>> {
    let change = Change::Put(DocumentWrite {
        collection: "notes".to_owned(),
        document_id: id.to_owned(),
        document: RawValue::from_string(format!(r#"{{"body":"{body}"}}"#))?,
    });
    Ok(mutate(storage, change)?)
}

fn model() -> ModelName {
    ModelName {
        name: "stand-in".to_owned(),
        revision: "1".to_owned(),
        manifest_sha256: "abc".to_owned(),
    }
}

fn declare(storage: &mut Storage) -> Result<(), Box<dyn Error>> {
    let declaration = Declaration {
        kind: VectorKind::Dense,
        source: "body".to_owned(),
        model: model(),
        dimensions: NonZeroUsize::new(2).ok_or("two is not zero")?,
        distance: Metric::Cosine,
    };
    let bytes = embedded::encode_declaration(&declaration)?;
    let key = field_key("notes", FIELD);
    storage.commit_namespace(&branch_name("main"), &[(key, Some(bytes))])?;
    Ok(())
}

/// The record's revision at the head.
fn revision(storage: &Storage, id: &str) -> Result<Version, Box<dyn Error>> {
    let key = revision_key("notes", id);
    let bytes = reads::get_bytes(&storage.reader, &head(storage)?, &key)?.ok_or("no revision")?;
    Ok(serde_json::from_slice::<DocumentRevision>(&bytes)?.written_at)
}

/// Writes each record with its vector, made from the revision it is at, in
/// one walk ending at the head.
fn records(storage: &mut Storage, records: &[(&str, [f32; 2])]) -> TestResult {
    for (id, _vector) in records {
        put(storage, id, "text")?;
    }
    embed(storage, records)
}

/// One walk to the head with these vectors, each made from its record's
/// current revision.
fn embed(storage: &mut Storage, records: &[(&str, [f32; 2])]) -> TestResult {
    let mut made = Vec::new();
    for (id, vector) in records {
        made.push(Made {
            id: (*id).to_owned(),
            made_from: revision(storage, id)?,
            vector: vector.to_vec(),
        });
    }
    let page = Page {
        collection: "notes".to_owned(),
        field: FIELD.to_owned(),
        model: model(),
        made,
        removed: Vec::new(),
        cleared: Vec::new(),
        windowed: Vec::new(),
        toward: head(storage)?.version,
        reach: Reach::End,
    };
    storage.embed_commit("main", &page, NonZeroUsize::MIN.saturating_add(63))?;
    Ok(())
}

fn find(storage: &Storage, access: &Access, question: &[f32], k: usize) -> Result<Found, Failure> {
    find_within(storage, access, question, k, BOUND)
}

/// `find`, comparing at most `bound` vectors.
fn find_within(
    storage: &Storage,
    access: &Access,
    question: &[f32],
    k: usize,
    bound: u64,
) -> Result<Found, Failure> {
    let at = head(storage)?;
    let (declaration, companion) = storage.embedded_field(&at, "notes", FIELD)?;
    let searched = Searched {
        collection: "notes",
        field: FIELD,
        declaration: &declaration,
        at: &at,
        companion: companion.as_ref(),
    };
    let k = NonZeroUsize::new(k).ok_or_else(|| Failure::new("k", "k is zero"))?;
    search(access, &storage.reader, &searched, question, k, bound)
}

fn ids(found: &Found) -> Vec<&str> {
    found.hits.iter().map(|hit| hit.id.as_str()).collect()
}

/// Fails while the hits are not the best first with equal scores in id
/// order: two at cosine 1, then one at 0, of four compared.
#[test]
fn the_best_are_returned_first_with_ties_by_id() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    let vectors = [
        ("n-3", [1.0, 0.0]),
        ("n-2", [0.0, 1.0]),
        ("n-1", [2.0, 0.0]),
        ("n-4", [-1.0, 0.0]),
    ];
    records(&mut storage, &vectors)?;
    let found = find(&storage, &request(&Standalone)?, &[1.0, 0.0], 3)?;
    assert_eq!(ids(&found), ["n-1", "n-3", "n-2"]);
    assert_eq!((found.compared, found.stale), (4, 0));
    Ok(())
}

/// Fails while a vector of text that has changed is returned as current: it
/// is left out and counted as stale.
#[test]
fn a_stale_vector_is_left_out_and_counted() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    records(&mut storage, &[("n-1", [1.0, 0.0]), ("n-2", [0.5, 0.5])])?;
    put(&mut storage, "n-1", "rewritten")?;
    let found = find(&storage, &request(&Standalone)?, &[1.0, 0.0], 5)?;
    assert_eq!(ids(&found), ["n-2"]);
    assert_eq!((found.compared, found.stale), (1, 1));
    Ok(())
}

/// Fails while a record the caller may not read is compared, returned or
/// counted.
#[test]
fn a_hidden_record_is_never_compared_or_counted() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    records(&mut storage, &[("n-1", [1.0, 0.0]), ("n-2", [0.5, 0.5])])?;
    put(&mut storage, "n-2", "rewritten")?;
    let hidden = TestSource::denying(&[("notes", "n-1"), ("notes", "n-2")]);
    let found = find(&storage, &request(&hidden)?, &[1.0, 0.0], 5)?;
    assert_eq!(found.hits, []);
    assert_eq!((found.compared, found.stale), (0, 0));
    Ok(())
}

/// Fails while vectors of two models can answer together: one of another
/// model refuses the search by name.
#[test]
fn another_models_vector_refuses_the_search() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    records(&mut storage, &[("n-1", [1.0, 0.0])])?;
    let foreign = Provenance {
        made_from: revision(&storage, "n-1")?,
        tag: ModelTag([9; 16]),
    };
    let key = embedded::vector_key("notes", FIELD, "n-1");
    let value = embedded::encode_embedded(&foreign, &[1.0, 0.0]);
    let companion = embedded::companion_name("main");
    storage.commit_namespace(&companion, &[(key, Some(value))])?;
    let refused = find(&storage, &request(&Standalone)?, &[1.0, 0.0], 5);
    assert_eq!(
        refused.err().map(|failure| failure.code).as_deref(),
        Some("embed_model_mixed")
    );
    Ok(())
}

/// Fails while a question of another length than declared is scored.
#[test]
fn a_question_of_another_length_is_refused() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    records(&mut storage, &[("n-1", [1.0, 0.0])])?;
    let refused = find(&storage, &request(&Standalone)?, &[1.0], 5);
    assert_eq!(
        refused.err().map(|failure| failure.code).as_deref(),
        Some("embed_question_length")
    );
    Ok(())
}

/// Fails while the join takes a vector whose record is gone at the searched
/// version for another record's: it is not compared, and not stale.
#[test]
fn a_vector_whose_record_is_gone_is_not_compared() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    records(&mut storage, &[("n-1", [1.0, 0.0]), ("n-2", [0.5, 0.5])])?;
    // Removed directly: the delete verb refuses a declared field until V6.
    let removed = [
        (document_key("notes", "n-1"), None),
        (revision_key("notes", "n-1"), None),
    ];
    storage.commit_namespace(&branch_name("main"), &removed)?;
    let found = find(&storage, &request(&Standalone)?, &[1.0, 0.0], 5)?;
    assert_eq!(ids(&found), ["n-2"]);
    assert_eq!((found.compared, found.stale), (1, 0));
    Ok(())
}

/// Fails while a shard's scan stops after its first page, or the merge after
/// a refill drops or repeats a vector: 200 records over two shards, more than
/// a page in each, written in one import and embedded in four pages of 64,
/// are joined whole. Record `i` has the vector `[1, 199 - i]`, so the best
/// three for `[1, 0]` are the last three; one record past the first page is
/// rewritten, so one vector is stale.
#[test]
fn a_field_larger_than_one_page_is_joined_whole() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    let names: Vec<String> = (0..200_u16).map(|index| format!("n-{index:03}")).collect();
    let mut documents = Vec::new();
    for id in &names {
        documents.push(ImportedDocument {
            document_id: id.clone(),
            document: RawValue::from_string(r#"{"body":"text"}"#.to_owned())?,
        });
    }
    let import = Mutation {
        request_id: "search-import".to_owned(),
        branch: "main".to_owned(),
        expected: head(&storage)?.version,
        change: Change::PutMany(DocumentImport {
            collection: "notes".to_owned(),
            documents,
        }),
    };
    storage.mutate_bounded(&import, NonZeroU32::new(4).ok_or("four is not zero")?)?;
    let toward = head(&storage)?.version;
    let written = revision(&storage, "n-000")?;
    let mut index = 0_u16;
    let pages: Vec<&[String]> = names.chunks(64).collect();
    for (number, page) in pages.iter().enumerate() {
        let mut made = Vec::new();
        for id in *page {
            made.push(Made {
                id: id.clone(),
                made_from: written,
                vector: vec![1.0, f32::from(199 - index)],
            });
            index += 1;
        }
        let last = page.last().ok_or("an empty page")?;
        let reach = if number + 1 == pages.len() {
            Reach::End
        } else {
            Reach::Upto(last.clone())
        };
        let page = Page {
            collection: "notes".to_owned(),
            field: FIELD.to_owned(),
            model: model(),
            made,
            removed: Vec::new(),
            cleared: Vec::new(),
            windowed: Vec::new(),
            toward,
            reach,
        };
        storage.embed_commit("main", &page, NonZeroUsize::MIN.saturating_add(63))?;
    }
    put(&mut storage, "n-150", "rewritten")?;
    let found = find(&storage, &request(&Standalone)?, &[1.0, 0.0], 3)?;
    assert_eq!((found.compared, found.stale), (199, 1));
    assert_eq!(ids(&found), ["n-199", "n-198", "n-197"]);
    Ok(())
}

fn five(storage: &mut Storage) -> TestResult {
    declare(storage)?;
    let given = [
        ("a", [1.0, 0.0]),
        ("b", [0.0, 1.0]),
        ("c", [1.0, 1.0]),
        ("d", [1.0, 0.5]),
        ("e", [0.5, 1.0]),
    ];
    records(storage, &given)
}

fn code<T>(outcome: Result<T, Failure>) -> Option<String> {
    outcome.err().map(|failure| failure.code)
}

/// Fails while the search reads past its bound: five current vectors with
/// a bound of 4 refuse `query_scan_limit`, and with a bound of 5 all five
/// are compared.
#[test]
fn a_search_past_its_bound_is_refused() -> TestResult {
    let (_directory, mut storage) = create()?;
    five(&mut storage)?;
    let access = request(&Standalone)?;
    let refused = find_within(&storage, &access, &[1.0, 0.0], 3, 4);
    assert_eq!(code(refused).as_deref(), Some("query_scan_limit"));
    let found = find_within(&storage, &access, &[1.0, 0.0], 3, 5)?;
    assert_eq!(found.compared, 5);
    Ok(())
}

/// Fails while the bound is not the one handed in: the same five vectors
/// are searched whole with a bound of 6 and refused with a bound of 2.
#[test]
fn the_bound_is_the_value_handed_in() -> TestResult {
    let (_directory, mut storage) = create()?;
    five(&mut storage)?;
    let access = request(&Standalone)?;
    let found = find_within(&storage, &access, &[1.0, 0.0], 3, 6)?;
    assert_eq!(found.compared, 5);
    let refused = find_within(&storage, &access, &[1.0, 0.0], 3, 2);
    assert_eq!(code(refused).as_deref(), Some("query_scan_limit"));
    Ok(())
}

fn state(storage: &Storage, at: &Captured) -> Result<State, Failure> {
    Ok(storage.embedding_state(at, "notes", FIELD)?.state)
}

/// Writes `position` for the field into `owner`'s companion.
fn place(storage: &mut Storage, owner: &str, position: &Position) -> TestResult {
    storage.position_for_tests(owner, "notes", FIELD, position)?;
    Ok(())
}

/// The control for the three below: right after a walk ends at the head,
/// the state is complete, and it names the branch, `through` and the
/// searched version.
#[test]
fn a_walk_ended_at_the_searched_version_is_complete() -> TestResult {
    let (_directory, mut storage) = create()?;
    five(&mut storage)?;
    let at = head(&storage)?;
    let embedding = storage.embedding_state(&at, "notes", FIELD)?;
    assert_eq!(embedding.state, State::Complete);
    assert_eq!(embedding.branch, "main");
    assert_eq!(embedding.through, Some(at.version));
    assert_eq!(embedding.toward, None);
    assert_eq!(embedding.searched, at.version);
    Ok(())
}

/// Fails while an open walk is not looked at: `through` is the searched
/// version on the searched branch, but a walk is open, so it is partial.
#[test]
fn an_open_walk_is_partial() -> TestResult {
    let (_directory, mut storage) = create()?;
    five(&mut storage)?;
    let at = head(&storage)?;
    let open = Position {
        branch: "main".to_owned(),
        through: Some(at.version),
        toward: Some(Toward {
            version: at.version,
            after: "c".to_owned(),
        }),
    };
    place(&mut storage, "main", &open)?;
    assert_eq!(state(&storage, &at)?, State::Partial);
    Ok(())
}

/// Fails while the position's branch is not looked at: a fork's position
/// inherited from main, whose `through` equals the fork's searched version
/// and with no walk open, is partial, as it names main.
#[test]
fn an_inherited_position_is_partial() -> TestResult {
    let (_directory, mut storage) = create()?;
    five(&mut storage)?;
    let from = head(&storage)?.version;
    mutate(
        &mut storage,
        Change::Fork(Fork {
            name: "feature".to_owned(),
            from,
        }),
    )?;
    let at = storage.capture(&Selection {
        branch: "feature".to_owned(),
        version: None,
    })?;
    let inherited = Position {
        branch: "main".to_owned(),
        through: Some(at.version),
        toward: None,
    };
    place(&mut storage, "feature", &inherited)?;
    assert_eq!(state(&storage, &at)?, State::Partial);
    Ok(())
}

/// Fails while any `through` passes for the searched version: a record
/// written after the walk ended moves the head, and the same position is
/// partial at the new head.
#[test]
fn an_older_through_is_partial() -> TestResult {
    let (_directory, mut storage) = create()?;
    five(&mut storage)?;
    put(&mut storage, "f", "later")?;
    let at = head(&storage)?;
    assert_eq!(state(&storage, &at)?, State::Partial);
    Ok(())
}

/// Fails while nothing embedded yet answers: with no companion, and with a
/// companion holding no position for the field, the state is refused
/// `embed_search_not_ready`.
#[test]
fn nothing_embedded_yet_is_refused() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    put(&mut storage, "a", "text")?;
    let at = head(&storage)?;
    assert_eq!(
        code(state(&storage, &at)).as_deref(),
        Some("embed_search_not_ready")
    );
    let other = Position {
        branch: "main".to_owned(),
        through: Some(at.version),
        toward: None,
    };
    storage.position_for_tests("main", "notes", "title_vec", &other)?;
    assert_eq!(
        code(state(&storage, &at)).as_deref(),
        Some("embed_search_not_ready")
    );
    Ok(())
}

/// Fails while stale vectors are read without counting: 2 current and 5
/// stale vectors with a bound of 4 refuse `query_scan_limit`, though only 2
/// would be compared.
#[test]
fn stale_vectors_count_against_the_bound() -> TestResult {
    let (_directory, mut storage) = create()?;
    five(&mut storage)?;
    records(&mut storage, &[("f", [1.0, 0.0]), ("g", [0.0, 1.0])])?;
    for id in ["a", "b", "c", "d", "e"] {
        put(&mut storage, id, "rewritten")?;
    }
    let access = request(&Standalone)?;
    let found = find_within(&storage, &access, &[1.0, 0.0], 3, 7)?;
    assert_eq!((found.compared, found.stale), (2, 5));
    let refused = find_within(&storage, &access, &[1.0, 0.0], 3, 4);
    assert_eq!(code(refused).as_deref(), Some("query_scan_limit"));
    Ok(())
}

/// Fails while hidden vectors are read without counting, or the refusal
/// tells how many were hidden: 2 current and 2 hidden vectors with a bound
/// of 3 refuse with the very words 4 current vectors with a bound of 3 do.
#[test]
fn hidden_vectors_count_against_the_bound_and_are_never_told() -> TestResult {
    let (_directory, mut storage) = create()?;
    declare(&mut storage)?;
    let given = [
        ("a", [1.0, 0.0]),
        ("b", [0.0, 1.0]),
        ("c", [1.0, 1.0]),
        ("d", [1.0, 0.5]),
    ];
    records(&mut storage, &given)?;
    let hidden = TestSource::denying(&[("notes", "c"), ("notes", "d")]);
    let hidden = find_within(&storage, &request(&hidden)?, &[1.0, 0.0], 3, 3).err();
    let open = find_within(&storage, &request(&Standalone)?, &[1.0, 0.0], 3, 3).err();
    let hidden = hidden.ok_or("the hidden search was not refused")?;
    let open = open.ok_or("the open search was not refused")?;
    assert_eq!(hidden.code, "query_scan_limit");
    assert_eq!(hidden.message, open.message);
    Ok(())
}

#[path = "embed_search_windows_tests.rs"]
mod windows;

const fn fixture_long_wait_interval() -> std::time::Duration {
    const LONG_WAIT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);
    LONG_WAIT_INTERVAL
}