use super::{invalid, DiskANNQuery, DiskANNQueryWork};
use crate::diskann_index::scoring::ScoringStatsExt;
use crate::diskann_index::{
format::{DiskANNCanonicalOrigin, DiskANNNode, DiskANNVectorVersion},
scoring::selection::TopK,
DiskANNCanonicalScorer, DiskANNDocumentScore,
};
use crate::{read_control::StorageReadControl, StorageBackendResult};
use uqa_core::{memory::BudgetedMap, DocId, PostingList};
#[derive(Clone, Copy)]
struct Candidate {
origin: DiskANNCanonicalOrigin,
score: Option<DiskANNDocumentScore>,
}
impl Candidate {
fn validate(
self,
ordinal: u32,
version: DiskANNVectorVersion,
) -> StorageBackendResult<Option<DiskANNDocumentScore>> {
if self.origin.version() != version || u64::from(ordinal) >= self.origin.count() {
return Err(invalid(
"physical candidate differs from its generation origin",
));
}
Ok(self.score)
}
}
pub(super) struct Candidates<'a, 's> {
query: &'a DiskANNQuery<'s>,
scorer: DiskANNCanonicalScorer<'a>,
selected: TopK,
graph: BudgetedMap<DocId, Candidate>,
control: &'a StorageReadControl,
work: DiskANNQueryWork,
}
impl<'a, 's> Candidates<'a, 's> {
pub(super) fn new(
query: &'a DiskANNQuery<'s>,
scorer: DiskANNCanonicalScorer<'a>,
k: usize,
control: &'a StorageReadControl,
) -> Self {
Self {
query,
scorer,
selected: TopK::deduplicating(k, control),
graph: BudgetedMap::new(control.memory()),
control,
work: DiskANNQueryWork::default(),
}
}
pub(super) fn len(&self) -> usize {
self.selected.len()
}
pub(super) fn work(&self) -> DiskANNQueryWork {
self.work
}
pub(super) fn offer_unversioned(
&mut self,
document: DocId,
score: f32,
vectors: u64,
) -> StorageBackendResult<()> {
self.query.check(self.control)?;
self.work.unversioned.record(vectors)?;
self.selected.offer_raw(document, score)
}
fn candidate(&mut self, document: DocId) -> StorageBackendResult<Candidate> {
self.query.check(self.control)?;
let origin = self
.query
.origins
.origin(document, self.control)?
.ok_or_else(|| invalid("physical document has no generation origin"))?;
if origin.count() == 0 {
return Err(invalid("physical candidate belongs to an empty tensor"));
}
let score = self.scorer.score_candidate(document, 0, origin.version())?;
if score.is_some_and(|score| score.vector_count() != origin.count()) {
return Err(invalid(
"canonical tensor count differs from its original version",
));
}
if let Some(score) = score {
self.work.reranked.record(score.vector_count())?;
}
Ok(Candidate { origin, score })
}
pub(super) fn nodes(&mut self, nodes: &[DiskANNNode]) -> StorageBackendResult<()> {
for node in nodes {
self.query.check(self.control)?;
let candidate = if let Some(candidate) = self.graph.get(&node.doc_id()) {
*candidate
} else {
let candidate = self.candidate(node.doc_id())?;
self.graph.insert(node.doc_id(), candidate)?;
candidate
};
if let Some(score) = candidate.validate(node.ordinal(), node.version())? {
self.selected.offer(score)?;
}
}
Ok(())
}
pub(super) fn side(&mut self) -> StorageBackendResult<()> {
let mut previous: Option<(DocId, Candidate)> = None;
let query = self.query;
query.reader.visit_side(self.control, &mut |entry| {
self.query.check(self.control)?;
self.work.side_entries = self
.work
.side_entries
.checked_add(1)
.ok_or(uqa_core::memory::MemoryError::SizeOverflow)?;
let candidate = match previous {
Some((document, candidate)) if document == entry.doc_id() => candidate,
_ => self
.graph
.get(&entry.doc_id())
.copied()
.map_or_else(|| self.candidate(entry.doc_id()), Ok)?,
};
if let Some(score) = candidate.validate(entry.ordinal(), entry.version())? {
self.selected.offer(score)?;
}
previous = Some((entry.doc_id(), candidate));
Ok(())
})
}
pub(super) fn changes(&mut self) -> StorageBackendResult<()> {
let mut after = None;
while let Some(change) = self
.query
.canonical
.next_change_after(after, self.control)?
{
self.query.check(self.control)?;
let document = change.document();
if after.is_some_and(|previous| document <= previous) {
return Err(invalid("query change cursor did not advance"));
}
after = Some(document);
if self.query.canonical.origin(document, self.control)? != Some(change.version()) {
return Err(invalid("query change differs from its canonical origin"));
}
if self
.query
.origins
.origin(document, self.control)?
.is_some_and(|origin| origin.version() == change.version())
{
continue;
}
if let Some(score) = self.scorer.score_document(document)? {
if score.version() != change.version() {
return Err(invalid("canonical change moved while scoring"));
}
self.work.changed.record(score.vector_count())?;
self.selected.offer(score)?;
}
}
self.query.check(self.control)
}
pub(super) fn finish(self) -> StorageBackendResult<PostingList> {
drop(self.graph);
self.selected.finish(self.control)
}
}