sz_orm_vector/
lib.rs

1//! SZ-ORM pgvector 扩展
2//!
3//! 提供 PostgreSQL pgvector 向量相似度搜索能力，支持三种实现：
4//!
5//! - **内存实现**（`InMemoryVectorStore`）：纯 Rust 向量计算，不连接数据库，适用于测试和基准
6//! - **Stub 实现**（`StubVectorStore`）：所有方法返回 Unsupported，适用于调试占位
7//! - **真实实现**（`RealPgVectorStore`，需启用 `real-pg` feature）：通过 tokio-postgres 连接 PostgreSQL + pgvector
8//!
9//! # 快速入门
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 memory;
30pub mod stub;
31
32#[cfg(feature = "real-pg")]
33pub mod real_pg;
34
35pub use error::VectorError;
36pub use memory::InMemoryVectorStore;
37pub use stub::StubVectorStore;
38
39#[cfg(feature = "real-pg")]
40pub use real_pg::{RealPgConfig, RealPgVectorStore};
41
42use async_trait::async_trait;
43use std::collections::HashMap;
44use std::str::FromStr;
45
46/// M-16 修复：top_k 最大限制
47///
48/// 限制 top_k 上限以防止：
49/// - 大 k 值导致内存爆炸（每个 SearchResult 包含完整向量）
50/// - 数据库/向量引擎执行超大 k 查询的性能问题
51/// - 恶意调用方通过 top_k=usize::MAX 触发 OOM
52pub const MAX_TOP_K: usize = 10_000;
53
54/// M-16 修复：校验 top_k 是否在合理范围内
55///
56/// - `top_k = 0`：返回 `TopKExceeded` 错误（无意义的查询）
57/// - `top_k > MAX_TOP_K`：返回 `TopKExceeded` 错误
58/// - `1 <= top_k <= MAX_TOP_K`：返回 Ok
59pub fn validate_top_k(top_k: usize) -> Result<usize, VectorError> {
60    if top_k == 0 {
61        return Err(VectorError::TopKExceeded {
62            requested: top_k,
63            max: MAX_TOP_K,
64        });
65    }
66    if top_k > MAX_TOP_K {
67        return Err(VectorError::TopKExceeded {
68            requested: top_k,
69            max: MAX_TOP_K,
70        });
71    }
72    Ok(top_k)
73}
74
75/// 向量记录
76#[derive(Debug, Clone)]
77pub struct VectorRecord {
78    pub id: String,
79    pub vector: Vec<f32>,
80    pub score: Option<f32>,
81    pub metadata: Option<HashMap<String, serde_json::Value>>,
82}
83
84impl VectorRecord {
85    pub fn new(id: impl Into<String>, vector: Vec<f32>) -> Self {
86        Self {
87            id: id.into(),
88            vector,
89            score: None,
90            metadata: None,
91        }
92    }
93
94    pub fn with_score(mut self, score: f32) -> Self {
95        self.score = Some(score);
96        self
97    }
98
99    pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
100        self.metadata = Some(metadata);
101        self
102    }
103}
104
105/// 搜索结果
106#[derive(Debug, Clone)]
107pub struct SearchResult {
108    pub id: String,
109    pub score: f32,
110    pub vector: Vec<f32>,
111    pub text: Option<String>,
112    pub metadata: Option<HashMap<String, serde_json::Value>>,
113}
114
115impl SearchResult {
116    pub fn new(id: impl Into<String>, score: f32, vector: Vec<f32>) -> Self {
117        Self {
118            id: id.into(),
119            score,
120            vector,
121            text: None,
122            metadata: None,
123        }
124    }
125
126    pub fn with_text(mut self, text: impl Into<String>) -> Self {
127        self.text = Some(text.into());
128        self
129    }
130
131    pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
132        self.metadata = Some(metadata);
133        self
134    }
135}
136
137/// 向量距离度量
138#[derive(Debug, Clone, Copy, Default, PartialEq)]
139pub enum VectorMetric {
140    #[default]
141    Cosine,
142    Euclidean,
143    DotProduct,
144}
145
146impl VectorMetric {
147    /// pgvector 操作符映射
148    pub fn pg_operator(&self) -> &'static str {
149        match self {
150            VectorMetric::Cosine => "<=>",
151            VectorMetric::Euclidean => "<->",
152            VectorMetric::DotProduct => "<#>",
153        }
154    }
155
156    pub fn as_str(&self) -> &'static str {
157        match self {
158            VectorMetric::Cosine => "cosine",
159            VectorMetric::Euclidean => "euclidean",
160            VectorMetric::DotProduct => "dotproduct",
161        }
162    }
163}
164
165impl FromStr for VectorMetric {
166    type Err = String;
167
168    fn from_str(s: &str) -> Result<Self, Self::Err> {
169        match s {
170            "cosine" => Ok(VectorMetric::Cosine),
171            "euclidean" => Ok(VectorMetric::Euclidean),
172            "dotproduct" => Ok(VectorMetric::DotProduct),
173            _ => Err(format!("unknown vector metric: {}", s)),
174        }
175    }
176}
177
178/// Vector Store 核心 trait
179///
180/// 提供向量集合的 CRUD 和相似度搜索能力。
181/// 所有方法均为 async，适用于真实数据库 I/O。
182#[async_trait]
183pub trait PgVectorStore: Send + Sync {
184    /// 创建集合
185    async fn create_collection(
186        &self,
187        name: &str,
188        dimension: usize,
189        metric: Option<VectorMetric>,
190    ) -> Result<(), VectorError>;
191
192    /// 删除集合
193    async fn delete_collection(&self, name: &str) -> Result<(), VectorError>;
194
195    /// 插入向量记录（upsert 语义：相同 id 会覆盖）
196    async fn insert(&self, collection: &str, records: Vec<VectorRecord>)
197        -> Result<(), VectorError>;
198
199    /// 相似度搜索
200    ///
201    /// M-16 修复：`top_k` 必须在 `[1, MAX_TOP_K]` 范围内。
202    /// 实现方应在执行搜索前调用 `validate_top_k(top_k)?` 进行校验。
203    async fn search(
204        &self,
205        collection: &str,
206        query: &[f32],
207        top_k: usize,
208    ) -> Result<Vec<SearchResult>, VectorError>;
209
210    /// 获取单个记录
211    async fn get(&self, collection: &str, id: &str) -> Result<Option<VectorRecord>, VectorError>;
212
213    /// 删除记录
214    async fn delete(&self, collection: &str, ids: Vec<String>) -> Result<u64, VectorError>;
215
216    /// 统计记录数
217    async fn count(&self, collection: &str) -> Result<usize, VectorError>;
218}
sz_orm_vector/lib.rs

sz_orm_vector/
lib.rs
