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

uqa_operators/
vector.rs
