haematite-cli 0.21.8

//! V5 (ruled 5 Oct 08:15, Q5 to Q7): exact search over an embedded field's
//! vectors in the branch's companion.
//!
//! Every vector of the field is compared, so the search reads in proportion
//! to the field's records, as exact search does. The companion's vectors and
//! the branch's revisions are both read in id order, a page at a time from
//! each shard, and joined: no point read for each vector. A vector is
//! compared only if its record exists at the searched version, the caller may
//! read it, and it was made from the record's revision there; one made from
//! any other revision is stale, left out and counted. A record the caller
//! may not read is never compared, returned or counted (Q6). A vector of any
//! model but the declared one refuses the search by name: models are never
//! mixed in one answer (Q7).
//!
//! V8a (ruled 5 Oct 09:34, the bound corrected 09:37): a search reads at
//! most the bound it is handed of the field's vectors, a setting's value
//! (`query_scan_records`, less what earlier pages read), and refuses
//! `query_scan_limit` past it, never answering from part of the field. Every
//! vector row read counts against it, current, stale, hidden or with no
//! record alike, so the bound is on work; what is hidden is never counted in
//! anything a caller sees, and the refusal names only the field and the
//! bound. `Storage::embedding_state` says how far the job's position
//! has come against the searched version, so an answer tells everything
//! searched from what is embedded so far; with nothing embedded yet it
//! refuses `embed_search_not_ready`, never answering empty.
//!
//! EMB-E2 (ruled 6 Oct 21:19): a record embedded as windows (`vw\0`) is
//! read beside the one-vector records, in the same id order. Each window is
//! a vector read, compared or stale, and counted as one; the record's score
//! is its best window's, and its hit answers that window's span.

use std::collections::VecDeque;
use std::num::NonZeroUsize;

use haematite::db::ReadOnlyDatabase;
use serde::Serialize;

use super::embedded::{self, Declaration};
use super::{Captured, DocumentRevision, Failure, Storage, reads, version};
use crate::serve::access::Access;
use crate::serve::protocol::{Metric, Version};

/// Keys read from one shard at a time.
const SCAN_PAGE: usize = 64;

/// One hit: the record, its score, and for a record embedded as windows
/// the span of the best one (EMB-E2).
#[derive(Debug, Clone, PartialEq)]
pub struct Hit {
    pub id: String,
    pub score: f64,
    pub span: Option<HitSpan>,
}

/// EMB-E2: where a hit's best window lies in its record's text, in bytes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct HitSpan {
    pub start: u64,
    pub end: u64,
}

/// What one exact search answers.
#[derive(Debug, Clone, PartialEq)]
pub struct Found {
    /// Best first; equal scores by record id.
    pub hits: Vec<Hit>,
    /// Vectors compared with the question.
    pub compared: u64,
    /// Vectors left out because their record's text has changed since.
    pub stale: u64,
}

/// What one exact search is over.
pub struct Searched<'a> {
    pub collection: &'a str,
    pub field: &'a str,
    pub declaration: &'a Declaration,
    /// The branch at the searched version.
    pub at: &'a Captured,
    /// The branch's companion, or `None` before the job's first page.
    pub companion: Option<&'a Captured>,
}

/// How far the job has embedded a field against the searched version.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum State {
    /// Every vector was made from the revision searched: the position's
    /// branch is the searched branch, its `through` is the searched version,
    /// and no walk is open.
    Complete,
    /// Anything else: a search answers from what is embedded so far.
    Partial,
}

/// The job's position against the searched version, as an answer carries it.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct Embedding {
    pub state: State,
    /// The branch the position names: the searched branch, or the parent it
    /// was inherited from.
    pub branch: String,
    pub through: Option<Version>,
    /// The open walk's target, if one is open.
    pub toward: Option<Version>,
    pub searched: Version,
}

