haematite_cli/serve/storage/
embed_search.rs

1//! V5 (ruled 5 Oct 08:15, Q5 to Q7): exact search over an embedded field's
2//! vectors in the branch's companion.
3//!
4//! Every vector of the field is compared, so the search reads in proportion
5//! to the field's records, as exact search does. The companion's vectors and
6//! the branch's revisions are both read in id order, a page at a time from
7//! each shard, and joined: no point read for each vector. A vector is
8//! compared only if its record exists at the searched version, the caller may
9//! read it, and it was made from the record's revision there; one made from
10//! any other revision is stale, left out and counted. A record the caller
11//! may not read is never compared, returned or counted (Q6). A vector of any
12//! model but the declared one refuses the search by name: models are never
13//! mixed in one answer (Q7).
14//!
15//! V8a (ruled 5 Oct 09:34, the bound corrected 09:37): a search reads at
16//! most the bound it is handed of the field's vectors, a setting's value
17//! (`query_scan_records`, less what earlier pages read), and refuses
18//! `query_scan_limit` past it, never answering from part of the field. Every
19//! vector row read counts against it, current, stale, hidden or with no
20//! record alike, so the bound is on work; what is hidden is never counted in
21//! anything a caller sees, and the refusal names only the field and the
22//! bound. `Storage::embedding_state` says how far the job's position
23//! has come against the searched version, so an answer tells everything
24//! searched from what is embedded so far; with nothing embedded yet it
25//! refuses `embed_search_not_ready`, never answering empty.
26//!
27//! EMB-E2 (ruled 6 Oct 21:19): a record embedded as windows (`vw\0`) is
28//! read beside the one-vector records, in the same id order. Each window is
29//! a vector read, compared or stale, and counted as one; the record's score
30//! is its best window's, and its hit answers that window's span.
31
32use std::collections::VecDeque;
33use std::num::NonZeroUsize;
34
35use haematite::db::ReadOnlyDatabase;
36use serde::Serialize;
37
38use super::embedded::{self, Declaration};
39use super::{Captured, DocumentRevision, Failure, Storage, reads, version};
40use crate::serve::access::Access;
41use crate::serve::protocol::{Metric, Version};
42
43/// Keys read from one shard at a time.
44const SCAN_PAGE: usize = 64;
45
46/// One hit: the record, its score, and for a record embedded as windows
47/// the span of the best one (EMB-E2).
48#[derive(Debug, Clone, PartialEq)]
49pub struct Hit {
50    pub id: String,
51    pub score: f64,
52    pub span: Option<HitSpan>,
53}
54
55/// EMB-E2: where a hit's best window lies in its record's text, in bytes.
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub struct HitSpan {
58    pub start: u64,
59    pub end: u64,
60}
61
62/// What one exact search answers.
63#[derive(Debug, Clone, PartialEq)]
64pub struct Found {
65    /// Best first; equal scores by record id.
66    pub hits: Vec<Hit>,
67    /// Vectors compared with the question.
68    pub compared: u64,
69    /// Vectors left out because their record's text has changed since.
70    pub stale: u64,
71}
72
73/// What one exact search is over.
74pub struct Searched<'a> {
75    pub collection: &'a str,
76    pub field: &'a str,
77    pub declaration: &'a Declaration,
78    /// The branch at the searched version.
79    pub at: &'a Captured,
80    /// The branch's companion, or `None` before the job's first page.
81    pub companion: Option<&'a Captured>,
82}
83
84/// How far the job has embedded a field against the searched version.
85#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
86#[serde(rename_all = "snake_case")]
87pub enum State {
88    /// Every vector was made from the revision searched: the position's
89    /// branch is the searched branch, its `through` is the searched version,
90    /// and no walk is open.
91    Complete,
92    /// Anything else: a search answers from what is embedded so far.
93    Partial,
94}
95
96/// The job's position against the searched version, as an answer carries it.
97#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
98pub struct Embedding {
99    pub state: State,
100    /// The branch the position names: the searched branch, or the parent it
101    /// was inherited from.
102    pub branch: String,
103    pub through: Option<Version>,
104    /// The open walk's target, if one is open.
105    pub toward: Option<Version>,
106    pub searched: Version,
107}
108
109impl Storage {
110    /// How far the job has embedded `collection.field` against `at`.
111    ///
112    /// # Errors
113    /// `embed_search_not_ready` when `at`'s branch has no companion or the
114    /// field no position: nothing is embedded yet, and a search never answers
115    /// empty for it. Whatever a read answers.
116    pub fn embedding_state(
117        &self,
118        at: &Captured,
119        collection: &str,
120        field: &str,
121    ) -> Result<Embedding, Failure> {
122        let not_ready = || {
123            Failure::new(
124                "embed_search_not_ready",
125                format!(
126                    "nothing of {collection}.{field} is embedded yet on {}",
127                    at.branch
128                ),
129            )
130        };
131        let name = embedded::companion_name(&at.branch);
132        if self.refs.get(&name).is_none() {
133            return Err(not_ready());
134        }
135        let key = embedded::position_key(collection, field);
136        let position = self.position(&name, &key)?.ok_or_else(not_ready)?;
137        let toward = position.toward.map(|open| open.version);
138        let complete = position.branch == at.branch
139            && position.through == Some(at.version)
140            && toward.is_none();
141        let state = if complete {
142            State::Complete
143        } else {
144            State::Partial
145        };
146        Ok(Embedding {
147            state,
148            branch: position.branch,
149            through: position.through,
150            toward,
151            searched: at.version,
152        })
153    }
154
155    /// The declaration of `collection.field` at `at`, and the companion of
156    /// `at`'s branch as it is now.
157    ///
158    /// # Errors
159    /// `embed_field_not_declared`; whatever a read answers.
160    pub fn embedded_field(
161        &self,
162        at: &Captured,
163        collection: &str,
164        field: &str,
165    ) -> Result<(Declaration, Option<Captured>), Failure> {
166        let declaration = self.declared(at, collection, field)?;
167        let name = embedded::companion_name(&at.branch);
168        let companion = self.refs.get(&name).map(|record| Captured {
169            branch: name.clone(),
170            version: version(record),
171            shards: record.shards.clone(),
172        });
173        Ok((declaration, companion))
174    }
175
176    /// Every (collection, field) declared embedded at `at` (EW1: the walks
177    /// the embedding thread starts with).
178    ///
179    /// # Errors
180    /// `vector_declaration_corrupt`; whatever a read answers.
181    pub fn embedded_fields(&self, at: &Captured) -> Result<Vec<(String, String)>, Failure> {
182        let mut fields = Vec::new();
183        reads::scan_rows(&self.reader, at, b"w\0", 2, |parts, bytes| {
184            if let ([collection, field], embedded::Field::Embedded(_)) =
185                (parts, embedded::read_field(bytes)?)
186            {
187                fields.push(((*collection).to_owned(), (*field).to_owned()));
188            }
189            Ok(())
190        })?;
191        Ok(fields)
192    }
193}
194
195/// The best `k` of the field's current vectors for `question`, reading at
196/// most `bound` of the field's vectors.
197///
198/// # Errors
199/// `embed_question_length` for a question of another length than declared;
200/// `query_scan_limit` as soon as one more vector than `bound` would be read;
201/// `embed_model_mixed` for a vector of another model; `vector_corrupt`;
202/// whatever a read answers.
203pub fn search(
204    access: &Access,
205    reader: &ReadOnlyDatabase,
206    searched: &Searched<'_>,
207    question: &[f32],
208    k: NonZeroUsize,
209    bound: u64,
210) -> Result<Found, Failure> {
211    let declaration = searched.declaration;
212    let dimensions = declaration.dimensions.get();
213    if question.len() != dimensions {
214        return Err(Failure::new(
215            "embed_question_length",
216            format!(
217                "the question has {} members; the field is declared with {dimensions}",
218                question.len()
219            ),
220        ));
221    }
222    let mut found = Found {
223        hits: Vec::new(),
224        compared: 0,
225        stale: 0,
226    };
227    let Some(companion) = searched.companion else {
228        return Ok(found);
229    };
230    let judge = Judge {
231        tag: embedded::model_tag(&declaration.model)?,
232        dimensions,
233        question,
234        metric: declaration.distance,
235    };
236    let revisions_prefix = format!("m\0{}\0", searched.collection);
237    let mut items = Items::new(reader, companion, searched, bound);
238    let mut revisions = Ordered::new(reader, searched.at, revisions_prefix.as_bytes());
239    let mut revision: Option<(Vec<u8>, Vec<u8>)> = revisions.next()?;
240    while let Some((id, item)) = items.next()? {
241        // Revisions before this id belong to records with no vector.
242        while revision.as_ref().is_some_and(|(at, _)| *at < id) {
243            revision = revisions.next()?;
244        }
245        let Some(current) = revision
246            .as_ref()
247            .filter(|(at, _)| *at == id)
248            .map(|(_, current)| current)
249        else {
250            continue;
251        };
252        let id = String::from_utf8(id).map_err(|_not_utf8| {
253            Failure::new("embed_record_id_corrupt", "a vector key names no record")
254        })?;
255        if !access.can_read_document(searched.collection, &id) {
256            continue;
257        }
258        let written_at = serde_json::from_slice::<DocumentRevision>(current)
259            .map_err(Failure::storage)?
260            .written_at;
261        let Some((score, span)) = judge.best(&id, &item, written_at, &mut found)? else {
262            continue;
263        };
264        let hit = Hit { id, score, span };
265        let at = found
266            .hits
267            .partition_point(|kept| better(kept, &hit, declaration.distance));
268        if at < k.get() {
269            found.hits.insert(at, hit);
270            found.hits.truncate(k.get());
271        }
272    }
273    Ok(found)
274}
275
276/// What a record's vectors are compared with.
277struct Judge<'a> {
278    tag: embedded::ModelTag,
279    dimensions: usize,
280    question: &'a [f32],
281    metric: Metric,
282}
283
284/// One vector of a record, with its window's span if it is one.
285type Decoded = (embedded::Provenance, Vec<f32>, Option<HitSpan>);
286
287impl Judge<'_> {
288    /// The best score of `item`'s vectors and, for windows, the best one's
289    /// span; `None` when they were made from another revision than
290    /// `written_at`. Every vector counts as compared or stale.
291    ///
292    /// # Errors
293    /// As [`Self::read`].
294    fn best(
295        &self,
296        id: &str,
297        item: &Item,
298        written_at: Version,
299        found: &mut Found,
300    ) -> Result<Option<(f64, Option<HitSpan>)>, Failure> {
301        let vectors = self.read(id, item)?;
302        let rows = u64::try_from(vectors.len()).unwrap_or(u64::MAX);
303        if vectors
304            .iter()
305            .any(|(provenance, ..)| provenance.made_from != written_at)
306        {
307            found.stale = found.stale.saturating_add(rows);
308            return Ok(None);
309        }
310        found.compared = found.compared.saturating_add(rows);
311        let mut best: Option<(f64, Option<HitSpan>)> = None;
312        for (_, vector, span) in &vectors {
313            let score = score(self.question, vector, self.metric);
314            if best.is_none_or(|(kept, _)| ahead(score, kept, self.metric)) {
315                best = Some((score, *span));
316            }
317        }
318        Ok(best)
319    }
320
321    /// `item`'s vectors decoded, a record's windows each in its place.
322    ///
323    /// # Errors
324    /// `embed_model_mixed` for a vector of another model; `vector_corrupt`
325    /// for windows out of place or another count than they name.
326    fn read(&self, id: &str, item: &Item) -> Result<Vec<Decoded>, Failure> {
327        let mut vectors = Vec::new();
328        match item {
329            Item::One(bytes) => {
330                let (provenance, vector) = embedded::decode_embedded(bytes, self.dimensions)?;
331                vectors.push((provenance, vector, None));
332            }
333            Item::Windows(rows) => {
334                for (place, (index, bytes)) in rows.iter().enumerate() {
335                    let (span, provenance, vector) =
336                        embedded::decode_window(bytes, self.dimensions)?;
337                    let placed = usize::try_from(*index).is_ok_and(|at| at == place);
338                    let counted = usize::try_from(span.count).is_ok_and(|n| n == rows.len());
339                    if !placed || !counted {
340                        return Err(Failure::new(
341                            "vector_corrupt",
342                            format!("the windows of {id} are not the ones they name"),
343                        ));
344                    }
345                    let span = HitSpan {
346                        start: span.start,
347                        end: span.end,
348                    };
349                    vectors.push((provenance, vector, Some(span)));
350                }
351            }
352        }
353        if vectors
354            .iter()
355            .any(|(provenance, ..)| provenance.tag != self.tag)
356        {
357            return Err(Failure::new(
358                "embed_model_mixed",
359                format!("the vector of {id} was made by another model than the declared one"),
360            ));
361        }
362        Ok(vectors)
363    }
364}
365
366/// Whether `score` is better than `kept` by `metric`; a tie keeps `kept`.
367fn ahead(score: f64, kept: f64, metric: Metric) -> bool {
368    match metric {
369        Metric::L2 => score.total_cmp(&kept).is_lt(),
370        Metric::Cosine | Metric::Dot => score.total_cmp(&kept).is_gt(),
371    }
372}
373
374/// Whether `kept` comes before `hit`: the better score, then the lower id.
375fn better(kept: &Hit, hit: &Hit, metric: Metric) -> bool {
376    let by_score = match metric {
377        Metric::L2 => kept.score.total_cmp(&hit.score),
378        Metric::Cosine | Metric::Dot => hit.score.total_cmp(&kept.score),
379    };
380    by_score.then_with(|| kept.id.cmp(&hit.id)).is_lt()
381}
382
383/// The measure of `vector` against `question`, as `vectors::score` measures
384/// a bare vector, line for line: the dot product, the distance, or the
385/// cosine (0 when either has no length).
386fn score(question: &[f32], vector: &[f32], metric: Metric) -> f64 {
387    let mut dot = 0.0_f64;
388    let mut question_norm = 0.0_f64;
389    let mut vector_norm = 0.0_f64;
390    let mut squared_distance = 0.0_f64;
391    for (left, right) in question.iter().zip(vector) {
392        let left = f64::from(*left);
393        let right = f64::from(*right);
394        dot = left.mul_add(right, dot);
395        question_norm = left.mul_add(left, question_norm);
396        vector_norm = right.mul_add(right, vector_norm);
397        let difference = left - right;
398        squared_distance = difference.mul_add(difference, squared_distance);
399    }
400    match metric {
401        Metric::Dot => dot,
402        Metric::L2 => squared_distance.sqrt(),
403        Metric::Cosine if question_norm <= 0.0 || vector_norm <= 0.0 => 0.0,
404        Metric::Cosine => dot / (question_norm.sqrt() * vector_norm.sqrt()),
405    }
406}
407
408type StoredRow = (Vec<u8>, Vec<u8>);
409
410/// A record's vectors as the companion holds them: one, or EMB-E2's windows
411/// in key order with their indices.
412enum Item {
413    One(Vec<u8>),
414    Windows(Vec<(u32, Vec<u8>)>),
415}
416
417/// A window row: its record id, its index and its value.
418type WindowRow = (Vec<u8>, u32, Vec<u8>);
419
420/// Every record's vectors of one field in id order: its `v\0` vector and
421/// its `vw\0` windows read side by side, each row read counted against
422/// `bound` and refused past it.
423struct Items<'a> {
424    vectors: Ordered<'a>,
425    windows: Ordered<'a>,
426    vector: Option<StoredRow>,
427    window: Option<WindowRow>,
428    started: bool,
429    read: u64,
430    bound: u64,
431    field: String,
432}
433
434impl<'a> Items<'a> {
435    fn new(
436        reader: &'a ReadOnlyDatabase,
437        companion: &Captured,
438        searched: &Searched<'_>,
439        bound: u64,
440    ) -> Self {
441        let (collection, field) = (searched.collection, searched.field);
442        let vectors = format!("v\0{collection}\0{field}\0");
443        let windows = embedded::window_prefix(collection, field);
444        Self {
445            vectors: Ordered::new(reader, companion, vectors.as_bytes()),
446            windows: Ordered::new(reader, companion, &windows),
447            vector: None,
448            window: None,
449            started: false,
450            read: 0,
451            bound,
452            field: format!("{collection}.{field}"),
453        }
454    }
455
456    /// One more row read: `query_scan_limit` past the bound.
457    fn count(&mut self) -> Result<(), Failure> {
458        if self.read >= self.bound {
459            return Err(Failure::new(
460                "query_scan_limit",
461                format!(
462                    "the search would read more than {} vectors of {}",
463                    self.bound, self.field
464                ),
465            ));
466        }
467        self.read = self.read.saturating_add(1);
468        Ok(())
469    }
470
471    fn pull_vector(&mut self) -> Result<Option<StoredRow>, Failure> {
472        let row = self.vectors.next()?;
473        if row.is_some() {
474            self.count()?;
475        }
476        Ok(row)
477    }
478
479    fn pull_window(&mut self) -> Result<Option<WindowRow>, Failure> {
480        let Some((key, value)) = self.windows.next()? else {
481            return Ok(None);
482        };
483        self.count()?;
484        let (id, index) = embedded::window_of(&[], &key)?;
485        Ok(Some((id.into_bytes(), index, value)))
486    }
487
488    /// The next record's id and vectors.
489    ///
490    /// # Errors
491    /// `query_scan_limit`; `vector_corrupt` for a record with a vector and
492    /// windows both, or more windows than any record keeps; whatever a read
493    /// answers.
494    fn next(&mut self) -> Result<Option<(Vec<u8>, Item)>, Failure> {
495        if !self.started {
496            self.started = true;
497            self.vector = self.pull_vector()?;
498            self.window = self.pull_window()?;
499        }
500        let corrupt = |what: &str| Failure::new("vector_corrupt", format!("a record {what}"));
501        let window_first = match (&self.vector, &self.window) {
502            (None, None) => return Ok(None),
503            (Some((id, _)), Some((window, ..))) if id == window => {
504                return Err(corrupt("holds a vector and windows both"));
505            }
506            (Some((id, _)), Some((window, ..))) => window < id,
507            (Some(_), None) => false,
508            (None, Some(_)) => true,
509        };
510        if !window_first {
511            let next = self.pull_vector()?;
512            let one = std::mem::replace(&mut self.vector, next);
513            return Ok(one.map(|(id, bytes)| (id, Item::One(bytes))));
514        }
515        let Some((id, index, bytes)) = self.window.take() else {
516            return Ok(None);
517        };
518        let mut rows = vec![(index, bytes)];
519        self.window = loop {
520            match self.pull_window()? {
521                Some((next, index, bytes)) if next == id => {
522                    let held = u32::try_from(rows.len()).unwrap_or(u32::MAX);
523                    if held >= embedded::WINDOWS_MOST {
524                        return Err(corrupt("holds more windows than any record keeps"));
525                    }
526                    rows.push((index, bytes));
527                }
528                other => break other,
529            }
530        };
531        Ok(Some((id, Item::Windows(rows))))
532    }
533}
534
535/// One shard's place in an ordered read: its buffered page and where the
536/// next page starts.
537struct ShardScan {
538    shard_id: usize,
539    root: haematite::tree::Hash,
540    buffered: VecDeque<StoredRow>,
541    next_from: Option<Vec<u8>>,
542}
543
544/// Every key under a prefix at one version, across its shards, in key order,
545/// with the prefix taken off: one page from each shard at a time.
546struct Ordered<'a> {
547    reader: &'a ReadOnlyDatabase,
548    prefix: Vec<u8>,
549    end: Vec<u8>,
550    shards: Vec<ShardScan>,
551}
552
553impl<'a> Ordered<'a> {
554    fn new(reader: &'a ReadOnlyDatabase, at: &Captured, prefix: &[u8]) -> Self {
555        let mut end = prefix.to_vec();
556        // Every prefix here ends in a NUL separator, so one past it is the
557        // first key beyond every key under it.
558        if let Some(last) = end.last_mut() {
559            *last = last.saturating_add(1);
560        }
561        let shards = at
562            .shards
563            .iter()
564            .map(|shard| ShardScan {
565                shard_id: shard.shard_id,
566                root: shard.head,
567                buffered: VecDeque::new(),
568                next_from: Some(prefix.to_vec()),
569            })
570            .collect();
571        Self {
572            reader,
573            prefix: prefix.to_vec(),
574            end,
575            shards,
576        }
577    }
578
579    /// The next key, without the prefix, and its value.
580    fn next(&mut self) -> Result<Option<StoredRow>, Failure> {
581        let limit = NonZeroUsize::new(SCAN_PAGE)
582            .ok_or_else(|| Failure::new("scan_page", "a scan page must be nonzero"))?;
583        for shard in &mut self.shards {
584            if shard.buffered.is_empty()
585                && let Some(from) = shard.next_from.take()
586            {
587                let page = self
588                    .reader
589                    .range_page_at(shard.shard_id, shard.root, &from, &self.end, limit)
590                    .map_err(Failure::observer)?;
591                shard.buffered.extend(page.entries);
592                shard.next_from = page.next_from;
593            }
594        }
595        let least = self
596            .shards
597            .iter_mut()
598            .filter(|shard| !shard.buffered.is_empty())
599            .min_by(|left, right| left.buffered.front().cmp(&right.buffered.front()));
600        let Some((key, value)) = least.and_then(|shard| shard.buffered.pop_front()) else {
601            return Ok(None);
602        };
603        let id = key
604            .strip_prefix(self.prefix.as_slice())
605            .map(<[u8]>::to_vec)
606            .ok_or_else(|| {
607                Failure::new("scan_key_outside", "a scanned key is outside its prefix")
608            })?;
609        Ok(Some((id, value)))
610    }
611}
612
613#[cfg(test)]
614#[path = "embed_search_tests.rs"]
615mod tests;
haematite_cli/serve/storage/embed_search.rs

haematite_cli/serve/storage/
embed_search.rs
