hawdb_executor/
vector.rs

1// Copyright 2026 Nowledge
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use hawdb_plan_cypher::{VectorBackendSelectionReason, VectorCandidateSource, VectorPhysicalPlan};
16use std::cmp::Ordering;
17use std::collections::BTreeSet;
18use std::fmt::{Display, Formatter};
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum VectorScoreSource {
22    Unavailable,
23    RawVector,
24    AnnApproximate,
25    QuantizedApproximate,
26}
27
28impl VectorScoreSource {
29    pub const fn as_str(self) -> &'static str {
30        match self {
31            Self::Unavailable => "unavailable",
32            Self::RawVector => "raw_vector",
33            Self::AnnApproximate => "ann_approximate",
34            Self::QuantizedApproximate => "quantized_approximate",
35        }
36    }
37}
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum VectorExecutionBackend {
41    ScalarFlat,
42    AnnProjection,
43    QuantizedProjection,
44    Unavailable,
45}
46
47impl VectorExecutionBackend {
48    pub const fn as_str(self) -> &'static str {
49        match self {
50            Self::ScalarFlat => "scalar_flat",
51            Self::AnnProjection => "ann_projection",
52            Self::QuantizedProjection => "quantized_projection",
53            Self::Unavailable => "unavailable",
54        }
55    }
56
57    const fn from_candidate_source(source: VectorCandidateSource) -> Self {
58        match source {
59            VectorCandidateSource::Scalar => Self::ScalarFlat,
60            VectorCandidateSource::Ann => Self::AnnProjection,
61            VectorCandidateSource::Quantized => Self::QuantizedProjection,
62        }
63    }
64}
65
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67pub enum VectorCompressionMode {
68    Unspecified,
69    Disabled,
70    Preferred,
71    Required,
72}
73
74impl VectorCompressionMode {
75    pub const fn as_str(self) -> &'static str {
76        match self {
77            Self::Unspecified => "unspecified",
78            Self::Disabled => "disabled",
79            Self::Preferred => "preferred",
80            Self::Required => "required",
81        }
82    }
83}
84
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub enum VectorFallbackReasonCode {
87    VectorDimensionMismatch,
88    VectorIndexEmpty,
89    CompressedVectorProjectionUnavailable,
90    CompressedVectorBudgetExceeded,
91    DiskAnnCacheBudgetExceeded,
92    QueryEmbeddingMissing,
93}
94
95impl VectorFallbackReasonCode {
96    pub const fn as_str(self) -> &'static str {
97        match self {
98            Self::VectorDimensionMismatch => "vector_dimension_mismatch",
99            Self::VectorIndexEmpty => "vector_index_empty",
100            Self::CompressedVectorProjectionUnavailable => {
101                "compressed_vector_projection_unavailable"
102            }
103            Self::CompressedVectorBudgetExceeded => "compressed_vector_budget_exceeded",
104            Self::DiskAnnCacheBudgetExceeded => "diskann_cache_budget_exceeded",
105            Self::QueryEmbeddingMissing => "query_embedding_missing",
106        }
107    }
108}
109
110#[derive(Debug, Clone, PartialEq)]
111pub struct VectorCandidate {
112    pub id: String,
113    pub score: f64,
114}
115
116#[derive(Debug, Clone, PartialEq)]
117pub struct VectorCandidateBatch {
118    pub score_source: VectorScoreSource,
119    pub candidates: Vec<VectorCandidate>,
120}
121
122#[derive(Debug, Clone, PartialEq)]
123pub struct VectorRawScore {
124    pub id: String,
125    pub score: f64,
126}
127
128#[derive(Debug, Clone, Copy)]
129pub struct VectorCandidateScanRequest<'a> {
130    pub source: VectorCandidateSource,
131    pub embedding_dimension: usize,
132    pub candidate_limit: usize,
133    pub filter_fields: &'a [String],
134}
135
136#[derive(Debug, Clone, Copy)]
137pub struct VectorRawRerankRequest<'a> {
138    pub embedding_dimension: usize,
139    pub candidates: &'a [VectorCandidate],
140}
141
142#[derive(Debug)]
143pub struct VectorResidualFilterRequest<'a> {
144    pub fields: &'a [String],
145    pub candidates: Vec<VectorCandidate>,
146}
147
148pub trait VectorExecutionSource {
149    type Error;
150
151    fn scan_candidates(
152        &mut self,
153        request: VectorCandidateScanRequest<'_>,
154    ) -> Result<VectorCandidateBatch, Self::Error>;
155
156    fn filter_residual(
157        &mut self,
158        request: VectorResidualFilterRequest<'_>,
159    ) -> Result<Vec<VectorCandidate>, Self::Error>;
160
161    fn rerank_raw(
162        &mut self,
163        request: VectorRawRerankRequest<'_>,
164    ) -> Result<Vec<VectorRawScore>, Self::Error>;
165
166    fn raw_vector_bytes_read(&self) -> u64 {
167        0
168    }
169
170    fn candidate_scan_metrics(&self) -> Option<VectorCandidateScanMetrics> {
171        None
172    }
173
174    fn capture_candidate_ids(&self) -> bool {
175        false
176    }
177}
178
179#[derive(Debug, Clone, PartialEq, Eq)]
180pub struct VectorCandidateScanMetrics {
181    pub kernel: String,
182    pub worker_count: usize,
183    pub segment_count: usize,
184    pub scanned_segment_count: usize,
185    pub scored_document_count: usize,
186    pub filtered_document_count: usize,
187    pub scanned_block_count: usize,
188    pub skipped_block_count: usize,
189    pub payload_bytes_read: u64,
190    pub admitted_working_bytes: usize,
191}
192
193#[derive(Debug, Clone, PartialEq, Eq)]
194pub struct VectorExecutionReport {
195    pub backend: VectorExecutionBackend,
196    pub compression_mode: VectorCompressionMode,
197    pub candidate_source: VectorCandidateSource,
198    pub backend_selection_reason: Option<VectorBackendSelectionReason>,
199    pub estimated_raw_vector_bytes: Option<u64>,
200    pub filter_selectivity_per_million: Option<u32>,
201    pub candidate_score_source: VectorScoreSource,
202    pub final_score_source: VectorScoreSource,
203    pub generated_candidate_count: usize,
204    pub descriptor_pruned_count: usize,
205    pub scalar_filtered_count: usize,
206    pub residual_filtered_count: usize,
207    pub candidate_scan_rounds: usize,
208    pub reranked_candidate_count: usize,
209    pub returned_count: usize,
210    pub raw_vector_bytes_read: u64,
211    pub candidate_scan_metrics: Option<VectorCandidateScanMetrics>,
212    pub index_covered_document_count: Option<usize>,
213    pub index_candidate_document_count: Option<usize>,
214    pub index_coverage_complete: Option<bool>,
215    pub fallback_reason_codes: Vec<VectorFallbackReasonCode>,
216}
217
218#[derive(Debug, Clone, PartialEq)]
219pub struct VectorExecutionOutput {
220    pub scores: Vec<VectorRawScore>,
221    pub candidate_ids: Vec<String>,
222    pub report: VectorExecutionReport,
223}
224
225#[derive(Debug, Clone, PartialEq, Eq)]
226pub enum VectorExecutionError<E> {
227    InvalidPlan(&'static str),
228    Source(E),
229    DuplicateCandidate,
230    DuplicateRawScore,
231    NonFiniteCandidateScore,
232    NonFiniteRawScore,
233    ResidualFilterAddedCandidate,
234    RawScoreOutsideCandidateSet,
235}
236
237impl<E: Display> Display for VectorExecutionError<E> {
238    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
239        match self {
240            Self::InvalidPlan(reason) => write!(formatter, "invalid vector plan: {reason}"),
241            Self::Source(error) => write!(formatter, "vector execution source failed: {error}"),
242            Self::DuplicateCandidate => formatter.write_str("vector candidate ids must be unique"),
243            Self::DuplicateRawScore => formatter.write_str("raw vector score ids must be unique"),
244            Self::NonFiniteCandidateScore => {
245                formatter.write_str("vector candidate score must be finite")
246            }
247            Self::NonFiniteRawScore => formatter.write_str("raw vector score must be finite"),
248            Self::ResidualFilterAddedCandidate => {
249                formatter.write_str("residual vector filter may only retain generated candidates")
250            }
251            Self::RawScoreOutsideCandidateSet => {
252                formatter.write_str("raw vector score must reference a generated candidate")
253            }
254        }
255    }
256}
257
258impl<E: std::error::Error + 'static> std::error::Error for VectorExecutionError<E> {}
259
260pub fn execute_vector_plan<S: VectorExecutionSource>(
261    plan: &VectorPhysicalPlan,
262    source: &mut S,
263) -> Result<VectorExecutionOutput, VectorExecutionError<S::Error>> {
264    let VectorPhysicalPlan::TopK {
265        limit: top_k,
266        input,
267    } = plan
268    else {
269        return Err(VectorExecutionError::InvalidPlan("missing TopK root"));
270    };
271    let VectorPhysicalPlan::RawVectorRerank {
272        embedding_dimension: rerank_dimension,
273        input,
274    } = input.as_ref()
275    else {
276        return Err(VectorExecutionError::InvalidPlan(
277            "TopK must consume RawVectorRerank",
278        ));
279    };
280    let (residual_filter, input) = match input.as_ref() {
281        VectorPhysicalPlan::ResidualFilter {
282            fields,
283            initial_candidate_limit,
284            input,
285        } => (
286            Some((fields.as_slice(), *initial_candidate_limit)),
287            input.as_ref(),
288        ),
289        input => (None, input),
290    };
291    let VectorPhysicalPlan::VectorCandidateScan {
292        source: candidate_source,
293        embedding_dimension: candidate_dimension,
294        candidate_limit,
295        input,
296    } = input
297    else {
298        return Err(VectorExecutionError::InvalidPlan(
299            "RawVectorRerank must consume VectorCandidateScan",
300        ));
301    };
302    let VectorPhysicalPlan::Filter { fields } = input.as_ref() else {
303        return Err(VectorExecutionError::InvalidPlan(
304            "VectorCandidateScan must consume Filter",
305        ));
306    };
307    if candidate_dimension != rerank_dimension {
308        return Err(VectorExecutionError::InvalidPlan(
309            "candidate and rerank dimensions differ",
310        ));
311    }
312    if *top_k == 0 || *candidate_limit < *top_k {
313        return Err(VectorExecutionError::InvalidPlan(
314            "candidate and top-k limits are inconsistent",
315        ));
316    }
317    if let Some((_, initial_limit)) = residual_filter
318        && (initial_limit < *top_k || initial_limit > *candidate_limit)
319    {
320        return Err(VectorExecutionError::InvalidPlan(
321            "residual candidate window is outside the top-k candidate budget",
322        ));
323    }
324
325    let mut scan_limit = residual_filter
326        .map(|(_, initial_limit)| initial_limit)
327        .unwrap_or(*candidate_limit);
328    let mut candidate_scan_rounds = 0;
329    let (batch, generated_candidate_count) = loop {
330        candidate_scan_rounds += 1;
331        let mut batch = source
332            .scan_candidates(VectorCandidateScanRequest {
333                source: *candidate_source,
334                embedding_dimension: *candidate_dimension,
335                candidate_limit: scan_limit,
336                filter_fields: fields,
337            })
338            .map_err(VectorExecutionError::Source)?;
339        validate_candidates(&batch.candidates)?;
340        sort_candidates(&mut batch.candidates);
341        batch.candidates.truncate(scan_limit);
342        let generated_candidate_count = batch.candidates.len();
343        if let Some((residual_fields, _)) = residual_filter {
344            let generated_ids = batch
345                .candidates
346                .iter()
347                .map(|candidate| candidate.id.clone())
348                .collect::<BTreeSet<_>>();
349            batch.candidates = source
350                .filter_residual(VectorResidualFilterRequest {
351                    fields: residual_fields,
352                    candidates: batch.candidates,
353                })
354                .map_err(VectorExecutionError::Source)?;
355            validate_candidates(&batch.candidates)?;
356            if batch
357                .candidates
358                .iter()
359                .any(|candidate| !generated_ids.contains(&candidate.id))
360            {
361                return Err(VectorExecutionError::ResidualFilterAddedCandidate);
362            }
363        }
364        if residual_filter.is_none()
365            || batch.candidates.len() >= *top_k
366            || scan_limit >= *candidate_limit
367            || generated_candidate_count < scan_limit
368        {
369            break (batch, generated_candidate_count);
370        }
371        scan_limit = scan_limit.saturating_mul(2).min(*candidate_limit);
372    };
373    let residual_filtered_count = generated_candidate_count.saturating_sub(batch.candidates.len());
374    let candidate_ids = if source.capture_candidate_ids() {
375        batch
376            .candidates
377            .iter()
378            .map(|candidate| candidate.id.clone())
379            .collect()
380    } else {
381        Vec::new()
382    };
383
384    let mut raw_scores = source
385        .rerank_raw(VectorRawRerankRequest {
386            embedding_dimension: *rerank_dimension,
387            candidates: &batch.candidates,
388        })
389        .map_err(VectorExecutionError::Source)?;
390    validate_raw_scores(&batch.candidates, &raw_scores)?;
391    sort_raw_scores(&mut raw_scores);
392    let reranked_candidate_count = raw_scores.len();
393    raw_scores.truncate(*top_k);
394
395    let final_score_source = if batch.score_source == VectorScoreSource::Unavailable {
396        VectorScoreSource::Unavailable
397    } else {
398        VectorScoreSource::RawVector
399    };
400    Ok(VectorExecutionOutput {
401        candidate_ids,
402        report: VectorExecutionReport {
403            backend: VectorExecutionBackend::from_candidate_source(*candidate_source),
404            compression_mode: VectorCompressionMode::Unspecified,
405            candidate_source: *candidate_source,
406            backend_selection_reason: None,
407            estimated_raw_vector_bytes: None,
408            filter_selectivity_per_million: None,
409            candidate_score_source: batch.score_source,
410            final_score_source,
411            generated_candidate_count,
412            descriptor_pruned_count: 0,
413            scalar_filtered_count: residual_filtered_count,
414            residual_filtered_count,
415            candidate_scan_rounds,
416            reranked_candidate_count,
417            returned_count: raw_scores.len(),
418            raw_vector_bytes_read: source.raw_vector_bytes_read(),
419            candidate_scan_metrics: source.candidate_scan_metrics(),
420            index_covered_document_count: None,
421            index_candidate_document_count: None,
422            index_coverage_complete: None,
423            fallback_reason_codes: Vec::new(),
424        },
425        scores: raw_scores,
426    })
427}
428
429fn validate_candidates<E>(candidates: &[VectorCandidate]) -> Result<(), VectorExecutionError<E>> {
430    let mut ids = BTreeSet::new();
431    for candidate in candidates {
432        if !candidate.score.is_finite() {
433            return Err(VectorExecutionError::NonFiniteCandidateScore);
434        }
435        if !ids.insert(candidate.id.as_str()) {
436            return Err(VectorExecutionError::DuplicateCandidate);
437        }
438    }
439    Ok(())
440}
441
442fn validate_raw_scores<E>(
443    candidates: &[VectorCandidate],
444    raw_scores: &[VectorRawScore],
445) -> Result<(), VectorExecutionError<E>> {
446    let candidate_ids = candidates
447        .iter()
448        .map(|candidate| candidate.id.as_str())
449        .collect::<BTreeSet<_>>();
450    let mut score_ids = BTreeSet::new();
451    for score in raw_scores {
452        if !score.score.is_finite() {
453            return Err(VectorExecutionError::NonFiniteRawScore);
454        }
455        if !candidate_ids.contains(score.id.as_str()) {
456            return Err(VectorExecutionError::RawScoreOutsideCandidateSet);
457        }
458        if !score_ids.insert(score.id.as_str()) {
459            return Err(VectorExecutionError::DuplicateRawScore);
460        }
461    }
462    Ok(())
463}
464
465fn sort_candidates(candidates: &mut [VectorCandidate]) {
466    candidates.sort_by(|left, right| {
467        descending_score_order(left.score, &left.id, right.score, &right.id)
468    });
469}
470
471fn sort_raw_scores(scores: &mut [VectorRawScore]) {
472    scores.sort_by(|left, right| {
473        descending_score_order(left.score, &left.id, right.score, &right.id)
474    });
475}
476
477fn descending_score_order(
478    left_score: f64,
479    left_id: &str,
480    right_score: f64,
481    right_id: &str,
482) -> Ordering {
483    right_score
484        .partial_cmp(&left_score)
485        .unwrap_or(Ordering::Equal)
486        .then_with(|| left_id.cmp(right_id))
487}
488
489#[cfg(test)]
490mod tests;
hawdb_executor/vector.rs

hawdb_executor/
vector.rs