impl Storage {
    /// How far the job has embedded `collection.field` against `at`.
    ///
    /// # Errors
    /// `embed_search_not_ready` when `at`'s branch has no companion or the
    /// field no position: nothing is embedded yet, and a search never answers
    /// empty for it. Whatever a read answers.
    pub fn embedding_state(
        &self,
        at: &Captured,
        collection: &str,
        field: &str,
    ) -> Result<Embedding, Failure> {
        let not_ready = || {
            Failure::new(
                "embed_search_not_ready",
                format!(
                    "nothing of {collection}.{field} is embedded yet on {}",
                    at.branch
                ),
            )
        };
        let name = embedded::companion_name(&at.branch);
        if self.refs.get(&name).is_none() {
            return Err(not_ready());
        }
        let key = embedded::position_key(collection, field);
        let position = self.position(&name, &key)?.ok_or_else(not_ready)?;
        let toward = position.toward.map(|open| open.version);
        let complete = position.branch == at.branch
            && position.through == Some(at.version)
            && toward.is_none();
        let state = if complete {
            State::Complete
        } else {
            State::Partial
        };
        Ok(Embedding {
            state,
            branch: position.branch,
            through: position.through,
            toward,
            searched: at.version,
        })
    }

    /// The declaration of `collection.field` at `at`, and the companion of
    /// `at`'s branch as it is now.
    ///
    /// # Errors
    /// `embed_field_not_declared`; whatever a read answers.
    pub fn embedded_field(
        &self,
        at: &Captured,
        collection: &str,
        field: &str,
    ) -> Result<(Declaration, Option<Captured>), Failure> {
        let declaration = self.declared(at, collection, field)?;
        let name = embedded::companion_name(&at.branch);
        let companion = self.refs.get(&name).map(|record| Captured {
            branch: name.clone(),
            version: version(record),
            shards: record.shards.clone(),
        });
        Ok((declaration, companion))
    }

    /// Every (collection, field) declared embedded at `at` (EW1: the walks
    /// the embedding thread starts with).
    ///
    /// # Errors
    /// `vector_declaration_corrupt`; whatever a read answers.
    pub fn embedded_fields(&self, at: &Captured) -> Result<Vec<(String, String)>, Failure> {
        let mut fields = Vec::new();
        reads::scan_rows(&self.reader, at, b"w\0", 2, |parts, bytes| {
            if let ([collection, field], embedded::Field::Embedded(_)) =
                (parts, embedded::read_field(bytes)?)
            {
                fields.push(((*collection).to_owned(), (*field).to_owned()));
            }
            Ok(())
        })?;
        Ok(fields)
    }
}

/// The best `k` of the field's current vectors for `question`, reading at
/// most `bound` of the field's vectors.
///
/// # Errors
/// `embed_question_length` for a question of another length than declared;
/// `query_scan_limit` as soon as one more vector than `bound` would be read;
/// `embed_model_mixed` for a vector of another model; `vector_corrupt`;
/// whatever a read answers.
pub fn search(
    access: &Access,
    reader: &ReadOnlyDatabase,
    searched: &Searched<'_>,
    question: &[f32],
    k: NonZeroUsize,
    bound: u64,
) -> Result<Found, Failure> {
    let declaration = searched.declaration;
    let dimensions = declaration.dimensions.get();
    if question.len() != dimensions {
        return Err(Failure::new(
            "embed_question_length",
            format!(
                "the question has {} members; the field is declared with {dimensions}",
                question.len()
            ),
        ));
    }
    let mut found = Found {
        hits: Vec::new(),
        compared: 0,
        stale: 0,
    };
    let Some(companion) = searched.companion else {
        return Ok(found);
    };
    let judge = Judge {
        tag: embedded::model_tag(&declaration.model)?,
        dimensions,
        question,
        metric: declaration.distance,
    };
    let revisions_prefix = format!("m\0{}\0", searched.collection);
    let mut items = Items::new(reader, companion, searched, bound);
    let mut revisions = Ordered::new(reader, searched.at, revisions_prefix.as_bytes());
    let mut revision: Option<(Vec<u8>, Vec<u8>)> = revisions.next()?;
    while let Some((id, item)) = items.next()? {
        // Revisions before this id belong to records with no vector.
        while revision.as_ref().is_some_and(|(at, _)| *at < id) {
            revision = revisions.next()?;
        }
        let Some(current) = revision
            .as_ref()
            .filter(|(at, _)| *at == id)
            .map(|(_, current)| current)
        else {
            continue;
        };
        let id = String::from_utf8(id).map_err(|_not_utf8| {
            Failure::new("embed_record_id_corrupt", "a vector key names no record")
        })?;
        if !access.can_read_document(searched.collection, &id) {
            continue;
        }
        let written_at = serde_json::from_slice::<DocumentRevision>(current)
            .map_err(Failure::storage)?
            .written_at;
        let Some((score, span)) = judge.best(&id, &item, written_at, &mut found)? else {
            continue;
        };
        let hit = Hit { id, score, span };
        let at = found
            .hits
            .partition_point(|kept| better(kept, &hit, declaration.distance));
        if at < k.get() {
            found.hits.insert(at, hit);
            found.hits.truncate(k.get());
        }
    }
    Ok(found)
}

