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 adaptive_exact_scan_count: usize,
190    pub payload_bytes_read: u64,
191    pub admitted_working_bytes: usize,
192}
193
194#[derive(Debug, Clone, PartialEq, Eq)]
195pub struct VectorExecutionReport {
196    pub backend: VectorExecutionBackend,
197    pub compression_mode: VectorCompressionMode,
198    pub candidate_source: VectorCandidateSource,
199    pub backend_selection_reason: Option<VectorBackendSelectionReason>,
200    pub estimated_raw_vector_bytes: Option<u64>,
201    pub filter_selectivity_per_million: Option<u32>,
202    pub candidate_score_source: VectorScoreSource,
203    pub final_score_source: VectorScoreSource,
204    pub generated_candidate_count: usize,
205    pub descriptor_pruned_count: usize,
206    pub scalar_filtered_count: usize,
207    pub residual_filtered_count: usize,
208    pub candidate_scan_rounds: usize,
209    pub reranked_candidate_count: usize,
210    pub returned_count: usize,
211    pub raw_vector_bytes_read: u64,
212    pub candidate_scan_metrics: Option<VectorCandidateScanMetrics>,
213    pub index_covered_document_count: Option<usize>,
214    pub index_candidate_document_count: Option<usize>,
215    pub index_coverage_complete: Option<bool>,
216    pub fallback_reason_codes: Vec<VectorFallbackReasonCode>,
217}
218
219#[derive(Debug, Clone, PartialEq)]
220pub struct VectorExecutionOutput {
221    pub scores: Vec<VectorRawScore>,
222    pub candidate_ids: Vec<String>,
223    pub report: VectorExecutionReport,
224}
225
226#[derive(Debug, Clone, PartialEq, Eq)]
227pub enum VectorExecutionError<E> {
228    InvalidPlan(&'static str),
229    Source(E),
230    DuplicateCandidate,
231    DuplicateRawScore,
232    NonFiniteCandidateScore,
233    NonFiniteRawScore,
234    ResidualFilterAddedCandidate,
235    RawScoreOutsideCandidateSet,
236}
237
238impl<E: Display> Display for VectorExecutionError<E> {
239    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
240        match self {
241            Self::InvalidPlan(reason) => write!(formatter, "invalid vector plan: {reason}"),
242            Self::Source(error) => write!(formatter, "vector execution source failed: {error}"),
243            Self::DuplicateCandidate => formatter.write_str("vector candidate ids must be unique"),
244            Self::DuplicateRawScore => formatter.write_str("raw vector score ids must be unique"),
245            Self::NonFiniteCandidateScore => {
246                formatter.write_str("vector candidate score must be finite")
247            }
248            Self::NonFiniteRawScore => formatter.write_str("raw vector score must be finite"),
249            Self::ResidualFilterAddedCandidate => {
250                formatter.write_str("residual vector filter may only retain generated candidates")
251            }
252            Self::RawScoreOutsideCandidateSet => {
253                formatter.write_str("raw vector score must reference a generated candidate")
254            }
255        }
256    }
257}
258
259impl<E: std::error::Error + 'static> std::error::Error for VectorExecutionError<E> {}
260
261pub fn execute_vector_plan<S: VectorExecutionSource>(
262    plan: &VectorPhysicalPlan,
263    source: &mut S,
264) -> Result<VectorExecutionOutput, VectorExecutionError<S::Error>> {
265    let VectorPhysicalPlan::TopK {
266        limit: top_k,
267        input,
268    } = plan
269    else {
270        return Err(VectorExecutionError::InvalidPlan("missing TopK root"));
271    };
272    let VectorPhysicalPlan::RawVectorRerank {
273        embedding_dimension: rerank_dimension,
274        input,
275    } = input.as_ref()
276    else {
277        return Err(VectorExecutionError::InvalidPlan(
278            "TopK must consume RawVectorRerank",
279        ));
280    };
281    let (residual_filter, input) = match input.as_ref() {
282        VectorPhysicalPlan::ResidualFilter {
283            fields,
284            initial_candidate_limit,
285            input,
286        } => (
287            Some((fields.as_slice(), *initial_candidate_limit)),
288            input.as_ref(),
289        ),
290        input => (None, input),
291    };
292    let VectorPhysicalPlan::VectorCandidateScan {
293        source: candidate_source,
294        embedding_dimension: candidate_dimension,
295        candidate_limit,
296        input,
297    } = input
298    else {
299        return Err(VectorExecutionError::InvalidPlan(
300            "RawVectorRerank must consume VectorCandidateScan",
301        ));
302    };
303    let VectorPhysicalPlan::Filter { fields } = input.as_ref() else {
304        return Err(VectorExecutionError::InvalidPlan(
305            "VectorCandidateScan must consume Filter",
306        ));
307    };
308    if candidate_dimension != rerank_dimension {
309        return Err(VectorExecutionError::InvalidPlan(
310            "candidate and rerank dimensions differ",
311        ));
312    }
313    if *top_k == 0 || *candidate_limit < *top_k {
314        return Err(VectorExecutionError::InvalidPlan(
315            "candidate and top-k limits are inconsistent",
316        ));
317    }
318    if let Some((_, initial_limit)) = residual_filter
319        && (initial_limit < *top_k || initial_limit > *candidate_limit)
320    {
321        return Err(VectorExecutionError::InvalidPlan(
322            "residual candidate window is outside the top-k candidate budget",
323        ));
324    }
325
326    let mut scan_limit = residual_filter
327        .map(|(_, initial_limit)| initial_limit)
328        .unwrap_or(*candidate_limit);
329    let mut candidate_scan_rounds = 0;
330    let (batch, generated_candidate_count) = loop {
331        candidate_scan_rounds += 1;
332        let mut batch = source
333            .scan_candidates(VectorCandidateScanRequest {
334                source: *candidate_source,
335                embedding_dimension: *candidate_dimension,
336                candidate_limit: scan_limit,
337                filter_fields: fields,
338            })
339            .map_err(VectorExecutionError::Source)?;
340        validate_candidates(&batch.candidates)?;
341        sort_candidates(&mut batch.candidates);
342        batch.candidates.truncate(scan_limit);
343        let generated_candidate_count = batch.candidates.len();
344        if let Some((residual_fields, _)) = residual_filter {
345            let generated_ids = batch
346                .candidates
347                .iter()
348                .map(|candidate| candidate.id.clone())
349                .collect::<BTreeSet<_>>();
350            batch.candidates = source
351                .filter_residual(VectorResidualFilterRequest {
352                    fields: residual_fields,
353                    candidates: batch.candidates,
354                })
355                .map_err(VectorExecutionError::Source)?;
356            validate_candidates(&batch.candidates)?;
357            if batch
358                .candidates
359                .iter()
360                .any(|candidate| !generated_ids.contains(&candidate.id))
361            {
362                return Err(VectorExecutionError::ResidualFilterAddedCandidate);
363            }
364        }
365        if residual_filter.is_none()
366            || batch.candidates.len() >= *top_k
367            || scan_limit >= *candidate_limit
368            || generated_candidate_count < scan_limit
369        {
370            break (batch, generated_candidate_count);
371        }
372        scan_limit = scan_limit.saturating_mul(2).min(*candidate_limit);
373    };
374    let residual_filtered_count = generated_candidate_count.saturating_sub(batch.candidates.len());
375    let candidate_ids = if source.capture_candidate_ids() {
376        batch
377            .candidates
378            .iter()
379            .map(|candidate| candidate.id.clone())
380            .collect()
381    } else {
382        Vec::new()
383    };
384
385    let mut raw_scores = source
386        .rerank_raw(VectorRawRerankRequest {
387            embedding_dimension: *rerank_dimension,
388            candidates: &batch.candidates,
389        })
390        .map_err(VectorExecutionError::Source)?;
391    validate_raw_scores(&batch.candidates, &raw_scores)?;
392    sort_raw_scores(&mut raw_scores);
393    let reranked_candidate_count = raw_scores.len();
394    raw_scores.truncate(*top_k);
395
396    let final_score_source = if batch.score_source == VectorScoreSource::Unavailable {
397        VectorScoreSource::Unavailable
398    } else {
399        VectorScoreSource::RawVector
400    };
401    Ok(VectorExecutionOutput {
402        candidate_ids,
403        report: VectorExecutionReport {
404            backend: VectorExecutionBackend::from_candidate_source(*candidate_source),
405            compression_mode: VectorCompressionMode::Unspecified,
406            candidate_source: *candidate_source,
407            backend_selection_reason: None,
408            estimated_raw_vector_bytes: None,
409            filter_selectivity_per_million: None,
410            candidate_score_source: batch.score_source,
411            final_score_source,
412            generated_candidate_count,
413            descriptor_pruned_count: 0,
414            scalar_filtered_count: residual_filtered_count,
415            residual_filtered_count,
416            candidate_scan_rounds,
417            reranked_candidate_count,
418            returned_count: raw_scores.len(),
419            raw_vector_bytes_read: source.raw_vector_bytes_read(),
420            candidate_scan_metrics: source.candidate_scan_metrics(),
421            index_covered_document_count: None,
422            index_candidate_document_count: None,
423            index_coverage_complete: None,
424            fallback_reason_codes: Vec::new(),
425        },
426        scores: raw_scores,
427    })
428}
429
430fn validate_candidates<E>(candidates: &[VectorCandidate]) -> Result<(), VectorExecutionError<E>> {
431    let mut ids = BTreeSet::new();
432    for candidate in candidates {
433        if !candidate.score.is_finite() {
434            return Err(VectorExecutionError::NonFiniteCandidateScore);
435        }
436        if !ids.insert(candidate.id.as_str()) {
437            return Err(VectorExecutionError::DuplicateCandidate);
438        }
439    }
440    Ok(())
441}
442
443fn validate_raw_scores<E>(
444    candidates: &[VectorCandidate],
445    raw_scores: &[VectorRawScore],
446) -> Result<(), VectorExecutionError<E>> {
447    let candidate_ids = candidates
448        .iter()
449        .map(|candidate| candidate.id.as_str())
450        .collect::<BTreeSet<_>>();
451    let mut score_ids = BTreeSet::new();
452    for score in raw_scores {
453        if !score.score.is_finite() {
454            return Err(VectorExecutionError::NonFiniteRawScore);
455        }
456        if !candidate_ids.contains(score.id.as_str()) {
457            return Err(VectorExecutionError::RawScoreOutsideCandidateSet);
458        }
459        if !score_ids.insert(score.id.as_str()) {
460            return Err(VectorExecutionError::DuplicateRawScore);
461        }
462    }
463    Ok(())
464}
465
466fn sort_candidates(candidates: &mut [VectorCandidate]) {
467    candidates.sort_by(|left, right| {
468        descending_score_order(left.score, &left.id, right.score, &right.id)
469    });
470}
471
472fn sort_raw_scores(scores: &mut [VectorRawScore]) {
473    scores.sort_by(|left, right| {
474        descending_score_order(left.score, &left.id, right.score, &right.id)
475    });
476}
477
478fn descending_score_order(
479    left_score: f64,
480    left_id: &str,
481    right_score: f64,
482    right_id: &str,
483) -> Ordering {
484    right_score
485        .partial_cmp(&left_score)
486        .unwrap_or(Ordering::Equal)
487        .then_with(|| left_id.cmp(right_id))
488}
489
490#[cfg(test)]
491mod tests;
hawdb_executor/vector.rs

hawdb_executor/
vector.rs
