uqa_operators/
vector.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Vector similarity and KNN operators.
8
9pub mod diagnostics;
10
11use std::sync::Arc;
12
13use uqa_core::{FieldName, IndexStats, Payload, PostingEntry, PostingList};
14use uqa_scoring::cosine_to_probability;
15use uqa_storage::{StorageBackendError, StorageBackendResult};
16
17use crate::base::{
18    missing_backend, require_finite_score, ExecutionContext, Operator, OperatorResult,
19};
20
21/// `V_theta(q)`: returns documents with cosine similarity at least
22/// `threshold` (Definition 3.1.2). Returns an empty posting list when the
23/// field has no vector index registered.
24pub struct VectorSimilarityOperator {
25    pub query_vector: Vec<f32>,
26    pub threshold: f32,
27    pub field: FieldName,
28}
29
30impl VectorSimilarityOperator {
31    pub fn new(query_vector: Vec<f32>, threshold: f32, field: impl Into<FieldName>) -> Self {
32        Self {
33            query_vector,
34            threshold,
35            field: field.into(),
36        }
37    }
38}
39
40impl Operator for VectorSimilarityOperator {
41    fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
42        validate_vector_query(&self.query_vector, "vector similarity")?;
43        if !self.threshold.is_finite() || !(-1.0..=1.0).contains(&self.threshold) {
44            return Err(StorageBackendError::Other(format!(
45                "vector similarity threshold must be finite and in [-1, 1], got {}",
46                self.threshold
47            )));
48        }
49        let index = ctx
50            .vector_indexes
51            .get(&self.field)
52            .ok_or_else(|| missing_backend("vector-index", "vector similarity"))?;
53        diagnostics::search_threshold(
54            index.as_ref(),
55            &self.field,
56            &self.query_vector,
57            self.threshold,
58            ctx.vector_observer.as_deref(),
59        )
60    }
61
62    fn cost_estimate(&self, stats: &IndexStats) -> f64 {
63        let n = (stats.total_docs + 1) as f64;
64        f64::from(stats.dimensions) * n.log2()
65    }
66}
67
68/// `KNN_k(q)`: top-`k` nearest neighbors by cosine similarity
69/// (Definition 3.1.3).
70pub struct KNNOperator {
71    pub query_vector: Vec<f32>,
72    pub k: usize,
73    pub field: FieldName,
74}
75
76impl KNNOperator {
77    pub fn new(query_vector: Vec<f32>, k: usize, field: impl Into<FieldName>) -> Self {
78        Self {
79            query_vector,
80            k,
81            field: field.into(),
82        }
83    }
84}
85
86impl Operator for KNNOperator {
87    fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
88        validate_vector_query(&self.query_vector, "KNN search")?;
89        let index = ctx
90            .vector_indexes
91            .get(&self.field)
92            .ok_or_else(|| missing_backend("vector-index", "KNN search"))?;
93        diagnostics::search_knn(
94            index.as_ref(),
95            &self.field,
96            &self.query_vector,
97            self.k,
98            ctx.vector_observer.as_deref(),
99        )
100    }
101
102    fn cost_estimate(&self, stats: &IndexStats) -> f64 {
103        let n = (stats.total_docs + 1) as f64;
104        f64::from(stats.dimensions) * n.log2()
105    }
106}
107
108/// Marks a vector operator (KNN, threshold, ...) for probability-domain
109/// evidence conversion. Direct operator execution applies the uncalibrated
110/// `(1 + score) / 2` bridge from cosine similarity in `[-1, 1]` to `[0, 1]`
111/// (Definition 7.1.2, Paper 3). At an engine fusion boundary the driver
112/// intercepts this marker and fits the Paper 5 query-pool likelihood-ratio
113/// transform instead; reusable calibrated models live in
114/// `uqa_scoring::VectorCalibrationModel`.
115pub struct CosineProbabilityOperator {
116    pub source: Arc<dyn Operator>,
117}
118
119impl CosineProbabilityOperator {
120    pub fn new(source: Arc<dyn Operator>) -> Self {
121        Self { source }
122    }
123}
124
125impl Operator for CosineProbabilityOperator {
126    fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
127        let pl = self.source.execute(ctx)?;
128        let mut entries = Vec::with_capacity(pl.len());
129        for entry in pl.entries() {
130            require_finite_score(entry.payload.score, "cosine probability projection")?;
131            if !(-1.0..=1.0).contains(&entry.payload.score) {
132                return Err(StorageBackendError::Other(format!(
133                    "cosine probability projection requires scores in [-1, 1], got {}",
134                    entry.payload.score
135                )));
136            }
137            entries.push(PostingEntry::new(
138                entry.doc_id,
139                Payload {
140                    positions: entry.payload.positions.clone(),
141                    score: cosine_to_probability(entry.payload.score),
142                    fields: entry.payload.fields.clone(),
143                },
144            ));
145        }
146        Ok(PostingList::from_sorted_unchecked(entries))
147    }
148
149    fn cost_estimate(&self, stats: &IndexStats) -> f64 {
150        self.source.cost_estimate(stats)
151    }
152}
153
154fn validate_vector_query(query: &[f32], operation: &str) -> StorageBackendResult<()> {
155    if query.is_empty() {
156        return Err(StorageBackendError::Other(format!(
157            "{operation} requires a non-empty query vector"
158        )));
159    }
160    if query.iter().any(|component| !component.is_finite()) {
161        return Err(StorageBackendError::Other(format!(
162            "{operation} query vector must contain only finite values"
163        )));
164    }
165    Ok(())
166}
uqa_operators/vector.rs

uqa_operators/
vector.rs