/// What a record's vectors are compared with.
struct Judge<'a> {
    tag: embedded::ModelTag,
    dimensions: usize,
    question: &'a [f32],
    metric: Metric,
}

/// One vector of a record, with its window's span if it is one.
type Decoded = (embedded::Provenance, Vec<f32>, Option<HitSpan>);

impl Judge<'_> {
    /// The best score of `item`'s vectors and, for windows, the best one's
    /// span; `None` when they were made from another revision than
    /// `written_at`. Every vector counts as compared or stale.
    ///
    /// # Errors
    /// As [`Self::read`].
    fn best(
        &self,
        id: &str,
        item: &Item,
        written_at: Version,
        found: &mut Found,
    ) -> Result<Option<(f64, Option<HitSpan>)>, Failure> {
        let vectors = self.read(id, item)?;
        let rows = u64::try_from(vectors.len()).unwrap_or(u64::MAX);
        if vectors
            .iter()
            .any(|(provenance, ..)| provenance.made_from != written_at)
        {
            found.stale = found.stale.saturating_add(rows);
            return Ok(None);
        }
        found.compared = found.compared.saturating_add(rows);
        let mut best: Option<(f64, Option<HitSpan>)> = None;
        for (_, vector, span) in &vectors {
            let score = score(self.question, vector, self.metric);
            if best.is_none_or(|(kept, _)| ahead(score, kept, self.metric)) {
                best = Some((score, *span));
            }
        }
        Ok(best)
    }

    /// `item`'s vectors decoded, a record's windows each in its place.
    ///
    /// # Errors
    /// `embed_model_mixed` for a vector of another model; `vector_corrupt`
    /// for windows out of place or another count than they name.
    fn read(&self, id: &str, item: &Item) -> Result<Vec<Decoded>, Failure> {
        let mut vectors = Vec::new();
        match item {
            Item::One(bytes) => {
                let (provenance, vector) = embedded::decode_embedded(bytes, self.dimensions)?;
                vectors.push((provenance, vector, None));
            }
            Item::Windows(rows) => {
                for (place, (index, bytes)) in rows.iter().enumerate() {
                    let (span, provenance, vector) =
                        embedded::decode_window(bytes, self.dimensions)?;
                    let placed = usize::try_from(*index).is_ok_and(|at| at == place);
                    let counted = usize::try_from(span.count).is_ok_and(|n| n == rows.len());
                    if !placed || !counted {
                        return Err(Failure::new(
                            "vector_corrupt",
                            format!("the windows of {id} are not the ones they name"),
                        ));
                    }
                    let span = HitSpan {
                        start: span.start,
                        end: span.end,
                    };
                    vectors.push((provenance, vector, Some(span)));
                }
            }
        }
        if vectors
            .iter()
            .any(|(provenance, ..)| provenance.tag != self.tag)
        {
            return Err(Failure::new(
                "embed_model_mixed",
                format!("the vector of {id} was made by another model than the declared one"),
            ));
        }
        Ok(vectors)
    }
}

/// Whether `score` is better than `kept` by `metric`; a tie keeps `kept`.
fn ahead(score: f64, kept: f64, metric: Metric) -> bool {
    match metric {
        Metric::L2 => score.total_cmp(&kept).is_lt(),
        Metric::Cosine | Metric::Dot => score.total_cmp(&kept).is_gt(),
    }
}

