use std::collections::VecDeque;
use std::num::NonZeroU64;
use haematite::db::ReadOnlyDatabase;
use haematite_query::Refusal;
use haematite_query::handle::Error;
use haematite_query::page::Record;
use haematite_query::stage::rerank::{Measure, order, score};
use super::super::parser::Question;
use super::rerank_vectors::{Budget, Read};
use super::{Context, refused};
use crate::serve::access::Access;
use crate::serve::contracts::identifier;
use crate::serve::embed::question::question_for;
use crate::serve::protocol::Metric;
use crate::serve::service::{Service, not_configured};
use crate::serve::storage::embedded::{self, Declaration, Field};
use crate::serve::storage::{Captured, Failure, Held, reads};
pub(super) struct Request {
pub field: String,
pub question: Question,
pub measure: Metric,
pub collection: String,
pub name: &'static str,
pub memory: Option<NonZeroU64>,
}
struct Judge {
field: String,
question: Vec<f32>,
measure: Metric,
declaration: Option<Declaration>,
collection: String,
name: &'static str,
}
pub(super) struct Prepared<'a> {
access: &'a Access,
reader: &'a ReadOnlyDatabase,
at: &'a Captured,
collection: String,
companion: Option<Held>,
judges: Vec<Judge>,
limit: u64,
pub unscored: Option<u64>,
failure: Option<Failure>,
}
impl<'a> Prepared<'a> {
pub fn new(
context: &Context<'a>,
at: &'a Captured,
requests: &[Request],
service: Option<&Service>,
limit: u64,
) -> Result<Self, Failure> {
let mut prepared = Self {
access: context.access,
reader: context.reader,
at,
collection: String::new(),
companion: None,
judges: Vec::new(),
limit,
unscored: Some(0),
failure: None,
};
for request in requests {
prepared.collection.clone_from(&request.collection);
let judge = prepared.judge(context, request, service)?;
prepared.judges.push(judge);
if let Some(memory) = request.memory {
prepared.questions_memory(memory.get())?;
}
}
Ok(prepared)
}
fn judge(
&mut self,
context: &Context<'_>,
request: &Request,
service: Option<&Service>,
) -> Result<Judge, Failure> {
identifier(&request.field)?;
let key = format!("w\0{}\0{}", self.collection, request.field);
let field = reads::get_bytes(self.reader, self.at, key.as_bytes())?
.map(|bytes| embedded::read_field(&bytes))
.transpose()?;
let declaration = match field {
Some(Field::Embedded(declaration)) => {
if declaration.distance != request.measure {
return Err(Failure::new(
"query_measure_mismatch",
format!(
"{} names {} for .{}; the field declares {}",
request.name,
metric(request.measure),
request.field,
metric(declaration.distance),
),
));
}
if self.companion.is_none() {
self.companion = context.catalogue.capture_companion_held(&self.at.branch)?;
}
Some(declaration)
}
Some(Field::Supplied { .. }) | None => None,
};
let question = match &request.question {
Question::Text(text) => {
let service = service.ok_or_else(not_configured)?;
let asker = service.embedding().ok_or_else(not_configured)?.asker()?;
let declaration = declaration.as_ref().ok_or_else(|| {
Failure::new(
"embed_field_not_declared",
format!("{} words require a declared embedded field", request.name),
)
})?;
question_for(asker, text, declaration)?
}
Question::Record(id) => self.record_question(request, id, declaration.as_ref())?,
};
if question.is_empty() || question.iter().any(|member| !member.is_finite()) {
return Err(Failure::new(
"vector_non_finite",
format!(
"the {} question must be a nonempty finite vector",
request.name
),
));
}
Ok(Judge {
field: request.field.clone(),
question,
measure: request.measure,
declaration,
collection: request.collection.clone(),
name: request.name,
})
}
fn record_question(
&self,
request: &Request,
id: &str,
declaration: Option<&Declaration>,
) -> Result<Vec<f32>, Failure> {
identifier(id)?;
let read = Read {
access: self.access,
reader: self.reader,
at: self.at,
companion: self.companion.as_deref(),
collection: &self.collection,
field: &request.field,
declaration,
};
let mut budget = Budget {
remaining: u64::MAX,
limit: self.limit,
read: 0,
name: request.name,
memory: request.memory.map(NonZeroU64::get),
};
let mut vectors = read.vectors(id, &mut budget)?.into_iter();
match (vectors.next(), vectors.next()) {
(Some(vector), None) => Ok(vector),
_ if request.name == "similar" => Err(super::super::search::missing(&request.field)),
_ => Err(Failure::new(
"query_search_record_missing",
format!(
"the record named has no single vector in .{} that this request may read",
request.field,
),
)),
}
}
pub fn rank(
&mut self,
index: usize,
records: Vec<Record>,
remaining: u64,
) -> Result<(VecDeque<Record>, u64), Error> {
self.rank_with_memory(index, records, remaining, None)
}
pub fn rank_with_memory(
&mut self,
index: usize,
records: Vec<Record>,
remaining: u64,
memory: Option<u64>,
) -> Result<(VecDeque<Record>, u64), Error> {
match self.ranked(index, records, remaining, memory) {
Ok(result) => Ok(result),
Err(failure) => {
let code = match failure.code.as_str() {
"query_scan_limit" => Refusal::QueryScanLimit,
"query_memory_limit" => Refusal::QueryMemoryLimit,
_ => Refusal::QueryReadFailed,
};
let error = Error::new(code, format!("{}: {}", failure.code, failure.message));
self.failure = Some(failure);
Err(error)
}
}
}
fn ranked(
&mut self,
index: usize,
records: Vec<Record>,
remaining: u64,
memory: Option<u64>,
) -> Result<(VecDeque<Record>, u64), Failure> {
let memory = memory
.map(|memory| self.questions_memory(memory))
.transpose()?;
let judge = self.judges.get(index).ok_or_else(|| {
Failure::new("query_protocol", "rerank stage has no prepared question")
})?;
let read = Read {
access: self.access,
reader: self.reader,
at: self.at,
companion: self.companion.as_deref(),
collection: &judge.collection,
field: &judge.field,
declaration: judge.declaration.as_ref(),
};
let mut budget = Budget {
remaining,
limit: self.limit,
read: 0,
name: judge.name,
memory,
};
if judge.name == "similar" {
for record in &records {
if !record.value.get().trim_start().starts_with('{') {
return Err(super::super::search::missing(&judge.field));
}
}
}
let descending = judge.measure != Metric::L2;
let mut failure = None;
let result = order(records, descending, |origin| {
budget.memory = memory;
if origin.collection != judge.collection {
let error = Failure::new(
"query_protocol",
format!("{} candidate belongs to a different collection", judge.name),
);
let refusal = Error::new(Refusal::QueryReadFailed, &error.message);
failure = Some(error);
return Err(refusal);
}
let vectors = match read.vectors(&origin.id, &mut budget) {
Ok(vectors) => vectors,
Err(error) => {
let refusal = Error::new(Refusal::QueryReadFailed, &error.message);
failure = Some(error);
return Err(refusal);
}
};
let mut best: Option<f64> = None;
for vector in vectors {
let score = score(&judge.question, &vector, measure(judge.measure))?;
if best.is_none_or(|kept| {
if descending {
score > kept
} else {
score < kept
}
}) {
best = Some(score);
}
}
Ok(best)
});
let (records, unscored) =
result.map_err(|error| failure.unwrap_or_else(|| refused(&error)))?;
self.unscored = Some(unscored);
Ok((records, budget.read))
}
fn questions_memory(&self, limit: u64) -> Result<u64, Failure> {
self.judges.iter().try_fold(limit, |remaining, judge| {
let bytes = judge
.question
.capacity()
.checked_mul(std::mem::size_of::<f32>())
.and_then(|bytes| bytes.checked_add(std::mem::size_of::<Judge>()))
.and_then(|bytes| u64::try_from(bytes).ok())
.ok_or_else(|| memory_failure(limit))?;
remaining
.checked_sub(bytes)
.ok_or_else(|| memory_failure(limit))
})
}
pub fn failure(&mut self, error: &Error) -> Failure {
self.failure.take().unwrap_or_else(|| refused(error))
}
}
const fn measure(metric: Metric) -> Measure {
match metric {
Metric::Cosine => Measure::Cosine,
Metric::L2 => Measure::L2,
Metric::Dot => Measure::Dot,
}
}
const fn metric(metric: Metric) -> &'static str {
match metric {
Metric::Cosine => "cosine",
Metric::L2 => "l2",
Metric::Dot => "dot",
}
}
#[cfg(all(test, unix))]
#[path = "rerank_tests.rs"]
mod tests;
#[cfg(all(test, unix))]
#[path = "rerank_text_tests.rs"]
mod text_tests;
fn memory_failure(limit: u64) -> Failure {
Failure::new(
"query_memory_limit",
format!("similar questions exceed query memory limit {limit} bytes"),
)
}