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

hawdb_executor/
vector.rs