/// Whether `kept` comes before `hit`: the better score, then the lower id.
fn better(kept: &Hit, hit: &Hit, metric: Metric) -> bool {
    let by_score = match metric {
        Metric::L2 => kept.score.total_cmp(&hit.score),
        Metric::Cosine | Metric::Dot => hit.score.total_cmp(&kept.score),
    };
    by_score.then_with(|| kept.id.cmp(&hit.id)).is_lt()
}

/// The measure of `vector` against `question`, as `vectors::score` measures
/// a bare vector, line for line: the dot product, the distance, or the
/// cosine (0 when either has no length).
fn score(question: &[f32], vector: &[f32], metric: Metric) -> f64 {
    let mut dot = 0.0_f64;
    let mut question_norm = 0.0_f64;
    let mut vector_norm = 0.0_f64;
    let mut squared_distance = 0.0_f64;
    for (left, right) in question.iter().zip(vector) {
        let left = f64::from(*left);
        let right = f64::from(*right);
        dot = left.mul_add(right, dot);
        question_norm = left.mul_add(left, question_norm);
        vector_norm = right.mul_add(right, vector_norm);
        let difference = left - right;
        squared_distance = difference.mul_add(difference, squared_distance);
    }
    match metric {
        Metric::Dot => dot,
        Metric::L2 => squared_distance.sqrt(),
        Metric::Cosine if question_norm <= 0.0 || vector_norm <= 0.0 => 0.0,
        Metric::Cosine => dot / (question_norm.sqrt() * vector_norm.sqrt()),
    }
}

type StoredRow = (Vec<u8>, Vec<u8>);

/// A record's vectors as the companion holds them: one, or EMB-E2's windows
/// in key order with their indices.
enum Item {
    One(Vec<u8>),
    Windows(Vec<(u32, Vec<u8>)>),
}

/// A window row: its record id, its index and its value.
type WindowRow = (Vec<u8>, u32, Vec<u8>);

/// Every record's vectors of one field in id order: its `v\0` vector and
/// its `vw\0` windows read side by side, each row read counted against
/// `bound` and refused past it.
struct Items<'a> {
    vectors: Ordered<'a>,
    windows: Ordered<'a>,
    vector: Option<StoredRow>,
    window: Option<WindowRow>,
    started: bool,
    read: u64,
    bound: u64,
    field: String,
}

impl<'a> Items<'a> {
    fn new(
        reader: &'a ReadOnlyDatabase,
        companion: &Captured,
        searched: &Searched<'_>,
        bound: u64,
    ) -> Self {
        let (collection, field) = (searched.collection, searched.field);
        let vectors = format!("v\0{collection}\0{field}\0");
        let windows = embedded::window_prefix(collection, field);
        Self {
            vectors: Ordered::new(reader, companion, vectors.as_bytes()),
            windows: Ordered::new(reader, companion, &windows),
            vector: None,
            window: None,
            started: false,
            read: 0,
            bound,
            field: format!("{collection}.{field}"),
        }
    }

    /// One more row read: `query_scan_limit` past the bound.
    fn count(&mut self) -> Result<(), Failure> {
        if self.read >= self.bound {
            return Err(Failure::new(
                "query_scan_limit",
                format!(
                    "the search would read more than {} vectors of {}",
                    self.bound, self.field
                ),
            ));
        }
        self.read = self.read.saturating_add(1);
        Ok(())
    }

    fn pull_vector(&mut self) -> Result<Option<StoredRow>, Failure> {
        let row = self.vectors.next()?;
        if row.is_some() {
            self.count()?;
        }
        Ok(row)
    }

    fn pull_window(&mut self) -> Result<Option<WindowRow>, Failure> {
        let Some((key, value)) = self.windows.next()? else {
            return Ok(None);
        };
        self.count()?;
        let (id, index) = embedded::window_of(&[], &key)?;
        Ok(Some((id.into_bytes(), index, value)))
    }

