haematite-cli 0.21.8

//! F7: a companion loses the vectors of fields its branch does not declare.

use std::error::Error;
use std::num::NonZeroUsize;

use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;

use super::Pruned;
use crate::serve::protocol::{Metric, Version};
use crate::serve::storage::embed_job::field_key;
use crate::serve::storage::embedded::{self, Declaration, ModelName, Position, VectorKind};
use crate::serve::storage::{Captured, Storage, branch_name, reads, version};

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

const DECLARED: &str = "body_vec";
const UNDECLARED: &str = "title_vec";

/// A store whose `main` declares `body_vec` and has a companion holding its
/// position.
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(),
    )?;
    let declaration = Declaration {
        kind: VectorKind::Dense,
        source: "body".to_owned(),
        model: ModelName {
            name: "stand-in".to_owned(),
            revision: "1".to_owned(),
            manifest_sha256: "abc".to_owned(),
        },
        dimensions: NonZeroUsize::new(2).ok_or("two is not zero")?,
        distance: Metric::Cosine,
    };
    let bytes = embedded::encode_declaration(&declaration)?;
    let key = field_key("notes", DECLARED);
    storage.commit_namespace(&branch_name("main"), &[(key, Some(bytes))])?;
    position(&mut storage, DECLARED)?;
    Ok((directory, storage))
}

/// A finished position for `field` in main's companion.
fn position(storage: &mut Storage, field: &str) -> TestResult {
    let through = storage
        .refs
        .get(&branch_name("main"))
        .map(version)
        .ok_or("no main")?;
    let position = Position {
        branch: "main".to_owned(),
        through: Some(through),
        toward: None,
    };
    storage.position_for_tests("main", "notes", field, &position)?;
    Ok(())
}

/// `count` vectors of `field` in main's companion.
fn vectors(storage: &mut Storage, field: &str, count: usize) -> TestResult {
    let writes: Vec<(Vec<u8>, Option<Vec<u8>>)> = (0..count)
        .map(|index| {
            let key = embedded::vector_key("notes", field, &format!("r-{index}"));
            (key, Some(b"vector".to_vec()))
        })
        .collect();
    storage.commit_namespace(&embedded::companion_name("main"), &writes)?;
    Ok(())
}

fn companion(storage: &Storage) -> Result<Captured, Box<dyn Error>> {
    let name = embedded::companion_name("main");
    let record = storage.refs.get(&name).ok_or("no companion")?;
    Ok(Captured {
        branch: name.clone(),
        version: version(record),
        shards: record.shards.clone(),
    })
}

fn companion_version(storage: &Storage) -> Result<Version, Box<dyn Error>> {
    Ok(companion(storage)?.version)
}

/// How many of `count` vectors of `field` the companion still holds.
fn held(storage: &Storage, field: &str, count: usize) -> Result<usize, Box<dyn Error>> {
    let at = companion(storage)?;
    let mut held = 0;
    for index in 0..count {
        let key = embedded::vector_key("notes", field, &format!("r-{index}"));
        if reads::get_bytes(&storage.reader, &at, &key)?.is_some() {
            held += 1;
        }
    }
    Ok(held)
}

fn prune(storage: &mut Storage) -> Result<Pruned, Box<dyn Error>> {
    let page = NonZeroUsize::new(2).ok_or("two is not zero")?;
    Ok(storage.embed_prune("main", page)?)
}

/// Fails while a call removes more than its page or touches a declared
/// field: 5 vectors of an undeclared field go 2, 2 and 1, the last call
/// answering no more, and the declared field's 3 all stay.
#[test]
fn an_undeclared_fields_vectors_go_a_page_at_a_time() -> TestResult {
    let (_directory, mut storage) = create()?;
    vectors(&mut storage, DECLARED, 3)?;
    vectors(&mut storage, UNDECLARED, 5)?;
    let taken = [
        prune(&mut storage)?,
        prune(&mut storage)?,
        prune(&mut storage)?,
    ];
    let expected = [
        Pruned {
            removed: 2,
            more: true,
        },
        Pruned {
            removed: 2,
            more: true,
        },
        Pruned {
            removed: 1,
            more: false,
        },
    ];
    assert_eq!(taken, expected);
    assert_eq!(held(&storage, UNDECLARED, 5)?, 0);
    assert_eq!(held(&storage, DECLARED, 3)?, 3);
    Ok(())
}

