haematite-cli 0.21.8

//! A search by a stored record or words embedded by the service: nearest
//! vectors first, each record read through the request's access and
//! carrying its score. Job 2 of Part 8 of docs/design/product/vector-index.md.
//!
//! A scan that reaches a page's records-read limit stops there and gives
//! `next`; a search cannot. The search is exact: which hit is best is known
//! only once every vector has been compared, so there is no place part way
//! through to stop at, and a search that must compare more vectors than the
//! limit is refused `query_scan_limit`.

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

use haematite::db::ReadOnlyDatabase;
use haematite_query::handle::Error;
use haematite_query::page::{Origin, Record, RecordSource};

use super::parser::{self, Question};
use super::shaped_source::ShapedSource;
use super::source::AuthorizedSource;
use crate::serve::access::Access;
use crate::serve::protocol::{Metric, Selection, VectorSearch};
use crate::serve::service::{Service, Words, not_configured};
use crate::serve::storage::vectors::{self, Hit};
use crate::serve::storage::{Captured, Failure};

const RECORD_MISSING: &str = "query_search_record_missing";

/// A scored source, retaining either kind of parsed question.
pub struct Search {
    pub field: String,
    pub question: Question,
    pub measure: Metric,
    pub limit: NonZeroU32,
}

impl From<parser::Search> for Search {
    fn from(search: parser::Search) -> Self {
        Self {
            field: search.field,
            question: Question::Record(search.record),
            measure: search.measure,
            limit: search.limit,
        }
    }
}

/// The caller's read view and, when served, the service's embedding worker.
pub struct Read<'a> {
    pub access: &'a Access,
    pub reader: &'a ReadOnlyDatabase,
    pub service: Option<&'a Service>,
}

/// Where a page's records come from.
pub enum Records<'a> {
    /// The documents of the query's scope, in (collection, id) order.
    Scan(AuthorizedSource<'a>),
    /// Root order with each collection's own arguments and projection.
    Shaped(Box<ShapedSource<'a>>),
    /// The hits of a search, best first.
    Search(Hits<'a>),
}

impl Records<'_> {
    /// How many vectors this page's search compared; `None` without a search.
    pub const fn compared(&self) -> Option<u64> {
        match self {
            Self::Scan(_) | Self::Shaped(_) => None,
            Self::Search(hits) => Some(hits.compared),
        }
    }
}

impl RecordSource for Records<'_> {
    fn next_record(&mut self) -> Result<Option<Record>, Error> {
        match self {
            Self::Scan(source) => source.next_record(),
            Self::Shaped(source) => source.next_record(),
            Self::Search(hits) => hits.next_record(),
        }
    }

    fn passed(&self) -> Option<Origin> {
        match self {
            Self::Scan(source) => source.passed(),
            Self::Shaped(source) => source.passed(),
            Self::Search(_) => None,
        }
    }
}

/// A search's hits, each read as its record when it is taken.
pub struct Hits<'a> {
    source: AuthorizedSource<'a>,
    collection: String,
    remaining: std::vec::IntoIter<Hit>,
    compared: u64,
}

/// How far one page's search may go, and where its hits start.
pub struct Reach {
    /// The vectors this page may compare: its records-read limit less the
    /// records already read on earlier pages.
    pub bound: u64,
    /// The hits earlier pages gave in full, skipped without being read.
    pub skip: u64,
}

impl RecordSource for Hits<'_> {
    fn next_record(&mut self) -> Result<Option<Record>, Error> {
        for hit in self.remaining.by_ref() {
            let found = self
                .source
                .found(&self.collection, &hit.document_id, hit.score)?;
            if found.is_some() {
                return Ok(found);
            }
        }
        Ok(None)
    }
}

/// The hits of `search` over `collection` at `at`, read through `source`,
/// from the first hit `reach` does not skip.
///
/// A record question reads its stored vector through the caller's access.
/// A text question uses the service's question worker and embedded search.
/// Both give the best hits first, with ties by record id and caller access
/// applied. The same search at the same moment gives the same hits, so a later
/// page skips exactly the hits an earlier page gave.
///
/// # Errors
///
/// `query_search_record_missing` when no vector of the record in the field is
/// there for the request to read: an absent record and one the request may
/// not read are refused with the same words. `query_scan_limit` as soon as
/// the search would compare more vectors than `reach` allows. Any other
/// refusal of the search, including the existing `embed_*` refusals for
/// words and `query_measure_mismatch` when the measure differs from the field.
pub fn hits<'a>(
    source: AuthorizedSource<'a>,
    read: &Read<'_>,
    at: &Captured,
    collection: &str,
    search: &Search,
    reach: &Reach,
) -> Result<Hits<'a>, Failure> {
    let k = NonZeroUsize::try_from(search.limit)
        .map_err(|_| Failure::new("page_limit", "the limit does not fit this machine"))?;
    let (mut found, compared) = match &search.question {
        Question::Record(record) => {
            let question = vectors::get(
                read.access,
                read.reader,
                at,
                collection,
                &search.field,
                record,
            )?
            .ok_or_else(|| missing(&search.field))?;
            let request = VectorSearch {
                selection: Selection {
                    branch: at.branch.clone(),
                    version: Some(at.version),
                },
                collection: collection.to_owned(),
                field: search.field.clone(),
                query: question,
                k: search.limit,
                metric: search.measure,
            };
            vectors::search_within(read.access, read.reader, at, &request, k, Some(reach.bound))?
        }
        Question::Text(text) => {
            let service = read.service.ok_or_else(not_configured)?;
            let (found, _) = service.search_words(
                read.access,
                at,
                &Words {
                    collection,
                    field: &search.field,
                    text,
                    measure: Some(search.measure),
                    k,
                    bound: reach.bound,
                },
            )?;
            let hits = found
                .hits
                .into_iter()
                .map(|hit| Hit {
                    document_id: hit.id,
                    score: hit.score,
                })
                .collect();
            (hits, found.compared)
        }
    };
    let given = usize::try_from(reach.skip).map_or(found.len(), |given| given.min(found.len()));
    found.drain(..given);
    Ok(Hits {
        source,
        collection: collection.to_owned(),
        remaining: found.into_iter(),
        compared,
    })
}

/// One refusal for an absent record and for one the request may not read.
pub(super) fn missing(field: &str) -> Failure {
    Failure::new(
        RECORD_MISSING,
        format!("the record named has no vector in .{field} that this request may read"),
    )
}

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

#[cfg(all(test, unix))]
#[path = "search_text_tests.rs"]
mod text_tests;