    /// The next record's id and vectors.
    ///
    /// # Errors
    /// `query_scan_limit`; `vector_corrupt` for a record with a vector and
    /// windows both, or more windows than any record keeps; whatever a read
    /// answers.
    fn next(&mut self) -> Result<Option<(Vec<u8>, Item)>, Failure> {
        if !self.started {
            self.started = true;
            self.vector = self.pull_vector()?;
            self.window = self.pull_window()?;
        }
        let corrupt = |what: &str| Failure::new("vector_corrupt", format!("a record {what}"));
        let window_first = match (&self.vector, &self.window) {
            (None, None) => return Ok(None),
            (Some((id, _)), Some((window, ..))) if id == window => {
                return Err(corrupt("holds a vector and windows both"));
            }
            (Some((id, _)), Some((window, ..))) => window < id,
            (Some(_), None) => false,
            (None, Some(_)) => true,
        };
        if !window_first {
            let next = self.pull_vector()?;
            let one = std::mem::replace(&mut self.vector, next);
            return Ok(one.map(|(id, bytes)| (id, Item::One(bytes))));
        }
        let Some((id, index, bytes)) = self.window.take() else {
            return Ok(None);
        };
        let mut rows = vec![(index, bytes)];
        self.window = loop {
            match self.pull_window()? {
                Some((next, index, bytes)) if next == id => {
                    let held = u32::try_from(rows.len()).unwrap_or(u32::MAX);
                    if held >= embedded::WINDOWS_MOST {
                        return Err(corrupt("holds more windows than any record keeps"));
                    }
                    rows.push((index, bytes));
                }
                other => break other,
            }
        };
        Ok(Some((id, Item::Windows(rows))))
    }
}

/// One shard's place in an ordered read: its buffered page and where the
/// next page starts.
struct ShardScan {
    shard_id: usize,
    root: haematite::tree::Hash,
    buffered: VecDeque<StoredRow>,
    next_from: Option<Vec<u8>>,
}

/// Every key under a prefix at one version, across its shards, in key order,
/// with the prefix taken off: one page from each shard at a time.
struct Ordered<'a> {
    reader: &'a ReadOnlyDatabase,
    prefix: Vec<u8>,
    end: Vec<u8>,
    shards: Vec<ShardScan>,
}

impl<'a> Ordered<'a> {
    fn new(reader: &'a ReadOnlyDatabase, at: &Captured, prefix: &[u8]) -> Self {
        let mut end = prefix.to_vec();
        // Every prefix here ends in a NUL separator, so one past it is the
        // first key beyond every key under it.
        if let Some(last) = end.last_mut() {
            *last = last.saturating_add(1);
        }
        let shards = at
            .shards
            .iter()
            .map(|shard| ShardScan {
                shard_id: shard.shard_id,
                root: shard.head,
                buffered: VecDeque::new(),
                next_from: Some(prefix.to_vec()),
            })
            .collect();
        Self {
            reader,
            prefix: prefix.to_vec(),
            end,
            shards,
        }
    }

    /// The next key, without the prefix, and its value.
    fn next(&mut self) -> Result<Option<StoredRow>, Failure> {
        let limit = NonZeroUsize::new(SCAN_PAGE)
            .ok_or_else(|| Failure::new("scan_page", "a scan page must be nonzero"))?;
        for shard in &mut self.shards {
            if shard.buffered.is_empty()
                && let Some(from) = shard.next_from.take()
            {
                let page = self
                    .reader
                    .range_page_at(shard.shard_id, shard.root, &from, &self.end, limit)
                    .map_err(Failure::observer)?;
                shard.buffered.extend(page.entries);
                shard.next_from = page.next_from;
            }
        }
        let least = self
            .shards
            .iter_mut()
            .filter(|shard| !shard.buffered.is_empty())
            .min_by(|left, right| left.buffered.front().cmp(&right.buffered.front()));
        let Some((key, value)) = least.and_then(|shard| shard.buffered.pop_front()) else {
            return Ok(None);
        };
        let id = key
            .strip_prefix(self.prefix.as_slice())
            .map(<[u8]>::to_vec)
            .ok_or_else(|| {
                Failure::new("scan_key_outside", "a scanned key is outside its prefix")
            })?;
        Ok(Some((id, value)))
    }
}

#[cfg(test)]
#[path = "embed_search_tests.rs"]
mod tests;