hawdb_executor/
vector_eligibility.rs1use crate::{QueryMemoryAccount, QueryMemoryLease};
7use hawdb_core::{
8 HawDBError, Result, RuntimeTaskContext, VectorIndexDescriptor, VectorIndexQueryMode, VectorRef,
9};
10use std::collections::BTreeMap;
11
12#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
15pub struct VectorEligibilityReport {
16 pub input_rows: usize,
17 pub work_units: usize,
18 pub max_work: usize,
19 pub source_strategy: &'static str,
20 pub eligible_vectors: usize,
21 pub scored_vectors: usize,
22 pub complete: bool,
23}
24
25#[derive(Debug)]
26pub struct VectorEligibleDomain<'a> {
27 definition: &'a VectorIndexDescriptor,
28 identity: hawdb_core::VectorIndexReadIdentity,
29 embedding: VectorRef<'a>,
30 mode: VectorIndexQueryMode,
31 entries: BTreeMap<String, Option<f64>>,
32 report: VectorEligibilityReport,
33 task: RuntimeTaskContext,
34 memory: QueryMemoryLease,
36}
37
38impl<'a> VectorEligibleDomain<'a> {
39 pub fn new(
40 definition: &'a VectorIndexDescriptor,
41 identity: hawdb_core::VectorIndexReadIdentity,
42 embedding: VectorRef<'a>,
43 mode: VectorIndexQueryMode,
44 max_work: usize,
45 account: &QueryMemoryAccount,
46 task: Option<&RuntimeTaskContext>,
47 ) -> Result<Self> {
48 if !mode.requires_eligible_domain() {
49 return Err(HawDBError::Semantic(
50 "ann_window does not accept an eligible domain".into(),
51 ));
52 }
53 if embedding.dimension() != definition.dimension() {
54 return Err(HawDBError::Semantic(
55 "vector query dimension differs from the bound index".into(),
56 ));
57 }
58 definition
60 .distance()
61 .distance_with_work(embedding, embedding, || {
62 crate::pipeline::runtime_checkpoint(task)
63 })?;
64 let memory = account.reserve(std::mem::size_of::<Self>().saturating_add(256))?;
65 Ok(Self {
66 definition,
67 identity,
68 embedding,
69 mode,
70 entries: BTreeMap::new(),
71 report: VectorEligibilityReport {
72 max_work,
73 source_strategy: "bounded_canonical",
74 ..Default::default()
75 },
76 task: task.cloned().unwrap_or_default().with_work_limit(max_work),
77 memory,
78 })
79 }
80
81 pub fn observe_input(&mut self) -> Result<()> {
82 if self.report.complete {
83 return Err(HawDBError::Execution(
84 "vector eligible domain is already complete".into(),
85 ));
86 }
87 self.task.charge_work(1)?;
88 self.report.input_rows += 1;
89 Ok(())
90 }
91
92 pub fn insert(&mut self, identity: &str, vector: VectorRef<'_>) -> Result<()> {
96 crate::pipeline::runtime_checkpoint(Some(&self.task))?;
97 if self.report.complete {
98 return Err(HawDBError::Execution(
99 "vector eligible domain is already complete".into(),
100 ));
101 }
102 if self.entries.contains_key(identity) {
103 return Ok(());
104 }
105 if vector.dimension() != self.definition.dimension() {
106 return Err(HawDBError::StorageIntegrity(
107 "eligible vector dimension differs from the bound index".into(),
108 ));
109 }
110 let distance =
111 if self.mode == VectorIndexQueryMode::Exact {
112 Some(self.definition.distance().distance_with_work(
113 self.embedding,
114 vector,
115 || crate::pipeline::runtime_checkpoint(Some(&self.task)),
116 )?)
117 } else {
118 None
119 };
120 self.memory.grow(identity.len().saturating_add(256))?;
122 self.entries.insert(identity.to_owned(), distance);
123 self.report.eligible_vectors += 1;
124 self.report.scored_vectors += usize::from(distance.is_some());
125 Ok(())
126 }
127
128 pub fn finish(&mut self) -> Result<()> {
129 crate::pipeline::runtime_checkpoint(Some(&self.task))?;
130 self.task.check_work_limit()?;
131 self.report.complete = true;
132 Ok(())
133 }
134
135 pub fn validate(
136 &self,
137 definition: &VectorIndexDescriptor,
138 embedding: &[f32],
139 mode: VectorIndexQueryMode,
140 ) -> Result<()> {
141 crate::pipeline::runtime_checkpoint(Some(&self.task))?;
142 self.task.check_work_limit()?;
143 if !self.report.complete
144 || self.mode != mode
145 || self.definition != definition
146 || !self.embedding.coordinates().eq(embedding.iter().copied())
147 {
148 return Err(HawDBError::Execution(
149 "vector eligible domain differs from the completed query binding".into(),
150 ));
151 }
152 Ok(())
153 }
154
155 pub fn identity(&self) -> hawdb_core::VectorIndexReadIdentity {
156 self.identity
157 }
158
159 pub fn task(&self) -> &RuntimeTaskContext {
160 &self.task
161 }
162
163 pub fn report(&self) -> VectorEligibilityReport {
164 VectorEligibilityReport {
165 work_units: self.task.work_units(),
166 ..self.report
167 }
168 }
169
170 pub fn find_identity(&self, identity: &str) -> Option<&str> {
171 self.entries
172 .get_key_value(identity)
173 .map(|(key, _)| key.as_str())
174 }
175
176 pub fn distances(&self) -> impl Iterator<Item = Result<(&str, f64)>> {
177 self.entries.iter().map(|(id, distance)| {
178 distance
179 .map(|distance| (id.as_str(), distance))
180 .ok_or_else(|| HawDBError::Execution("eligible identity has no exact score".into()))
181 })
182 }
183}