llm_kernel/embedding/
pgvector.rs

1//! pgvector `AsyncVectorIndex` — PostgreSQL + the `pgvector` extension.
2//!
3//! `PgVectorIndex` implements [`AsyncVectorIndex`] over a `sqlx::PgPool`,
4//! mirroring `qdrant`/`elastic` (async remote vector backend). Vectors live in
5//! a `{table}(id BIGINT PK, vec vector)` relation; search uses the cosine
6//! distance operator `<=>`. Requires `CREATE EXTENSION vector` in the target DB
7//! (the table + HNSW index are created automatically by [`PgVectorIndex::new`]).
8
9use async_trait::async_trait;
10use sqlx::PgPool;
11use sqlx::QueryBuilder;
12use sqlx::postgres::PgPoolOptions;
13
14use crate::embedding::vector_index::SearchHit;
15use crate::error::{KernelError, Result};
16
17/// 검색 결과 행(id + cosine 유사도) — 튜플 대신 구조체로 sqlx `FromRow` 안정 매핑.
18#[derive(sqlx::FromRow)]
19struct ScoreRow {
20    id: i64,
21    score: f64,
22}
23
24/// f32 슬라이스 → pgvector 문자열 리터럴 `[1,2,3]` (text → vector 입력 캐스트).
25/// `pgvector::Vector`의 sqlx `Type` 바인드가 의존 환경에 따라 충돌해 문자열로 회피.
26fn vec_literal(v: &[f32]) -> String {
27    let mut s = String::from("[");
28    for (i, f) in v.iter().enumerate() {
29        if i > 0 {
30            s.push(',');
31        }
32        s.push_str(&f.to_string());
33    }
34    s.push(']');
35    s
36}
37
38/// PostgreSQL vector index backed by the `pgvector` extension.
39///
40/// All operations are async over a shared `PgPool` (MVCC, connection-pooled).
41/// The relation named `table` holds `(id BIGINT PK, vec vector(N))` rows.
42pub struct PgVectorIndex {
43    pool: PgPool,
44    table: String,
45    dim: usize,
46}
47
48impl PgVectorIndex {
49    /// Connect to `url` (libpq connstring / `postgresql://…`), create the
50    /// vector table + HNSW cosine index if missing, and return a ready index.
51    ///
52    /// `dim` is enforced by the fixed `vector(dim)` column; vectors whose
53    /// length differs from `dim` are rejected by pgvector on insert.
54    pub async fn new(url: &str, table: &str, dim: usize) -> Result<Self> {
55        validate_table_name(table)?;
56        let pool = PgPoolOptions::new()
57            .max_connections(8)
58            .connect(url)
59            .await
60            .map_err(|e| KernelError::Embedding(format!("pgvector connect: {e}")))?;
61        let idx = Self {
62            pool,
63            table: table.to_string(),
64            dim,
65        };
66        idx.init_schema().await?;
67        Ok(idx)
68    }
69
70    /// `CREATE TABLE IF NOT EXISTS {table} (id BIGINT PK, vec vector)` + HNSW
71    /// cosine index. Idempotent.
72    async fn init_schema(&self) -> Result<()> {
73        // Identifier is caller-controlled (not user input at runtime) and
74        // validated in `new`, so `format!` is acceptable here. PG cannot bind
75        // identifiers. Callers pass a fixed, validated table name.
76        sqlx::query(&format!(
77            "CREATE TABLE IF NOT EXISTS {} (id BIGINT PRIMARY KEY, vec vector({}) NOT NULL)",
78            self.table, self.dim
79        ))
80        .execute(&self.pool)
81        .await
82        .map_err(|e| KernelError::Embedding(format!("pgvector create table: {e}")))?;
83        sqlx::query(&format!(
84            "CREATE INDEX IF NOT EXISTS idx_{}_vec ON {} USING hnsw (vec vector_cosine_ops)",
85            self.table, self.table
86        ))
87        .execute(&self.pool)
88        .await
89        .map_err(|e| KernelError::Embedding(format!("pgvector hnsw index: {e}")))?;
90        Ok(())
91    }
92}
93
94#[async_trait]
95impl crate::embedding::AsyncVectorIndex for PgVectorIndex {
96    async fn add(&self, vectors: &[Vec<f32>], ids: &[u64]) -> Result<()> {
97        if vectors.len() != ids.len() {
98            return Err(KernelError::Embedding(format!(
99                "vectors.len() ({}) must equal ids.len() ({})",
100                vectors.len(),
101                ids.len()
102            )));
103        }
104        if vectors.is_empty() {
105            return Ok(());
106        }
107        // Map u64 → i64 up front: pgvector stores BIGINT (i64), so values
108        // above i64::MAX cannot be represented and must be rejected rather
109        // than silently wrapped. Single batched INSERT → one round trip.
110        let pg_ids: Vec<i64> = ids.iter().map(|&id| to_pg_id(id)).collect::<Result<_>>()?;
111        let mut q = QueryBuilder::new("INSERT INTO ");
112        q.push(self.table.as_str());
113        q.push(" (id, vec) ");
114        q.push_values(vectors.iter().zip(pg_ids.iter()), |mut b, (v, &id)| {
115            b.push_bind(id).push_bind(vec_literal(v));
116        });
117        q.push(" ON CONFLICT (id) DO UPDATE SET vec = EXCLUDED.vec");
118        q.build()
119            .execute(&self.pool)
120            .await
121            .map_err(|e| KernelError::Embedding(format!("pgvector add: {e}")))?;
122        Ok(())
123    }
124
125    async fn remove(&self, ids: &[u64]) -> Result<()> {
126        if ids.is_empty() {
127            return Ok(());
128        }
129        let ids: Vec<i64> = ids.iter().map(|&i| to_pg_id(i)).collect::<Result<_>>()?;
130        sqlx::query(&format!("DELETE FROM {} WHERE id = ANY($1)", self.table))
131            .bind(&ids)
132            .execute(&self.pool)
133            .await
134            .map_err(|e| KernelError::Embedding(format!("pgvector remove: {e}")))?;
135        Ok(())
136    }
137
138    async fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchHit>> {
139        let q = vec_literal(query);
140        // cosine distance <=> : 0 (동일) .. 2 (반대). score = 1 - distance.
141        let rows: Vec<ScoreRow> = sqlx::query_as(&format!(
142            "SELECT id, 1 - (vec <=> $1::vector) AS score FROM {} ORDER BY vec <=> $1::vector LIMIT $2",
143            self.table
144        ))
145        .bind(q)
146        .bind(k as i64)
147        .fetch_all(&self.pool)
148        .await
149        .map_err(|e| KernelError::Embedding(format!("pgvector search: {e}")))?;
150        Ok(rows
151            .into_iter()
152            .map(|r| SearchHit {
153                id: r.id as u64,
154                score: r.score as f32,
155            })
156            .collect())
157    }
158
159    async fn search_filtered(
160        &self,
161        query: &[f32],
162        k: usize,
163        allowlist: &[u64],
164    ) -> Result<Vec<SearchHit>> {
165        if allowlist.is_empty() {
166            return Ok(Vec::new());
167        }
168        let q = vec_literal(query);
169        let allow: Vec<i64> = allowlist
170            .iter()
171            .map(|&i| to_pg_id(i))
172            .collect::<Result<_>>()?;
173        let rows: Vec<ScoreRow> = sqlx::query_as(&format!(
174            "SELECT id, 1 - (vec <=> $1::vector) AS score FROM {} WHERE id = ANY($2) \
175             ORDER BY vec <=> $1::vector LIMIT $3",
176            self.table
177        ))
178        .bind(q)
179        .bind(&allow)
180        .bind(k as i64)
181        .fetch_all(&self.pool)
182        .await
183        .map_err(|e| KernelError::Embedding(format!("pgvector search_filtered: {e}")))?;
184        Ok(rows
185            .into_iter()
186            .map(|r| SearchHit {
187                id: r.id as u64,
188                score: r.score as f32,
189            })
190            .collect())
191    }
192
193    async fn len(&self) -> Result<usize> {
194        let n: i64 = sqlx::query_scalar(&format!("SELECT count(*) FROM {}", self.table))
195            .fetch_one(&self.pool)
196            .await
197            .map_err(|e| KernelError::Embedding(format!("pgvector len: {e}")))?;
198        Ok(n as usize)
199    }
200
201    fn dim(&self) -> usize {
202        self.dim
203    }
204}
205
206/// `u64` external ID → PG `BIGINT` (`i64`). IDs exceeding `i64::MAX` cannot
207/// be stored in a BIGINT column — reject rather than silently wrap.
208fn to_pg_id(id: u64) -> Result<i64> {
209    i64::try_from(id).map_err(|_| KernelError::Embedding(format!("id {id} exceeds BIGINT range")))
210}
211
212/// Validate that `table` is a plain, safe SQL identifier (ASCII alphanumeric +
213/// `_`, starting with a letter or `_`). It is interpolated into DDL/DML via
214/// `format!` (PG cannot bind identifiers), so reject anything that could break
215/// out of the identifier context.
216fn validate_table_name(table: &str) -> Result<()> {
217    let mut chars = table.chars();
218    let first_ok = matches!(chars.next(), Some(c) if c.is_ascii_alphabetic() || c == '_');
219    let valid = first_ok && chars.all(|c| c.is_ascii_alphanumeric() || c == '_');
220    if !valid {
221        return Err(KernelError::Embedding(format!(
222            "invalid table identifier: {table:?}"
223        )));
224    }
225    Ok(())
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use crate::embedding::AsyncVectorIndex;
232
233    /// `LLMKERNEL_PG_URL` 미설정 시 자동 skip (graph-pg pg.rs 패턴).
234    fn pg_url() -> Option<String> {
235        std::env::var("LLMKERNEL_PG_URL").ok()
236    }
237
238    /// add → search → search_filtered → remove 라운드트립. HNSW는 대량 데이터에서
239    /// 정확도가 보장되지만 소규모 테스트에선 정확 매칭이 간헐적일 수 있어
240    /// 여기선 id/회수 위주로 검증.
241    #[tokio::test]
242    async fn roundtrip_add_search_remove() {
243        let Some(url) = pg_url() else {
244            eprintln!("skip pgvector test: LLMKERNEL_PG_URL unset");
245            return;
246        };
247        let table = format!("lk_test_{}", line!());
248        let idx = PgVectorIndex::new(&url, &table, 3).await.expect("new");
249
250        let vecs = vec![
251            vec![1.0, 0.0, 0.0],
252            vec![0.0, 1.0, 0.0],
253            vec![0.0, 0.0, 1.0],
254        ];
255        let ids = vec![10u64, 20, 30];
256        idx.add(&vecs, &ids).await.expect("add");
257        assert_eq!(idx.len().await.unwrap(), 3);
258
259        // nearest to [1,0,0] → id 10 우선
260        let hits = idx.search(&[1.0, 0.0, 0.0], 1).await.unwrap();
261        assert_eq!(hits.len(), 1);
262        assert_eq!(hits[0].id, 10);
263
264        // filtered: allow [20,30] → 10 제외
265        let hits = idx
266            .search_filtered(&[1.0, 0.0, 0.0], 1, &[20, 30])
267            .await
268            .unwrap();
269        assert_eq!(hits.len(), 1);
270        assert_ne!(hits[0].id, 10);
271
272        idx.remove(&[10]).await.unwrap();
273        assert_eq!(idx.len().await.unwrap(), 2);
274
275        // cleanup
276        sqlx::query(&format!("DROP TABLE IF EXISTS {}", table))
277            .execute(&idx.pool)
278            .await
279            .ok();
280    }
281
282    #[test]
283    fn rejects_invalid_table_name() {
284        // valid identifiers accepted
285        assert!(validate_table_name("lk_test_1").is_ok());
286        assert!(validate_table_name("_vec").is_ok());
287        // rejected — would break out of identifier context
288        assert!(validate_table_name("").is_err());
289        assert!(validate_table_name("1bad").is_err());
290        assert!(validate_table_name("rm; DROP").is_err());
291        assert!(validate_table_name("weird\"name").is_err());
292        assert!(validate_table_name("sch.tbl").is_err());
293    }
294
295    #[test]
296    fn rejects_overflowing_id() {
297        assert_eq!(to_pg_id(0).unwrap(), 0);
298        assert_eq!(to_pg_id(42).unwrap(), 42);
299        assert_eq!(to_pg_id(i64::MAX as u64).unwrap(), i64::MAX);
300        assert!(to_pg_id((i64::MAX as u64) + 1).is_err());
301        assert!(to_pg_id(u64::MAX).is_err());
302    }
303}
llm_kernel/embedding/pgvector.rs

llm_kernel/embedding/
pgvector.rs