/// Fails while a call with nothing to prune commits: the companion's
/// version is unchanged.
#[test]
fn nothing_to_prune_commits_nothing() -> TestResult {
    let (_directory, mut storage) = create()?;
    vectors(&mut storage, DECLARED, 3)?;
    let before = companion_version(&storage)?;
    let nothing = Pruned {
        removed: 0,
        more: false,
    };
    assert_eq!(prune(&mut storage)?, nothing);
    assert_eq!(companion_version(&storage)?, before);
    assert_eq!(held(&storage, DECLARED, 3)?, 3);
    Ok(())
}

/// Fails while a field's position is not respected: an undeclared field
/// that still has a position keeps its vectors, and nothing is committed.
#[test]
fn an_undeclared_field_with_a_position_is_left_alone() -> TestResult {
    let (_directory, mut storage) = create()?;
    position(&mut storage, UNDECLARED)?;
    vectors(&mut storage, UNDECLARED, 3)?;
    let before = companion_version(&storage)?;
    assert_eq!(prune(&mut storage)?.removed, 0);
    assert_eq!(companion_version(&storage)?, before);
    assert_eq!(held(&storage, UNDECLARED, 3)?, 3);
    Ok(())
}

/// Fails while a corrupt vector key is guessed at: a key whose field part
/// is not UTF-8 makes the call refuse by name and commit nothing.
#[test]
fn a_corrupt_vector_key_is_refused() -> TestResult {
    let (_directory, mut storage) = create()?;
    let key = b"v\0notes\0\xff\0r-0".to_vec();
    let name = embedded::companion_name("main");
    storage.commit_namespace(&name, &[(key, Some(b"vector".to_vec()))])?;
    let before = companion_version(&storage)?;
    let refused = storage.embed_prune("main", NonZeroUsize::MIN).err();
    let code = refused.map(|failure| failure.code);
    assert_eq!(code.as_deref(), Some("embed_vector_key_corrupt"));
    assert_eq!(companion_version(&storage)?, before);
    Ok(())
}

/// Window indices a record holds here: the second's last byte is not UTF-8.
const WINDOWS: [u32; 2] = [0, 0xff];

/// `count` records of `field` in main's companion, each as two windows.
fn windows(storage: &mut Storage, field: &str, count: usize) -> TestResult {
    let mut writes = Vec::new();
    for index in 0..count {
        for window in WINDOWS {
            let key = embedded::window_key("notes", field, &format!("r-{index}"), window);
            writes.push((key, Some(b"window".to_vec())));
        }
    }
    storage.commit_namespace(&embedded::companion_name("main"), &writes)?;
    Ok(())
}

/// How many window keys of `count` records of `field` the companion holds.
fn windows_held(storage: &Storage, field: &str, count: usize) -> Result<usize, Box<dyn Error>> {
    let at = companion(storage)?;
    let mut held = 0;
    for index in 0..count {
        for window in WINDOWS {
            let key = embedded::window_key("notes", field, &format!("r-{index}"), window);
            if reads::get_bytes(&storage.reader, &at, &key)?.is_some() {
                held += 1;
            }
        }
    }
    Ok(held)
}

/// EMB-E2 (f), `embed_prune.rs:65-163`: fails while an undeclared field's
/// windows are kept, a declared field's are taken, or a window key whose
/// index is not UTF-8 is refused as corrupt: after the field's vector, its
/// four window keys go two a call, and the declared field's stay.
#[test]
fn an_undeclared_fields_windows_go_after_its_vectors() -> TestResult {
    let (_directory, mut storage) = create()?;
    vectors(&mut storage, UNDECLARED, 1)?;
    windows(&mut storage, UNDECLARED, 2)?;
    windows(&mut storage, DECLARED, 1)?;
    let taken = [
        prune(&mut storage)?,
        prune(&mut storage)?,
        prune(&mut storage)?,
    ];
    let expected = [
        Pruned {
            removed: 1,
            more: true,
        },
        Pruned {
            removed: 2,
            more: true,
        },
        Pruned {
            removed: 2,
            more: false,
        },
    ];
    assert_eq!(taken, expected);
    assert_eq!(held(&storage, UNDECLARED, 1)?, 0);
    assert_eq!(windows_held(&storage, UNDECLARED, 2)?, 0);
    assert_eq!(windows_held(&storage, DECLARED, 1)?, 2);
    Ok(())
}

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
}