sz_orm_vector/
lib.rs

1//! SZ-ORM pgvector Extension
2//!
3//! Provides PostgreSQL pgvector vector similarity search capabilities, supporting three implementations:
4//!
5//! - **In-memory implementation** (`InMemoryVectorStore`): Pure Rust vector computation, no database connection, suitable for testing and benchmarking
6//! - **Stub implementation** (`StubVectorStore`): All methods return Unsupported, suitable for debug placeholder
7//! - **Real implementation** (`RealPgVectorStore`, requires `real-pg` feature): Connects to PostgreSQL + pgvector via tokio-postgres
8//!
9//! # Quick Start
10//!
11//! ```rust
12//! use sz_orm_vector::{InMemoryVectorStore, PgVectorStore, VectorRecord, VectorMetric};
13//!
14//! # #[tokio::main]
15//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
16//! let store = InMemoryVectorStore::new();
17//! store.create_collection("docs", 3, None).await?;
18//!
19//! let record = VectorRecord::new("doc1", vec![1.0, 0.0, 0.0]);
20//! store.insert("docs", vec![record]).await?;
21//!
22//! let results = store.search("docs", &[1.0, 0.0, 0.0], 5).await?;
23//! println!("found {} results", results.len());
24//! # Ok(())
25//! # }
26//! ```
27
28pub mod error;
29pub mod extensions;
30pub mod memory;
31pub mod stub;
32
33#[cfg(feature = "embedding-store")]
34pub mod embedding_store;
35
36#[cfg(feature = "real-pg")]
37pub mod real_pg;
38
39// v4.0.0 M3：混合搜索（hybrid-search feature gate 隔离）
40#[cfg(feature = "hybrid-search")]
41pub mod hybrid_search;
42
43pub use error::VectorError;
44pub use extensions::{
45    AnnIndexDef, AnnIndexRegistry, AnnIndexType, BatchOpsExt, DimensionValidator, HnswParams,
46    IvfflatParams, MemoryBatchOps, SimilarityAlgorithms, VectorNormalizer, MAX_VECTOR_DIMENSION,
47    MIN_VECTOR_DIMENSION,
48};
49pub use memory::InMemoryVectorStore;
50pub use stub::StubVectorStore;
51
52#[cfg(feature = "real-pg")]
53pub use real_pg::{RealPgConfig, RealPgVectorStore};
54
55use async_trait::async_trait;
56use std::collections::HashMap;
57use std::str::FromStr;
58
59/// M-16 fix: top_k maximum limit
60///
61/// Limit top_k upper bound to prevent:
62/// - Large k values causing memory explosion (each SearchResult contains full vector)
63/// - Performance issues with database/vector engine executing huge k queries
64/// - Malicious caller triggering OOM via top_k=usize::MAX
65pub const MAX_TOP_K: usize = 10_000;
66
67/// M-16 fix: Validate whether top_k is within reasonable range
68///
69/// - `top_k = 0`: Returns `TopKExceeded` error (meaningless query)
70/// - `top_k > MAX_TOP_K`: Returns `TopKExceeded` error
71/// - `1 <= top_k <= MAX_TOP_K`: Returns Ok
72pub fn validate_top_k(top_k: usize) -> Result<usize, VectorError> {
73    if top_k == 0 {
74        return Err(VectorError::TopKExceeded {
75            requested: top_k,
76            max: MAX_TOP_K,
77        });
78    }
79    if top_k > MAX_TOP_K {
80        return Err(VectorError::TopKExceeded {
81            requested: top_k,
82            max: MAX_TOP_K,
83        });
84    }
85    Ok(top_k)
86}
87
88/// Vector record
89#[derive(Debug, Clone)]
90pub struct VectorRecord {
91    pub id: String,
92    pub vector: Vec<f32>,
93    pub score: Option<f32>,
94    pub metadata: Option<HashMap<String, serde_json::Value>>,
95}
96
97impl VectorRecord {
98    pub fn new(id: impl Into<String>, vector: Vec<f32>) -> Self {
99        Self {
100            id: id.into(),
101            vector,
102            score: None,
103            metadata: None,
104        }
105    }
106
107    pub fn with_score(mut self, score: f32) -> Self {
108        self.score = Some(score);
109        self
110    }
111
112    pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
113        self.metadata = Some(metadata);
114        self
115    }
116}
117
118/// Search result
119#[derive(Debug, Clone)]
120pub struct SearchResult {
121    pub id: String,
122    pub score: f32,
123    pub vector: Vec<f32>,
124    pub text: Option<String>,
125    pub metadata: Option<HashMap<String, serde_json::Value>>,
126}
127
128impl SearchResult {
129    pub fn new(id: impl Into<String>, score: f32, vector: Vec<f32>) -> Self {
130        Self {
131            id: id.into(),
132            score,
133            vector,
134            text: None,
135            metadata: None,
136        }
137    }
138
139    pub fn with_text(mut self, text: impl Into<String>) -> Self {
140        self.text = Some(text.into());
141        self
142    }
143
144    pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
145        self.metadata = Some(metadata);
146        self
147    }
148}
149
150/// Vector distance metric
151#[derive(Debug, Clone, Copy, Default, PartialEq)]
152pub enum VectorMetric {
153    #[default]
154    Cosine,
155    Euclidean,
156    DotProduct,
157}
158
159impl VectorMetric {
160    /// pgvector operator mapping
161    pub fn pg_operator(&self) -> &'static str {
162        match self {
163            VectorMetric::Cosine => "<=>",
164            VectorMetric::Euclidean => "<->",
165            VectorMetric::DotProduct => "<#>",
166        }
167    }
168
169    pub fn as_str(&self) -> &'static str {
170        match self {
171            VectorMetric::Cosine => "cosine",
172            VectorMetric::Euclidean => "euclidean",
173            VectorMetric::DotProduct => "dotproduct",
174        }
175    }
176}
177
178impl FromStr for VectorMetric {
179    type Err = String;
180
181    fn from_str(s: &str) -> Result<Self, Self::Err> {
182        match s {
183            "cosine" => Ok(VectorMetric::Cosine),
184            "euclidean" => Ok(VectorMetric::Euclidean),
185            "dotproduct" => Ok(VectorMetric::DotProduct),
186            _ => Err(format!("unknown vector metric: {}", s)),
187        }
188    }
189}
190
191/// Vector Store core trait
192///
193/// Provides CRUD and similarity search capabilities for vector collections.
194/// All methods are async, suitable for real database I/O.
195#[async_trait]
196pub trait PgVectorStore: Send + Sync {
197    /// Create collection
198    async fn create_collection(
199        &self,
200        name: &str,
201        dimension: usize,
202        metric: Option<VectorMetric>,
203    ) -> Result<(), VectorError>;
204
205    /// Delete collection
206    async fn delete_collection(&self, name: &str) -> Result<(), VectorError>;
207
208    /// Insert vector record (upsert semantics: same id overwrites)
209    async fn insert(&self, collection: &str, records: Vec<VectorRecord>)
210        -> Result<(), VectorError>;
211
212    /// Similarity search
213    ///
214    /// M-16 fix: `top_k` must be in `[1, MAX_TOP_K]` range.
215    /// Implementations should call `validate_top_k(top_k)?` before executing search.
216    async fn search(
217        &self,
218        collection: &str,
219        query: &[f32],
220        top_k: usize,
221    ) -> Result<Vec<SearchResult>, VectorError>;
222
223    /// Get single record
224    async fn get(&self, collection: &str, id: &str) -> Result<Option<VectorRecord>, VectorError>;
225
226    /// Delete record
227    async fn delete(&self, collection: &str, ids: Vec<String>) -> Result<u64, VectorError>;
228
229    /// Count records
230    async fn count(&self, collection: &str) -> Result<usize, VectorError>;
231}
232// v7.3.0 任务 3.5：向量 ANN 加速与召回指标与租户隔离
233// 启用 ann-accel feature 时编译，默认不启用
234
235/// ANN 搜索结果（含召回率指标）
236#[derive(Debug, Clone)]
237pub struct AnnSearchResult {
238    /// 检索结果
239    pub records: Vec<SearchResult>,
240    /// 召回率（ANN 结果 vs 暴力 kNN，0.0 ~ 1.0）
241    pub recall_rate: f64,
242    /// 延迟（毫秒）
243    pub latency_ms: u64,
244}
245
246/// ANN 加速检索 trait（继承 PgVectorStore）
247///
248/// 提供 ANN 索引加速检索能力，支持租户隔离过滤和召回率指标。
249/// 内存实现使用精确 kNN（召回率 = 1.0），真实 ANN 加速需启用 pgvector HNSW/IVF 索引。
250#[cfg(feature = "ann-accel")]
251#[async_trait]
252pub trait AnnAccelerated: PgVectorStore {
253    /// ANN 加速检索
254    ///
255    /// # 参数
256    /// - `collection`: 集合名称
257    /// - `query`: 查询向量
258    /// - `top_k`: 返回结果数
259    /// - `tenant_id`: 租户 ID（可选，用于租户隔离过滤）
260    ///
261    /// # 返回值
262    /// - `Ok(AnnSearchResult)`: 检索结果 + 召回率 + 延迟
263    /// - `Err(VectorError)`: 检索失败
264    async fn ann_search(
265        &self,
266        collection: &str,
267        query: &[f32],
268        top_k: usize,
269        tenant_id: Option<&str>,
270    ) -> Result<AnnSearchResult, VectorError>;
271}
272
273#[cfg(feature = "ann-accel")]
274#[async_trait]
275impl AnnAccelerated for InMemoryVectorStore {
276    async fn ann_search(
277        &self,
278        collection: &str,
279        query: &[f32],
280        top_k: usize,
281        tenant_id: Option<&str>,
282    ) -> Result<AnnSearchResult, VectorError> {
283        use std::time::Instant;
284        let start = Instant::now();
285
286        // 精确 kNN 检索（内存实现中 ANN = kNN，召回率 = 1.0）
287        let mut results = self.search(collection, query, top_k).await?;
288
289        // 租户隔离过滤：只保留 metadata 中 tenant_id 匹配的记录
290        if let Some(tid) = tenant_id {
291            results.retain(|r| {
292                r.metadata
293                    .as_ref()
294                    .and_then(|m| m.get("tenant_id"))
295                    .and_then(|v| v.as_str())
296                    == Some(tid)
297            });
298        }
299
300        let latency_ms = start.elapsed().as_millis() as u64;
301
302        // 召回率 = 1.0（内存实现使用精确 kNN）
303        let recall_rate = 1.0;
304
305        Ok(AnnSearchResult {
306            records: results,
307            recall_rate,
308            latency_ms,
309        })
310    }
311}
sz_orm_vector/lib.rs

sz_orm_vector/
lib.rs
