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

hawdb_executor/
vector.rs
