haematite_cli/serve/storage/
embedded.rs

1//! Stored embedding declarations, companion headers, model tags and keys.
2//!
3//! A `w\0` value in decimal is a field as it was before declarations: its
4//! vectors are supplied by the caller and stored bare on the branch, with that
5//! length. A field declared embedded has its vectors made by the job and kept
6//! in the branch's companion namespace, never on the branch (Q1); each value
7//! there is headed by the version of the record it was made from and the tag
8//! of the model that made it.
9
10use std::num::NonZeroUsize;
11
12use haematite::tree::Hash;
13use serde::{Deserialize, Serialize};
14
15use super::Failure;
16use super::vectors::{decode, parse_dimension};
17use crate::serve::protocol::{Metric, Version};
18
19/// Each branch's companion is the namespace `@haem/embed/{branch}`.
20pub const COMPANION_PREFIX: &str = "@haem/embed/";
21/// The only header format: format, generation, sequence, tag.
22pub const FORMAT: u8 = 1;
23pub const TAG_BYTES: usize = 16;
24pub const HEADER_BYTES: usize = 1 + 8 + 8 + TAG_BYTES;
25
26/// A model as a declaration names it, and as its worker's Ready names it.
27#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
28#[serde(deny_unknown_fields)]
29pub struct ModelName {
30    pub name: String,
31    pub revision: String,
32    pub manifest_sha256: String,
33}
34
35/// The kinds a field may be declared as; dense first (design item "Dense
36/// vectors first").
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
38#[serde(rename_all = "snake_case")]
39pub enum VectorKind {
40    Dense,
41}
42
43/// A field declared embedded: which text field its vectors are made from,
44/// by which model, of what length, and the measure a search uses.
45#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
46#[serde(deny_unknown_fields)]
47pub struct Declaration {
48    pub kind: VectorKind,
49    pub source: String,
50    pub model: ModelName,
51    pub dimensions: NonZeroUsize,
52    pub distance: Metric,
53}
54
55/// What a field's `w\0` value says it is.
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub enum Field {
58    /// Written before declarations, or supplied by the caller: bare vectors
59    /// on the branch, of this length.
60    Supplied { dimensions: usize },
61    /// Made by the job and kept in the companion.
62    Embedded(Declaration),
63}
64
65/// The field a `w\0` value describes: decimal is supplied, as every value
66/// written before declarations is; anything else is a declaration or
67/// refused by name.
68///
69/// # Errors
70/// `vector_dimension_corrupt` for decimal that is not a canonical positive
71/// length; `vector_declaration_corrupt` for a declaration that does not
72/// decode.
73pub fn read_field(bytes: &[u8]) -> Result<Field, Failure> {
74    if bytes.first().is_some_and(u8::is_ascii_digit) {
75        return parse_dimension(bytes).map(|dimensions| Field::Supplied { dimensions });
76    }
77    serde_json::from_slice(bytes)
78        .map(Field::Embedded)
79        .map_err(|error| Failure::new("vector_declaration_corrupt", error.to_string()))
80}
81
82/// The `w\0` value of `declaration`.
83///
84/// # Errors
85/// Whatever encoding it answers.
86pub fn encode_declaration(declaration: &Declaration) -> Result<Vec<u8>, Failure> {
87    serde_json::to_vec(declaration).map_err(Failure::storage)
88}
89
90/// The first 16 bytes of the store's hash of the model's name, revision and
91/// manifest, as JSON so no two models share the bytes hashed.
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub struct ModelTag(pub [u8; TAG_BYTES]);
94
95/// The tag of `model`.
96///
97/// # Errors
98/// Whatever encoding the name answers.
99pub fn model_tag(model: &ModelName) -> Result<ModelTag, Failure> {
100    let named = serde_json::to_vec(model).map_err(Failure::storage)?;
101    let hash = Hash::of(&named);
102    let tag = hash
103        .as_bytes()
104        .first_chunk::<TAG_BYTES>()
105        .ok_or_else(|| Failure::new("embed_tag_short", "the store's hash is under 16 bytes"))?;
106    Ok(ModelTag(*tag))
107}
108
109/// What an embedded vector was made from.
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
111pub struct Provenance {
112    /// The record's revision (`m\0` `written_at`) its text was read at.
113    pub made_from: Version,
114    pub tag: ModelTag,
115}
116
117/// A companion value: the header, then the members as little-endian f32s,
118/// as a bare vector stores them.
119#[must_use]
120pub fn encode_embedded(provenance: &Provenance, vector: &[f32]) -> Vec<u8> {
121    let capacity = vector.len().saturating_mul(4).saturating_add(HEADER_BYTES);
122    let mut bytes = Vec::with_capacity(capacity);
123    bytes.push(FORMAT);
124    bytes.extend_from_slice(&provenance.made_from.generation.to_be_bytes());
125    bytes.extend_from_slice(&provenance.made_from.sequence.to_be_bytes());
126    bytes.extend_from_slice(&provenance.tag.0);
127    for member in vector {
128        bytes.extend_from_slice(&member.to_le_bytes());
129    }
130    bytes
131}
132
133/// A companion value's provenance and vector of `dimensions` members.
134///
135/// # Errors
136/// `vector_corrupt` for a value shorter than its header, of another format,
137/// or whose members are not `dimensions` finite numbers.
138pub fn decode_embedded(bytes: &[u8], dimensions: usize) -> Result<(Provenance, Vec<f32>), Failure> {
139    let short = || {
140        Failure::new(
141            "vector_corrupt",
142            "an embedded vector is shorter than its header",
143        )
144    };
145    let (&format, rest) = bytes.split_first().ok_or_else(short)?;
146    if format != FORMAT {
147        return Err(Failure::new(
148            "vector_corrupt",
149            format!("an embedded vector has format {format}, not {FORMAT}"),
150        ));
151    }
152    let (generation, rest) = rest.split_first_chunk::<8>().ok_or_else(short)?;
153    let (sequence, rest) = rest.split_first_chunk::<8>().ok_or_else(short)?;
154    let (tag, members) = rest.split_first_chunk::<TAG_BYTES>().ok_or_else(short)?;
155    let provenance = Provenance {
156        made_from: Version {
157            generation: u64::from_be_bytes(*generation),
158            sequence: u64::from_be_bytes(*sequence),
159        },
160        tag: ModelTag(*tag),
161    };
162    Ok((provenance, decode(members, dimensions)?))
163}
164
165/// The job's field position, committed with its vectors in the companion.
166///
167/// Naming the branch lets a fork resume from its parent's position.
168#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
169#[serde(deny_unknown_fields)]
170pub struct Position {
171    pub branch: String,
172    /// Every change to `branch` up to and including this version has its
173    /// vector or its removal; `None` before the first walk ends, when the
174    /// walk starts from the empty tree.
175    pub through: Option<Version>,
176    /// The walk under way, if one is.
177    pub toward: Option<Toward>,
178}
179
180/// A walk of the difference from a position's `through` to `version`, in key
181/// order: every record up to and including `after` is finished.
182#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
183#[serde(deny_unknown_fields)]
184pub struct Toward {
185    pub version: Version,
186    /// A record id: the walk's order is its keys' order, `m\0` then the
187    /// collection then the id, so ids compare as bytes.
188    pub after: String,
189}
190
191/// The position a companion value holds.
192///
193/// # Errors
194/// `embed_position_corrupt` for a value that does not decode.
195pub fn read_position(bytes: &[u8]) -> Result<Position, Failure> {
196    serde_json::from_slice(bytes)
197        .map_err(|error| Failure::new("embed_position_corrupt", error.to_string()))
198}
199
200/// The companion value of `position`.
201///
202/// # Errors
203/// Whatever encoding it answers.
204pub fn encode_position(position: &Position) -> Result<Vec<u8>, Failure> {
205    serde_json::to_vec(position).map_err(Failure::storage)
206}
207
208/// The companion namespace of `branch`.
209#[must_use]
210pub fn companion_name(branch: &str) -> String {
211    format!("{COMPANION_PREFIX}{branch}")
212}
213
214/// The companion key of a field's position.
215#[must_use]
216pub fn position_key(collection: &str, field: &str) -> Vec<u8> {
217    format!("e\0{collection}\0{field}").into_bytes()
218}
219
220/// The companion key of a record's vector, laid out as a bare vector's key on
221/// the branch is.
222#[must_use]
223pub fn vector_key(collection: &str, field: &str, id: &str) -> Vec<u8> {
224    format!("v\0{collection}\0{field}\0{id}").into_bytes()
225}
226
227/// EMB-E2 (j): the companion key that marks a field's open walk a rewalk.
228///
229/// A rewalk reads from the empty tree and lists only texts with no vector.
230/// The key is not under `e\0`, so a scan of positions never reads it.
231#[must_use]
232pub fn rewalk_key(collection: &str, field: &str) -> Vec<u8> {
233    format!("er\0{collection}\0{field}").into_bytes()
234}
235
236/// EMB-E2: the most windows a record may hold, whatever the setting; a stored
237/// count past it is corrupt, so clearing a record's windows stays bounded.
238pub const WINDOWS_MOST: u32 = 1024;
239
240/// EMB-E2: the companion prefix of every window of a field.
241#[must_use]
242pub fn window_prefix(collection: &str, field: &str) -> Vec<u8> {
243    format!("vw\0{collection}\0{field}\0").into_bytes()
244}
245
246/// EMB-E2: the companion key of window `index` of a record's text.
247///
248/// The key is the field's window prefix, the id, a NUL and the index
249/// big-endian, so a record's windows sort together and in text order. Ids
250/// hold no NUL.
251#[must_use]
252pub fn window_key(collection: &str, field: &str, id: &str, index: u32) -> Vec<u8> {
253    let mut key = window_prefix(collection, field);
254    key.extend_from_slice(id.as_bytes());
255    key.push(0);
256    key.extend_from_slice(&index.to_be_bytes());
257    key
258}
259
260/// EMB-E2: the record id and window index of a key under `prefix`.
261///
262/// # Errors
263/// `vector_corrupt` for a key that is not a window key under `prefix`.
264pub fn window_of(prefix: &[u8], key: &[u8]) -> Result<(String, u32), Failure> {
265    let corrupt = || Failure::new("vector_corrupt", "a window key names no record and index");
266    let rest = key.strip_prefix(prefix).ok_or_else(corrupt)?;
267    let (named, index) = rest.split_last_chunk::<4>().ok_or_else(corrupt)?;
268    let id = named.strip_suffix(&[0]).ok_or_else(corrupt)?;
269    let id = String::from_utf8(id.to_vec()).map_err(|_not_utf8| corrupt())?;
270    Ok((id, u32::from_be_bytes(*index)))
271}
272
273/// EMB-E2: where a window lies in its record's text, `[start, end)` in
274/// bytes, and how many windows the record has.
275#[derive(Debug, Clone, Copy, PartialEq, Eq)]
276pub struct Span {
277    pub count: u32,
278    pub start: u64,
279    pub end: u64,
280}
281
282const SPAN_BYTES: usize = 20;
283
284/// EMB-E2: a window's value: its [`Span`], then its provenance and vector as
285/// [`encode_embedded`] lays them out.
286#[must_use]
287pub fn encode_window(span: Span, provenance: &Provenance, vector: &[f32]) -> Vec<u8> {
288    let embedded = encode_embedded(provenance, vector);
289    let mut bytes = Vec::with_capacity(SPAN_BYTES.saturating_add(embedded.len()));
290    bytes.extend_from_slice(&span.count.to_be_bytes());
291    bytes.extend_from_slice(&span.start.to_be_bytes());
292    bytes.extend_from_slice(&span.end.to_be_bytes());
293    bytes.extend_from_slice(&embedded);
294    bytes
295}
296
297/// EMB-E2: a window's span, provenance and vector of `dimensions` members.
298///
299/// # Errors
300/// `vector_corrupt` for a value shorter than its span, a count past
301/// [`WINDOWS_MOST`], or as [`decode_embedded`].
302pub fn decode_window(
303    bytes: &[u8],
304    dimensions: usize,
305) -> Result<(Span, Provenance, Vec<f32>), Failure> {
306    let (span, rest) = split_span(bytes)?;
307    let (provenance, vector) = decode_embedded(rest, dimensions)?;
308    Ok((span, provenance, vector))
309}
310
311/// EMB-E2: a window's span alone, without decoding its vector.
312///
313/// # Errors
314/// As [`decode_window`].
315pub fn window_span(bytes: &[u8]) -> Result<Span, Failure> {
316    split_span(bytes).map(|(span, _)| span)
317}
318
319fn split_span(bytes: &[u8]) -> Result<(Span, &[u8]), Failure> {
320    let short = || {
321        Failure::new(
322            "vector_corrupt",
323            "an embedded window is shorter than its span",
324        )
325    };
326    let (count, rest) = bytes.split_first_chunk::<4>().ok_or_else(short)?;
327    let (start, rest) = rest.split_first_chunk::<8>().ok_or_else(short)?;
328    let (end, rest) = rest.split_first_chunk::<8>().ok_or_else(short)?;
329    let span = Span {
330        count: u32::from_be_bytes(*count),
331        start: u64::from_be_bytes(*start),
332        end: u64::from_be_bytes(*end),
333    };
334    if span.count > WINDOWS_MOST || span.end <= span.start {
335        return Err(Failure::new(
336            "vector_corrupt",
337            "an embedded window names no span or too many windows",
338        ));
339    }
340    Ok((span, rest))
341}
342
343#[cfg(test)]
344#[path = "embedded_tests.rs"]
345mod tests;
haematite_cli/serve/storage/embedded.rs

haematite_cli/serve/storage/
embedded.rs
