use haematite::db::ReadOnlyDatabase;
use crate::serve::access::Access;
use crate::serve::storage::embedded::{self, Declaration};
use crate::serve::storage::{Captured, DocumentRevision, Failure, reads, revision_key, vectors};
pub(super) struct Budget {
pub remaining: u64,
pub limit: u64,
pub read: u64,
pub name: &'static str,
pub memory: Option<u64>,
}
impl Budget {
fn keep(&mut self, capacity: usize) -> Result<(), Failure> {
if let Some(remaining) = self.memory {
let bytes = capacity
.checked_mul(std::mem::size_of::<f32>())
.and_then(|bytes| bytes.checked_add(std::mem::size_of::<Vec<f32>>()))
.and_then(|bytes| u64::try_from(bytes).ok());
self.memory = Some(
bytes
.and_then(|bytes| remaining.checked_sub(bytes))
.ok_or_else(|| {
Failure::new(
"query_memory_limit",
format!(
"{} candidate vectors exceed query memory allowance",
self.name
),
)
})?,
);
}
Ok(())
}
fn charge(&mut self) -> Result<(), Failure> {
if self.read >= self.remaining {
return Err(Failure::new(
"query_scan_limit",
format!(
"{} records and vectors exceed scan limit {}",
self.name, self.limit
),
));
}
self.read += 1;
Ok(())
}
}
pub(super) struct Read<'a> {
pub access: &'a Access,
pub reader: &'a ReadOnlyDatabase,
pub at: &'a Captured,
pub companion: Option<&'a Captured>,
pub collection: &'a str,
pub field: &'a str,
pub declaration: Option<&'a Declaration>,
}
impl Read<'_> {
pub fn vectors(&self, id: &str, budget: &mut Budget) -> Result<Vec<Vec<f32>>, Failure> {
if !self.access.can_read_document(self.collection, id) {
return Ok(Vec::new());
}
let Some(declaration) = self.declaration else {
let vector = vectors::get(
self.access,
self.reader,
self.at,
self.collection,
self.field,
id,
)?;
return match vector {
Some(vector) => {
budget.charge()?;
budget.keep(vector.capacity())?;
Ok(vec![vector])
}
None => Ok(Vec::new()),
};
};
let Some(companion) = self.companion else {
return Ok(Vec::new());
};
let revision = reads::get_bytes(self.reader, self.at, &revision_key(self.collection, id))?;
let Some(revision) = revision else {
return Ok(Vec::new());
};
let revision: DocumentRevision =
serde_json::from_slice(&revision).map_err(Failure::storage)?;
let key = embedded::vector_key(self.collection, self.field, id);
let one = reads::get_bytes(self.reader, companion, &key)?;
let first_key = embedded::window_key(self.collection, self.field, id, 0);
let first = reads::get_bytes(self.reader, companion, &first_key)?;
let tag = embedded::model_tag(&declaration.model)?;
let dimensions = declaration.dimensions.get();
match (one, first) {
(Some(_), Some(_)) => Err(corrupt(
"a rerank candidate holds a vector and windows both",
)),
(Some(bytes), None) => {
budget.charge()?;
let (provenance, vector) = embedded::decode_embedded(&bytes, dimensions)?;
if provenance.tag != tag {
return Err(mixed());
}
if provenance.made_from == revision.written_at {
budget.keep(vector.capacity())?;
Ok(vec![vector])
} else {
Ok(Vec::new())
}
}
(None, Some(bytes)) => {
self.windows(id, companion, declaration, &revision, bytes, budget)
}
(None, None) => Ok(Vec::new()),
}
}
fn windows(
&self,
id: &str,
companion: &Captured,
declaration: &Declaration,
revision: &DocumentRevision,
first: Vec<u8>,
budget: &mut Budget,
) -> Result<Vec<Vec<f32>>, Failure> {
let count = embedded::window_span(&first)?.count;
let tag = embedded::model_tag(&declaration.model)?;
let mut vectors = Vec::new();
let mut stale = false;
let mut bytes = first;
for index in 0..count {
if index != 0 {
let key = embedded::window_key(self.collection, self.field, id, index);
bytes = reads::get_bytes(self.reader, companion, &key)?
.ok_or_else(|| corrupt("a rerank candidate is missing one of its windows"))?;
}
budget.charge()?;
let (span, provenance, vector) =
embedded::decode_window(&bytes, declaration.dimensions.get())?;
if span.count != count {
return Err(corrupt("rerank windows disagree about their count"));
}
if provenance.tag != tag {
return Err(mixed());
}
stale |= provenance.made_from != revision.written_at;
budget.keep(vector.capacity())?;
vectors
.try_reserve_exact(1)
.map_err(|error| Failure::new("query_memory_limit", error.to_string()))?;
vectors.push(vector);
}
let extra = embedded::window_key(self.collection, self.field, id, count);
if reads::get_bytes(self.reader, companion, &extra)?.is_some() {
return Err(corrupt(
"a rerank candidate holds more windows than it names",
));
}
if stale {
vectors.clear();
}
Ok(vectors)
}
}
fn corrupt(message: &str) -> Failure {
Failure::new("vector_corrupt", message)
}
fn mixed() -> Failure {
Failure::new(
"embed_model_mixed",
"a rerank vector was made by another model than declared",
)
}