diff --git a/Cargo.lock b/Cargo.lock index cfb8cb37..5f243ba7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5277,6 +5277,7 @@ dependencies = [ "anyhow", "chrono", "cron", + "futures-util", "pgvector", "serde", "serde_json", diff --git a/crates/utopia-server/src/main.rs b/crates/utopia-server/src/main.rs index 026c3d96..03b5c57f 100644 --- a/crates/utopia-server/src/main.rs +++ b/crates/utopia-server/src/main.rs @@ -460,6 +460,36 @@ async fn dispatch(st: &state::AppState, job: &utopia_store::jobs::Job) -> anyhow .ok_or_else(|| anyhow::anyhow!("payload 缺少 kb_id"))?; ontology_index::refresh(st, kb_id).await.map(|_| ()) } + // 向量索引（0035 / #512）：第一次写下某个维度的向量时排的，事务外 + // CONCURRENTLY 建。维度超过 HNSW 上限的到此为止，重试也建不出来 + "build_vector_index" => { + let target = job + .payload + .get("table") + .and_then(|v| v.as_str()) + .and_then(utopia_store::vector_index::Target::parse) + .ok_or_else(|| anyhow::anyhow!("payload 缺少 table"))?; + let dims = + job.payload + .get("dims") + .and_then(|v| v.as_u64()) + .ok_or_else(|| anyhow::anyhow!("payload 缺少 dims"))? as usize; + let built = utopia_store::vector_index::build(&st.pool, target, dims) + .await + .map_err(|e| match e { + utopia_core::AppError::Validation(_) => { + anyhow::Error::from(e).context(utopia_core::Terminal) + } + other => anyhow::Error::from(other), + })?; + tracing::info!( + index = %built.name, + created = built.created, + seconds = format!("{:.1}", built.seconds), + "向量索引就绪" + ); + Ok(()) + } "resolve_types" => { let kb_id: Uuid = job .payload diff --git a/crates/utopia-server/src/type_resolution.rs b/crates/utopia-server/src/type_resolution.rs index bb2c1770..cb90a6f1 100644 --- a/crates/utopia-server/src/type_resolution.rs +++ b/crates/utopia-server/src/type_resolution.rs @@ -180,6 +180,14 @@ async fn preview_with( slots.push((pi, slot)); } + // 邻居先一起取（#514）：六十个查询彼此无关，有界并发取回来，下面再按原顺序推理。 + // 后代集合按粗类记一批：同一个粗类的递归 CTE 一批里只发一次 + let ids: Vec = subjects.iter().map(|s| s.id).collect(); + let mut neighbour_lists = + utopia_store::resolution::nearest_typed_for_each(&state.pool, kb_id, &ids, NEIGHBOURS) + .await?; + let mut descendants_memo = utopia_store::resolution::DescendantsMemo::default(); + let mut out = Vec::with_capacity(subjects.len()); for (i, s) in subjects.iter().enumerate() { // 粗类的后代**排前面，但不是唯一能选的**。 @@ -193,13 +201,9 @@ async fn preview_with( // 该说不的是裁决那一步：它看得到描述、看得到粗类，能说出"这不是"。 // 还没有类的实体没有"粗类的后代"这个轴可用（0009），整张类表都是候选， // 排序就纯按检索顺序来 - let descendants: std::collections::HashSet<_> = match s.coarse_id { - Some(c) => utopia_store::resolution::descendants_of(&state.pool, kb_id, c) - .await? - .into_iter() - .collect(), - None => std::collections::HashSet::new(), - }; + let descendants = descendants_memo + .get(&state.pool, kb_id, s.coarse_id) + .await?; // **两路交替取，不按距离合并。** // // 距离在两路之间不可比：短查询（"医药集团"）产生的距离系统性地小于 @@ -242,10 +246,8 @@ async fn preview_with( // 剩下的顺序就是两路交替的检索序:检索决定端什么上去,裁决决定它是什么, // `crosses_axis` 决定要不要人看。一层一件事。 let candidates: Vec<_> = ranked.into_iter().take(CANDIDATES as usize).collect(); - // 第二路：语境相似的已定类实体，按类投票 - let raw = - utopia_store::resolution::nearest_typed_entities(&state.pool, kb_id, s.id, NEIGHBOURS) - .await?; + // 第二路：语境相似的已定类实体，按类投票（上面一起取回来的，这里按序拿） + let raw = std::mem::take(&mut neighbour_lists[i]); let mut votes: std::collections::BTreeMap, bool)> = std::collections::BTreeMap::new(); for (name, _tid, key, distance, same_doc) in raw { diff --git a/crates/utopia-store/Cargo.toml b/crates/utopia-store/Cargo.toml index fd10e4e9..2c81fb7e 100644 --- a/crates/utopia-store/Cargo.toml +++ b/crates/utopia-store/Cargo.toml @@ -10,6 +10,7 @@ utopia-ingest.workspace = true utopia-reason.workspace = true serde.workspace = true pgvector.workspace = true +futures-util.workspace = true cron.workspace = true sqlx.workspace = true tokio.workspace = true diff --git a/crates/utopia-store/src/documents.rs b/crates/utopia-store/src/documents.rs index f822f635..98f692c1 100644 --- a/crates/utopia-store/src/documents.rs +++ b/crates/utopia-store/src/documents.rs @@ -1210,10 +1210,16 @@ pub async fn set_embeddings(pool: &PgPool, items: &[(Uuid, Vec)]) -> AppRes .await?; } tx.commit().await?; + // 第一次写下这个维度的向量，索引就该排上了（0035）。在了的话这一句是一次查找 + if let Some((_, first)) = items.first() { + crate::vector_index::request(pool, crate::vector_index::Target::Chunks, first.len()) + .await?; + } Ok(()) } -/// 向量近邻检索（余弦距离，顺扫；P1 规模足够）。 +/// 向量近邻检索（余弦距离）。维度写成字面量、两侧 cast、`relaxed_order`——三条 +/// 规矩见 `vector_index`；走不走索引由规划器定 pub async fn vector_search( pool: &PgPool, kb_id: Uuid, @@ -1221,21 +1227,30 @@ pub async fn vector_search( limit: i64, as_of: Option>, ) -> AppResult> { + let dims = embedding.len(); + if dims == 0 { + return Ok(Vec::new()); + } let query_vec = Vector::from(embedding.to_vec()); + let mut tx = pool.begin().await?; + crate::vector_index::relaxed_order(pool, &mut tx).await?; let rows: Vec<(Uuid,)> = sqlx::query_as(&format!( "SELECT id FROM chunks c WHERE c.kb_id = $1 AND c.embedding IS NOT NULL AND {live} - AND vector_dims(c.embedding) = vector_dims($2) - ORDER BY c.embedding <=> $2 + AND {same_dims} + ORDER BY {distance} LIMIT $3", live = crate::record_axis::chunk_live_at("c", 4), + same_dims = crate::vector_index::same_dims("c.embedding", dims), + distance = crate::vector_index::distance("c.embedding", 2, dims), )) .bind(kb_id) .bind(&query_vec) .bind(limit) .bind(as_of) - .fetch_all(pool) + .fetch_all(&mut *tx) .await?; + tx.commit().await?; Ok(rows.into_iter().map(|(id,)| id).collect()) } diff --git a/crates/utopia-store/src/lib.rs b/crates/utopia-store/src/lib.rs index 332c6a06..e582fed0 100644 --- a/crates/utopia-store/src/lib.rs +++ b/crates/utopia-store/src/lib.rs @@ -36,5 +36,6 @@ pub mod sources; pub mod temporal; pub mod test_db; pub mod tokens; +pub mod vector_index; pub mod workspaces; pub mod world_axis; diff --git a/crates/utopia-store/src/resolution.rs b/crates/utopia-store/src/resolution.rs index e2ca12f1..f29ad63a 100644 --- a/crates/utopia-store/src/resolution.rs +++ b/crates/utopia-store/src/resolution.rs @@ -965,6 +965,7 @@ async fn update_profile(pool: &PgPool, id: Uuid, n: i32, ctx: &[f32]) -> AppResu } _ => (ctx.to_vec(), 1), }; + let dims = new_vec.len(); sqlx::query( "UPDATE entities SET profile_embedding = $2, profile_n = $3, updated_at = now() WHERE id = $1", @@ -974,6 +975,8 @@ async fn update_profile(pool: &PgPool, id: Uuid, n: i32, ctx: &[f32]) -> AppResu .bind(new_n) .execute(pool) .await?; + // 画像表也要索引（0035 / #514）：类型消解按主语逐个扫它。在了的话这一句是一次查找 + crate::vector_index::request(pool, crate::vector_index::Target::EntityProfiles, dims).await?; Ok(()) } @@ -2480,23 +2483,78 @@ pub async fn set_specific_type(pool: &PgPool, entity_id: Uuid, value: &str) -> A /// /// 已知弱点：只出现在一篇文档里的实体，语境向量就是那一块的向量，同文档的实体 /// 会互相成为近邻。调用方要看得到 `same_document`，别把它当成类型证据。 +/// 一批之内 [`descendants_of`] 的记忆（#514）。 +/// +/// 粗类来自抽取的小词表（person、organization、product 加几个），六十个主语里 +/// 同一个 `coarse_id` 反复出现，同一个递归 CTE 就反复发。批内本体不动，同一输入 +/// 同一结果，记住不改答案。**没有粗类的主语不进这张表**：它没有「后代」这个轴 +/// （0009），整张类表都是候选；用 `Option` 当键会把它和某个真实的类混在一起 +#[derive(Default)] +pub struct DescendantsMemo { + sets: std::collections::HashMap>, +} + +impl DescendantsMemo { + pub async fn get( + &mut self, + pool: &PgPool, + kb_id: Uuid, + root: Option, + ) -> AppResult> { + let Some(root) = root else { + return Ok(HashSet::new()); + }; + if let Some(set) = self.sets.get(&root) { + return Ok(set.clone()); + } + let set: HashSet = descendants_of(pool, kb_id, root) + .await? + .into_iter() + .collect(); + self.sets.insert(root, set.clone()); + Ok(set) + } + + /// 记住了几个根 + pub fn len(&self) -> usize { + self.sets.len() + } + + pub fn is_empty(&self) -> bool { + self.sets.is_empty() + } +} + +/// 一个主语的近邻：语境相似的已定类实体，连同「是否同一篇文档」。 +/// +/// 主语自己的向量先取出来再查：维度要写进 SQL（`vector_index` 规矩 1），SQL 里的 +/// 子查询给不了这个数。没有向量的主语没有近邻，回空 pub async fn nearest_typed_entities( pool: &PgPool, kb_id: Uuid, entity_id: Uuid, limit: i64, ) -> AppResult> { - Ok(sqlx::query_as( - "WITH me AS ( - SELECT profile_embedding AS v FROM entities WHERE id = $2 AND kb_id = $1 - ), - my_docs AS ( + let me: Option<(Option,)> = + sqlx::query_as("SELECT profile_embedding FROM entities WHERE id = $2 AND kb_id = $1") + .bind(kb_id) + .bind(entity_id) + .fetch_optional(pool) + .await?; + let Some(query_vec) = me.and_then(|(v,)| v) else { + return Ok(Vec::new()); + }; + let dims = query_vec.as_slice().len(); + let mut tx = pool.begin().await?; + crate::vector_index::relaxed_order(pool, &mut tx).await?; + let rows = sqlx::query_as(&format!( + "WITH my_docs AS ( SELECT DISTINCT ev.document_id FROM fact_evidence ev JOIN facts f ON f.id = ev.fact_id WHERE f.kb_id = $1 AND (f.subject_id = $2 OR f.object_id = $2) ) SELECT e.canonical_name, t.id, t.key, - (e.profile_embedding <=> (SELECT v FROM me))::float8 AS distance, + ({distance})::float8 AS distance, EXISTS (SELECT 1 FROM fact_evidence ev2 JOIN facts f2 ON f2.id = ev2.fact_id WHERE f2.kb_id = $1 AND (f2.subject_id = e.id OR f2.object_id = e.id) @@ -2507,16 +2565,54 @@ pub async fn nearest_typed_entities( FROM entities e JOIN entity_types t ON t.id = e.type_id WHERE e.kb_id = $1 AND e.merged_into IS NULL AND e.id <> $2 - AND e.profile_embedding IS NOT NULL - AND (SELECT v FROM me) IS NOT NULL - ORDER BY e.profile_embedding <=> (SELECT v FROM me) - LIMIT $3", - ) + AND e.profile_embedding IS NOT NULL AND {same_dims} + ORDER BY {distance} + LIMIT $4", + distance = crate::vector_index::distance("e.profile_embedding", 3, dims), + same_dims = crate::vector_index::same_dims("e.profile_embedding", dims), + )) .bind(kb_id) .bind(entity_id) + .bind(&query_vec) .bind(limit) - .fetch_all(pool) - .await?) + .fetch_all(&mut *tx) + .await?; + tx.commit().await?; + Ok(rows) +} + +/// 同时跑几个近邻查询。池是 32、按「同时跑的短查询」定的（`db.rs`）；六十个 +/// 全表扫一起上就是那里记的池子被吃空的样子——慢请求和超时，没有一句话说池小了。 +/// 八个留足了给请求的余量，而 HNSW 就位之后每个查询只有几毫秒，再高也没意义 +pub const NEIGHBOUR_SCANS: usize = 8; + +/// 一批主语各自的近邻，**按送进来的顺序回**：下游裁决按这个顺序读。 +/// +/// 六十个查询彼此无关，串行只因为循环是串行的（#514）。这里有界并发地取， +/// `buffered` 保序，推理仍在原顺序上做 +pub async fn nearest_typed_for_each( + pool: &PgPool, + kb_id: Uuid, + ids: &[Uuid], + limit: i64, +) -> AppResult>> { + nearest_typed_for_each_with(pool, kb_id, ids, limit, NEIGHBOUR_SCANS).await +} + +/// 同上，并发上限由调用方给（测试用它证明上限不改答案） +pub async fn nearest_typed_for_each_with( + pool: &PgPool, + kb_id: Uuid, + ids: &[Uuid], + limit: i64, + at_once: usize, +) -> AppResult>> { + use futures_util::{stream, StreamExt, TryStreamExt}; + stream::iter(ids.iter().copied()) + .map(|id| nearest_typed_entities(pool, kb_id, id, limit)) + .buffered(at_once.max(1)) + .try_collect() + .await } /// 按实体逐个改类，写进同一本账。返回 (批次 id, 改动数)。 diff --git a/crates/utopia-store/src/vector_index.rs b/crates/utopia-store/src/vector_index.rs new file mode 100644 index 00000000..0c5ec49f --- /dev/null +++ b/crates/utopia-store/src/vector_index.rs @@ -0,0 +1,293 @@ +//! 向量索引由任务来建，不是一条编号迁移（0035）。 +//! +//! 列是无维的（`embedding vector`，没有 `(N)`）：维度跟着工作区选的嵌入模型走， +//! 迁移跑的时候还不知道，pgvector 也建不了未知维度的 HNSW。另一头， +//! `CREATE INDEX CONCURRENTLY` 进不了 sqlx 给迁移套的事务，而不带 CONCURRENTLY +//! 就要在 `chunks` 上持 ACCESS EXCLUSIVE 锁到建完。所以：第一次写下某个维度的 +//! 向量时排一条任务，任务在事务外建一个**按维度的部分表达式索引**， +//! `IF NOT EXISTS` 让重排无害。 +//! +//! 读路径上有三条规矩，缺一条索引就悄悄不生效或悄悄少回： +//! 1. **维度写成字面量。** `vector_dims(col) = $2` 绑参数时自定义计划能用索引， +//! generic plan 退回顺扫；sqlx 的预处理语句跑五次就切 generic，第六次起索引 +//! 悄悄失效。用 [`same_dims`] / [`distance`] 把整数写进 SQL。 +//! 2. **两侧都 cast 到 `vector(N)`。** 索引建在表达式 `col::vector(N)` 上， +//! `ORDER BY` 必须一字不差地写同一个表达式。 +//! 3. **`hnsw.iterative_scan = relaxed_order`。** HNSW 先取 `ef_search` 个候选再过 +//! WHERE；一张按 kb 分租的表上，小库的行在候选里占不到几个，`LIMIT 10` 会回 +//! 三行甚至零行（实测 1 万行的库在 6 万行的表上：关着回 3/24，开着回满）。 +//! iterative_scan 让它继续往下走到凑够为止。它有两个停下来的条件，都在这里放宽： +//! `hnsw.scan_mem_multiplier` 从 1 提到 4（默认 work_mem 4 MB 时即 16 MB）——实测 +//! 这一个才是绑住它的：真实库各占索引 1%、旁边一个 5 万行的合成租户、强制走索引， +//! 倍数 1 时 522 问里 154 问回不满、recall 0.705、每问 32 ms；倍数 4 时全部回满、 +//! recall 1.0、每问 74 ms，16 与 4 无异。`hnsw.max_scan_tuples` 从 20,000 提到 +//! 100,000，实测里它没绑住，提的理由是把「到顶」留成看得见的慢，而不是悄悄少回。 +//! +//! 走不走索引由规划器定：有 `chunks_kb_idx` 时小库、中库它自己选精确路径，只有 +//! 占表大头的库才走 HNSW（实测 6 万行：20 行和 1 万行的库走精确，5 万的走索引）。 +//! 应用侧不设阈值——阈值是对规划器的猜测，猜错了两边都慢。 + +use sqlx::{Executor, PgPool, Postgres, Transaction}; +use std::collections::HashSet; +use std::sync::{Mutex, OnceLock}; +use utopia_core::{AppError, AppResult}; + +/// 任务种类，`main.rs` 的分发按这个名字认 +pub const JOB_KIND: &str = "build_vector_index"; + +/// pgvector 的 HNSW 对 `vector` 类型的上限。超过的维度（text-embedding-3-large +/// 是 3072）不建索引，查询照常走精确路径。`halfvec` 能到 4000，但那是另一种 +/// 精度，等有人用到再说 +pub const MAX_DIMS: usize = 2000; + +/// 哪一列。两张表同一套机制，名字不同而已 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Target { + /// `chunks.embedding`：文档分块，会无界增长的那张 + Chunks, + /// `entities.profile_embedding`：实体画像，类型消解按主语逐个扫它（#514） + EntityProfiles, +} + +impl Target { + pub fn table(self) -> &'static str { + match self { + Target::Chunks => "chunks", + Target::EntityProfiles => "entities", + } + } + + pub fn column(self) -> &'static str { + match self { + Target::Chunks => "embedding", + Target::EntityProfiles => "profile_embedding", + } + } + + /// 任务载荷里的名字 + pub fn key(self) -> &'static str { + self.table() + } + + pub fn parse(key: &str) -> Option { + match key { + "chunks" => Some(Target::Chunks), + "entities" => Some(Target::EntityProfiles), + _ => None, + } + } +} + +/// 索引名：`chunks_embedding_hnsw_1024` +pub fn index_name(target: Target, dims: usize) -> String { + format!("{}_{}_hnsw_{dims}", target.table(), target.column()) +} + +/// 部分索引的谓词；查询里要原样出现（规矩 1） +pub fn same_dims(column: &str, dims: usize) -> String { + format!("vector_dims({column}) = {dims}") +} + +/// `<=>` 两侧都 cast 到字面维度（规矩 1、2）；`param` 是查询向量的参数号 +pub fn distance(column: &str, param: usize, dims: usize) -> String { + format!("{column}::vector({dims}) <=> ${param}::vector({dims})") +} + +/// 索引现在的状态：`None` 没有；`Some(valid)` 有，`false` 是上次建到一半留下的 +pub async fn status(pool: &PgPool, target: Target, dims: usize) -> AppResult> { + let row: Option<(bool,)> = sqlx::query_as( + "SELECT i.indisvalid FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + WHERE c.relname = $1 AND c.relkind = 'i'", + ) + .bind(index_name(target, dims)) + .fetch_optional(pool) + .await?; + Ok(row.map(|(v,)| v)) +} + +/// 进程里记住的「已经在了」的索引：往后每次写入只是一次查找 +fn known() -> &'static Mutex> { + static KNOWN: OnceLock>> = OnceLock::new(); + KNOWN.get_or_init(|| Mutex::new(HashSet::new())) +} + +fn remember(name: &str) { + known().lock().unwrap().insert(name.to_string()); +} + +fn forget(name: &str) { + known().lock().unwrap().remove(name); +} + +fn is_known(name: &str) -> bool { + known().lock().unwrap().contains(name) +} + +/// 写下某个维度的向量时叫一声：索引不在就排一条建索引任务。 +/// 返回 `Some(job id)` = 这一次排上了。 +/// +/// 索引在的话记在进程里，往后只是一次 HashSet 查找；不在的时候每次写入一次目录 +/// 查询加一次「没排着才排」的插入——建索引那一两分钟里写入照常，代价可忽略。 +/// 超过 [`MAX_DIMS`] 的维度什么都不排：建不了，查询走精确路径 +pub async fn request(pool: &PgPool, target: Target, dims: usize) -> AppResult> { + if dims == 0 || dims > MAX_DIMS { + return Ok(None); + } + let name = index_name(target, dims); + if is_known(&name) { + return Ok(None); + } + if status(pool, target, dims).await? == Some(true) { + remember(&name); + return Ok(None); + } + crate::jobs::enqueue_unless_queued( + pool, + JOB_KIND, + serde_json::json!({ "table": target.key(), "dims": dims }), + ) + .await +} + +/// 一次 [`build`] 的结果 +#[derive(Debug)] +pub struct Built { + pub name: String, + /// `false` = 本来就在 + pub created: bool, + pub seconds: f64, +} + +/// 任务本体。**事务外**（CONCURRENTLY 的要求）、**串行**（并行建索引要共享内存， +/// Docker 默认 64 MB 的 `/dev/shm` 会让它报 could not resize shared memory segment； +/// 串行 6 万行 1024 维约 90 秒，到处能跑）。上一次建到一半留下的无效索引先删—— +/// `IF NOT EXISTS` 看见它会以为已经建好。 +pub async fn build(pool: &PgPool, target: Target, dims: usize) -> AppResult { + if dims == 0 || dims > MAX_DIMS { + return Err(AppError::Validation(format!( + "{dims} dims: HNSW on vector holds up to {MAX_DIMS}" + ))); + } + let name = index_name(target, dims); + let started = std::time::Instant::now(); + let mut conn = pool.acquire().await?; + let outcome = async { + conn.execute("SET max_parallel_maintenance_workers = 0") + .await?; + // 默认 64 MB 的 maintenance_work_mem 到一万四千行 1024 维就装不下图，之后 + // 每一行都要落盘再读，5 万行建了 4 分 20 秒；512 MB 只在建的这一会儿占着 + conn.execute("SET maintenance_work_mem = '512MB'").await?; + let before = status(pool, target, dims).await?; + if before == Some(false) { + conn.execute(format!("DROP INDEX CONCURRENTLY IF EXISTS {name}").as_str()) + .await?; + } + let existed = before == Some(true); + conn.execute( + format!( + "CREATE INDEX CONCURRENTLY IF NOT EXISTS {name} ON {table} \ + USING hnsw (({col}::vector({dims})) vector_cosine_ops) WHERE {pred}", + table = target.table(), + col = target.column(), + pred = same_dims(target.column(), dims), + ) + .as_str(), + ) + .await?; + Ok::(!existed) + } + .await; + // 会话级 SET 跟着连接回池，成败都复位 + let _ = conn.execute("RESET max_parallel_maintenance_workers").await; + let _ = conn.execute("RESET maintenance_work_mem").await; + let created = outcome?; + remember(&name); + Ok(Built { + name, + created, + seconds: started.elapsed().as_secs_f64(), + }) +} + +/// 删掉（测试与手工维护用；写路径不会走到这里） +pub async fn drop(pool: &PgPool, target: Target, dims: usize) -> AppResult<()> { + let name = index_name(target, dims); + forget(&name); + let mut conn = pool.acquire().await?; + conn.execute(format!("DROP INDEX CONCURRENTLY IF EXISTS {name}").as_str()) + .await?; + Ok(()) +} + +/// `hnsw.iterative_scan` 是 pgvector 0.8 才有的参数。旧版本第一次探一下、进程内 +/// 记住：探不到就不设——查询还是对的，只是共享表上的小库可能少回几行， +/// 那正是 0.8 修的事 +async fn iterative_scan_available(pool: &PgPool) -> bool { + static AVAILABLE: OnceLock = OnceLock::new(); + if let Some(v) = AVAILABLE.get() { + return *v; + } + let seen: Result, _> = + sqlx::query_scalar("SELECT current_setting('hnsw.iterative_scan', true)") + .fetch_one(pool) + .await; + let ok = matches!(seen, Ok(Some(_))); + let _ = AVAILABLE.set(ok); + ok +} + +/// 读路径的会话设置（规矩 3）。`SET LOCAL` 只在事务里生效，所以近邻查询都套一个事务 +pub async fn relaxed_order(pool: &PgPool, tx: &mut Transaction<'_, Postgres>) -> AppResult<()> { + if iterative_scan_available(pool).await { + (&mut **tx) + .execute("SET LOCAL hnsw.iterative_scan = relaxed_order") + .await?; + // 三个参数同一个版本来的（0.8）。两个停止条件见模块注释：内存那个是实测绑住 + // 扫描的，元组上限那个没绑住，放宽是为了把「到顶」留成看得见的慢 + (&mut **tx) + .execute("SET LOCAL hnsw.scan_mem_multiplier = 4") + .await?; + (&mut **tx) + .execute("SET LOCAL hnsw.max_scan_tuples = 100000") + .await?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_index_name_carries_table_column_and_dims() { + assert_eq!( + index_name(Target::Chunks, 1024), + "chunks_embedding_hnsw_1024" + ); + assert_eq!( + index_name(Target::EntityProfiles, 768), + "entities_profile_embedding_hnsw_768" + ); + } + + #[test] + fn the_sql_writes_the_dimension_as_a_literal() { + assert_eq!( + same_dims("c.embedding", 1024), + "vector_dims(c.embedding) = 1024" + ); + assert_eq!( + distance("c.embedding", 2, 1024), + "c.embedding::vector(1024) <=> $2::vector(1024)" + ); + } + + #[test] + fn a_target_round_trips_through_the_payload() { + for t in [Target::Chunks, Target::EntityProfiles] { + assert_eq!(Target::parse(t.key()), Some(t)); + } + assert_eq!(Target::parse("documents"), None); + } +} diff --git a/crates/utopia-store/tests/a_batch_gathers_its_neighbours_in_order.rs b/crates/utopia-store/tests/a_batch_gathers_its_neighbours_in_order.rs new file mode 100644 index 00000000..b138fd9b --- /dev/null +++ b/crates/utopia-store/tests/a_batch_gathers_its_neighbours_in_order.rs @@ -0,0 +1,352 @@ +//! 类型消解的近邻那一路（#514），打在真库上。 +//! +//! 规矩是：消解的答案是台账的函数，与候选按什么顺序、什么时机取回来无关。 +//! 所以并发地取只有在「同一个库给出同一批邻居」时才算安全。这里钉四样： +//! - **顺序**：一批主语的邻居按送进去的顺序回来，下游裁决按这个顺序读 +//! - **并发上限不改答案**：上限 1 与上限 8 给出同一份结果；有没有索引也一样 +//! - **SQL 里的三道门**：自己不是自己的邻居；没判出类型的不当证据；同一篇文档的要标出来 +//! - **后代集合的记忆**：一批之内同一个粗类只问一次，没有粗类的不进表 +//! +//! 还有一条只有连库才验得出：两个连接的小池子也跑得完一整批，不会撞上取连接的超时。 + +use pgvector::Vector; +use sqlx::postgres::PgPoolOptions; +use sqlx::PgPool; +use std::collections::HashSet; +use utopia_store::resolution::{self, DescendantsMemo}; +use utopia_store::vector_index::{self, Target}; +use uuid::Uuid; + +/// 这个测试独占的维度：其它连库测试用 3 维 +const DIMS: usize = 5; +const TYPED: usize = 40; + +fn lcg(seed: &mut u64) -> f32 { + *seed = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((*seed >> 33) as f32 / (1u64 << 31) as f32) - 0.5 +} + +struct Fixture { + org: Uuid, + kb: Uuid, + thing: Uuid, + person: Uuid, + company: Uuid, + employee: Uuid, + /// 三个待消解的主语 + subjects: Vec, + /// 自己的向量是 3 维的主语 + three_dims_subject: Uuid, +} + +async fn seed(pool: &PgPool) -> anyhow::Result { + let (org, ws, kb) = (Uuid::now_v7(), Uuid::now_v7(), Uuid::now_v7()); + let (thing, person, company, employee) = ( + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + ); + let tag = Uuid::now_v7(); + + sqlx::query("INSERT INTO organizations (id, name) VALUES ($1, 'gather-test')") + .bind(org) + .execute(pool) + .await?; + sqlx::query("INSERT INTO workspaces (id, org_id, name) VALUES ($1, $2, 'gather-test')") + .bind(ws) + .bind(org) + .execute(pool) + .await?; + sqlx::query( + "INSERT INTO knowledge_bases (id, workspace_id, name) VALUES ($1, $2, 'gather-test')", + ) + .bind(kb) + .bind(ws) + .execute(pool) + .await?; + for (id, key) in [ + (thing, "thing"), + (person, "person"), + (company, "company"), + (employee, "employee"), + ] { + sqlx::query("INSERT INTO entity_types (id, kb_id, key, label) VALUES ($1, $2, $3, $3)") + .bind(id) + .bind(kb) + .bind(key) + .execute(pool) + .await?; + } + for (child, parent) in [(person, thing), (company, thing), (employee, person)] { + sqlx::query("INSERT INTO entity_type_parents (child_id, parent_id) VALUES ($1, $2)") + .bind(child) + .bind(parent) + .execute(pool) + .await?; + } + + let mut seed = 0xfeed_u64; + let mut vector = |dims: usize| -> Vector { + Vector::from((0..dims).map(|_| lcg(&mut seed)).collect::>()) + }; + let insert = |id: Uuid, name: String, type_id: Option, v: Vector| { + let pool = pool.clone(); + async move { + sqlx::query( + "INSERT INTO entities (id, kb_id, type_id, canonical_name, profile_embedding, profile_n) + VALUES ($1, $2, $3, $4, $5, 1)", + ) + .bind(id) + .bind(kb) + .bind(type_id) + .bind(name) + .bind(v) + .execute(&pool) + .await + } + }; + let types = [person, company, employee]; + let mut typed = Vec::with_capacity(TYPED); + for i in 0..TYPED { + let id = Uuid::now_v7(); + insert(id, format!("typed {i}"), Some(types[i % 3]), vector(DIMS)).await?; + typed.push(id); + } + let subjects: Vec = (0..3).map(|_| Uuid::now_v7()).collect(); + for (i, s) in subjects.iter().enumerate() { + insert(*s, format!("subject {i}"), Some(thing), vector(DIMS)).await?; + } + let untyped = Uuid::now_v7(); + insert(untyped, "untyped".into(), None, vector(DIMS)).await?; + let three_dims = Uuid::now_v7(); + insert(three_dims, "three dims".into(), Some(person), vector(3)).await?; + let three_dims_subject = Uuid::now_v7(); + insert( + three_dims_subject, + "three dims subject".into(), + Some(thing), + vector(3), + ) + .await?; + + // 主语 0 和 typed[0] 在同一篇文档里：一条事实，证据落在那篇文档的一块上 + let same_doc_neighbour = typed[0]; + let (doc, chunk, fact) = (Uuid::now_v7(), Uuid::now_v7(), Uuid::now_v7()); + sqlx::query( + "INSERT INTO documents (id, kb_id, filename, sha256, status) + VALUES ($1, $2, 'shared.md', $3, 'ready')", + ) + .bind(doc) + .bind(kb) + .bind(format!("sha-{tag}")) + .execute(pool) + .await?; + sqlx::query( + "INSERT INTO chunks (id, kb_id, document_id, seq, text) VALUES ($1, $2, $3, 0, 'a sentence')", + ) + .bind(chunk) + .bind(kb) + .bind(doc) + .execute(pool) + .await?; + sqlx::query("INSERT INTO facts (id, kb_id, subject_id, object_id) VALUES ($1, $2, $3, $4)") + .bind(fact) + .bind(kb) + .bind(subjects[0]) + .bind(same_doc_neighbour) + .execute(pool) + .await?; + sqlx::query( + "INSERT INTO fact_evidence (fact_id, chunk_id, quote, document_id) + VALUES ($1, $2, 'a sentence', $3)", + ) + .bind(fact) + .bind(chunk) + .bind(doc) + .execute(pool) + .await?; + + Ok(Fixture { + org, + kb, + thing, + person, + company, + employee, + subjects, + three_dims_subject, + }) +} + +type Neighbours = Vec<(String, Uuid, String, f64, bool)>; + +fn names(list: &Neighbours) -> Vec<&str> { + list.iter().map(|(n, ..)| n.as_str()).collect() +} + +#[tokio::test] +async fn the_neighbours_come_back_in_order_and_the_gates_hold() -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = PgPool::connect(&url).await?; + let f = seed(&pool).await?; + + let run = async { + vector_index::drop(&pool, Target::EntityProfiles, DIMS).await?; + + // ---- 一、逐个问，作为基准 + let mut serial: Vec = Vec::new(); + for s in &f.subjects { + serial.push(resolution::nearest_typed_entities(&pool, f.kb, *s, 10).await?); + } + for (i, list) in serial.iter().enumerate() { + assert_eq!(list.len(), 10, "主语 {i} 该有十个邻居"); + let me = format!("subject {i}"); + assert!( + list.iter().all(|(n, ..)| *n != me), + "自己不是自己的邻居（别的主语可以是：它们挂着 thing，是已定类的）" + ); + assert!( + list.iter().all(|(n, ..)| n != "untyped"), + "没判出类型的不当证据" + ); + assert!( + list.iter().all(|(n, ..)| n != "three dims"), + "另一维度的画像不是候选" + ); + } + // 主语 0 的邻居里，typed 0 与它同一篇文档，别的都不是 + let flagged: Vec<&str> = serial[0] + .iter() + .filter(|(_, _, _, _, same)| *same) + .map(|(n, ..)| n.as_str()) + .collect(); + let has_typed0 = serial[0].iter().any(|(n, ..)| n == "typed 0"); + if has_typed0 { + assert_eq!(flagged, vec!["typed 0"], "同一篇文档的要标出来，且只标它"); + } else { + assert!(flagged.is_empty(), "不在邻居里就没什么可标"); + } + // 自己的向量是 3 维的主语：没有 3 维的已定类邻居可比……除了那一条 3 维的 + let odd = resolution::nearest_typed_entities(&pool, f.kb, f.three_dims_subject, 10).await?; + assert_eq!( + names(&odd), + vec!["three dims"], + "3 维的查询只找 3 维的，也不报错" + ); + + // ---- 二、一批取：顺序与上限 + let gathered_1 = + resolution::nearest_typed_for_each_with(&pool, f.kb, &f.subjects, 10, 1).await?; + let gathered_8 = + resolution::nearest_typed_for_each_with(&pool, f.kb, &f.subjects, 10, 8).await?; + assert_eq!(gathered_1.len(), f.subjects.len()); + for i in 0..f.subjects.len() { + assert_eq!( + names(&gathered_1[i]), + names(&serial[i]), + "上限 1：第 {i} 个位置是第 {i} 个主语的邻居" + ); + assert_eq!( + names(&gathered_8[i]), + names(&serial[i]), + "上限 8：第 {i} 个位置是第 {i} 个主语的邻居" + ); + } + + // ---- 三、两个连接的池子也跑得完六十个 + let tiny = PgPoolOptions::new() + .max_connections(2) + .acquire_timeout(std::time::Duration::from_secs(10)) + .connect(&url) + .await?; + let sixty: Vec = f.subjects.iter().cycle().take(60).copied().collect(); + let over_tiny = resolution::nearest_typed_for_each_with(&tiny, f.kb, &sixty, 10, 8).await?; + assert_eq!(over_tiny.len(), 60); + for (i, list) in over_tiny.iter().enumerate() { + assert_eq!(names(list), names(&serial[i % 3])); + } + tiny.close().await; + + // ---- 四、索引就位之后同一份答案 + let built = vector_index::build(&pool, Target::EntityProfiles, DIMS).await?; + assert!(built.created); + assert_eq!(built.name, "entities_profile_embedding_hnsw_5"); + let indexed = resolution::nearest_typed_for_each(&pool, f.kb, &f.subjects, 10).await?; + for i in 0..f.subjects.len() { + assert_eq!( + names(&indexed[i]), + names(&serial[i]), + "索引改了主语 {i} 的邻居" + ); + } + + // ---- 五、后代集合按粗类记一批 + let mut memo = DescendantsMemo::default(); + let direct = |root: Uuid| { + let pool = pool.clone(); + async move { + Ok::, anyhow::Error>( + resolution::descendants_of(&pool, f.kb, root) + .await? + .into_iter() + .collect(), + ) + } + }; + assert_eq!( + memo.get(&pool, f.kb, Some(f.thing)).await?, + direct(f.thing).await? + ); + assert_eq!( + memo.get(&pool, f.kb, Some(f.person)).await?, + [f.person, f.employee].into_iter().collect::>() + ); + assert_eq!( + memo.get(&pool, f.kb, Some(f.company)).await?, + direct(f.company).await? + ); + assert_eq!(memo.len(), 3); + assert!( + memo.get(&pool, f.kb, None).await?.is_empty(), + "没有粗类：没有后代这个轴" + ); + assert_eq!(memo.len(), 3, "没有粗类的不进表"); + // 批内本体动了：记忆给的还是问过的那一份，直接问给的是新的——这就是「记住」 + let intern = Uuid::now_v7(); + sqlx::query( + "INSERT INTO entity_types (id, kb_id, key, label) VALUES ($1, $2, 'intern', 'intern')", + ) + .bind(intern) + .bind(f.kb) + .execute(&pool) + .await?; + sqlx::query("INSERT INTO entity_type_parents (child_id, parent_id) VALUES ($1, $2)") + .bind(intern) + .bind(f.person) + .execute(&pool) + .await?; + assert!(!memo + .get(&pool, f.kb, Some(f.person)) + .await? + .contains(&intern)); + assert!(direct(f.person).await?.contains(&intern)); + Ok::<_, anyhow::Error>(()) + } + .await; + + let _ = vector_index::drop(&pool, Target::EntityProfiles, DIMS).await; + sqlx::query("DELETE FROM knowledge_bases WHERE id = $1") + .bind(f.kb) + .execute(&pool) + .await?; + sqlx::query("DELETE FROM organizations WHERE id = $1") + .bind(f.org) + .execute(&pool) + .await?; + run +} diff --git a/crates/utopia-store/tests/a_vector_index_is_built_by_a_job.rs b/crates/utopia-store/tests/a_vector_index_is_built_by_a_job.rs new file mode 100644 index 00000000..04d9dfb3 --- /dev/null +++ b/crates/utopia-store/tests/a_vector_index_is_built_by_a_job.rs @@ -0,0 +1,90 @@ +//! 索引由任务建（0035）：写入一侧只负责「排一条」，且排一次就够。 +//! +//! - 第一次写下某个维度：排上一条 `build_vector_index` +//! - 再写：不再排（队列里已经有一条） +//! - 建好之后再写：什么都不排（进程里记住了） +//! - 超过 HNSW 上限的维度：不排，也建不了——这种失败不该重试 +//! +//! 建索引本身与查询的一致性在 `the_nearest_chunk_is_found_however_it_is_reached` 里。 + +use sqlx::PgPool; +use utopia_store::vector_index::{self, Target, JOB_KIND, MAX_DIMS}; + +/// 这个测试独占的维度 +const DIMS: usize = 9; + +async fn queued(pool: &PgPool, dims: usize) -> anyhow::Result { + Ok(sqlx::query_scalar( + "SELECT count(*) FROM jobs WHERE kind = $1 AND payload = $2 AND status = 'queued'", + ) + .bind(JOB_KIND) + .bind(serde_json::json!({ "table": "chunks", "dims": dims })) + .fetch_one(pool) + .await?) +} + +#[tokio::test] +async fn the_first_write_of_a_dimension_queues_one_build() -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = PgPool::connect(&url).await?; + // 上一次跑留下的都清掉，从「什么都没有」开始 + vector_index::drop(&pool, Target::Chunks, DIMS).await?; + sqlx::query("DELETE FROM jobs WHERE kind = $1 AND payload->>'dims' = $2") + .bind(JOB_KIND) + .bind(DIMS.to_string()) + .execute(&pool) + .await?; + + let run = async { + let first = vector_index::request(&pool, Target::Chunks, DIMS).await?; + assert!(first.is_some(), "第一次写下 9 维：排上一条"); + assert_eq!(queued(&pool, DIMS).await?, 1); + let second = vector_index::request(&pool, Target::Chunks, DIMS).await?; + assert_eq!(second, None, "队列里已经有一条，不再排"); + assert_eq!(queued(&pool, DIMS).await?, 1); + + let built = vector_index::build(&pool, Target::Chunks, DIMS).await?; + assert!(built.created); + assert_eq!( + vector_index::status(&pool, Target::Chunks, DIMS).await?, + Some(true) + ); + // 任务跑完那条会被标 done；这里模拟一下，好让下一问只靠「进程里记住了」 + sqlx::query("UPDATE jobs SET status = 'done' WHERE kind = $1 AND payload->>'dims' = $2") + .bind(JOB_KIND) + .bind(DIMS.to_string()) + .execute(&pool) + .await?; + let third = vector_index::request(&pool, Target::Chunks, DIMS).await?; + assert_eq!(third, None, "建好了：什么都不排"); + assert_eq!(queued(&pool, DIMS).await?, 0); + + // 超过上限的维度 + let too_wide = MAX_DIMS + 1; + assert_eq!( + vector_index::request(&pool, Target::Chunks, too_wide).await?, + None, + "建不了的不排" + ); + assert_eq!(queued(&pool, too_wide).await?, 0); + let err = vector_index::build(&pool, Target::Chunks, too_wide) + .await + .expect_err("HNSW on vector holds up to 2000 dims"); + assert!( + matches!(err, utopia_core::AppError::Validation(_)), + "是 Validation，分发层据此判 Terminal：{err:?}" + ); + Ok::<_, anyhow::Error>(()) + } + .await; + + let _ = vector_index::drop(&pool, Target::Chunks, DIMS).await; + sqlx::query("DELETE FROM jobs WHERE kind = $1 AND payload->>'dims' = $2") + .bind(JOB_KIND) + .bind(DIMS.to_string()) + .execute(&pool) + .await?; + run +} diff --git a/crates/utopia-store/tests/the_nearest_chunk_is_found_however_it_is_reached.rs b/crates/utopia-store/tests/the_nearest_chunk_is_found_however_it_is_reached.rs new file mode 100644 index 00000000..e31558ba --- /dev/null +++ b/crates/utopia-store/tests/the_nearest_chunk_is_found_however_it_is_reached.rs @@ -0,0 +1,295 @@ +//! 向量召回的答案不随计划变（0035 / #512），打在真库上。 +//! +//! 索引对使用者是看不见的：契约是召回和记录轴。所以这里每一问都问两遍—— +//! 没有索引时一遍，`vector_index::build` 之后再一遍——两遍必须一致。走不走索引 +//! 由规划器定，这个测试不假装能替它选；它钉的是**答案**，改了答案就红，只改计划就静。 +//! +//! 六个问题各自会以不同的方式坏： +//! - 最近的那一条要回来（表达式索引与 `ORDER BY` 写得不一样，索引就悄悄不生效） +//! - 大库旁边的小库 `LIMIT 10` 要回满（HNSW 后置过滤会把小库回成零行，`relaxed_order` 救的正是它） +//! - 不串库（租户，不是性能） +//! - 另一维度的分块不是候选，也不报错（换过嵌入模型的库两种维度并存） +//! - 顶掉的分块不是命中 +//! - 带时刻的检索看到的是当时活着的（0019 对着索引再说一遍） + +use pgvector::Vector; +use sqlx::PgPool; +use utopia_store::vector_index::{self, Target}; +use uuid::Uuid; + +/// 这个测试独占的维度：其它连库测试用 3 维，索引按维度分开，互不打扰 +const DIMS: usize = 7; +const CROWD: usize = 2000; + +fn t(s: &str) -> chrono::DateTime { + s.parse().unwrap() +} + +/// 可复现的伪随机向量：同一粒种子同一批向量，失败了能重放 +fn lcg(seed: &mut u64) -> f32 { + *seed = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((*seed >> 33) as f32 / (1u64 << 31) as f32) - 0.5 +} + +struct Fixture { + org: Uuid, + big: Uuid, + small: Uuid, + /// 大库里与查询向量重合的那一条 + planted: Uuid, + /// 与查询重合但五月被顶掉的 + superseded: Uuid, + /// 小库里唯一的一条 + only: Uuid, + /// 大库里一条 5 维的（换过模型） + other_dims: Uuid, +} + +async fn seed(pool: &PgPool) -> anyhow::Result { + let (org, ws, big, small) = ( + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + ); + let (doc_big, doc_small) = (Uuid::now_v7(), Uuid::now_v7()); + let (planted, superseded, only, other_dims) = ( + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + Uuid::now_v7(), + ); + let tag = Uuid::now_v7(); + + sqlx::query("INSERT INTO organizations (id, name) VALUES ($1, 'plan-test')") + .bind(org) + .execute(pool) + .await?; + sqlx::query("INSERT INTO workspaces (id, org_id, name) VALUES ($1, $2, 'plan-test')") + .bind(ws) + .bind(org) + .execute(pool) + .await?; + for (kb, name, doc) in [(big, "big", doc_big), (small, "small", doc_small)] { + sqlx::query("INSERT INTO knowledge_bases (id, workspace_id, name) VALUES ($1, $2, $3)") + .bind(kb) + .bind(ws) + .bind(name) + .execute(pool) + .await?; + sqlx::query( + "INSERT INTO documents (id, kb_id, filename, sha256, status, created_at) + VALUES ($1, $2, $3, $4, 'ready', $5)", + ) + .bind(doc) + .bind(kb) + .bind(format!("{name}.md")) + .bind(format!("sha-{tag}-{name}")) + .bind(t("2026-02-01T00:00:00Z")) + .execute(pool) + .await?; + } + + let mut tx = pool.begin().await?; + let mut seed = 0x5eed_u64; + for i in 0..CROWD { + let v: Vec = (0..DIMS).map(|_| lcg(&mut seed)).collect(); + sqlx::query( + "INSERT INTO chunks (id, kb_id, document_id, seq, text, embedding, created_at) + VALUES ($1, $2, $3, $4, $5, $6, $7)", + ) + .bind(Uuid::now_v7()) + .bind(big) + .bind(doc_big) + .bind(i as i32) + .bind(format!("crowd {i}")) + .bind(Vector::from(v)) + .bind(t("2026-02-01T00:00:00Z")) + .execute(&mut *tx) + .await?; + } + // 与查询重合的一条：活着的、五月被顶掉的各一 + for (id, seq, text, superseded_at) in [ + (planted, CROWD as i32, "the planted one", None), + ( + superseded, + CROWD as i32 + 1, + "the one a reparse displaced", + Some("2026-05-01T00:00:00Z"), + ), + ] { + sqlx::query( + "INSERT INTO chunks (id, kb_id, document_id, seq, text, embedding, created_at, + superseded_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8)", + ) + .bind(id) + .bind(big) + .bind(doc_big) + .bind(seq) + .bind(text) + .bind(Vector::from(query())) + .bind(t("2026-02-01T00:00:00Z")) + .bind(superseded_at.map(t)) + .execute(&mut *tx) + .await?; + } + // 换过模型留下的 5 维一条 + sqlx::query( + "INSERT INTO chunks (id, kb_id, document_id, seq, text, embedding, created_at) + VALUES ($1, $2, $3, $4, $5, $6, $7)", + ) + .bind(other_dims) + .bind(big) + .bind(doc_big) + .bind(CROWD as i32 + 2) + .bind("five dims") + .bind(Vector::from(vec![1.0, 0.0, 0.0, 0.0, 0.0])) + .bind(t("2026-02-01T00:00:00Z")) + .execute(&mut *tx) + .await?; + // 小库：一条，离查询不近不远 + sqlx::query( + "INSERT INTO chunks (id, kb_id, document_id, seq, text, embedding, created_at) + VALUES ($1, $2, $3, 0, 'the only one', $4, $5)", + ) + .bind(only) + .bind(small) + .bind(doc_small) + .bind(Vector::from(vec![0.3, -0.2, 0.1, 0.4, -0.1, 0.2, 0.0])) + .bind(t("2026-02-01T00:00:00Z")) + .execute(&mut *tx) + .await?; + tx.commit().await?; + + Ok(Fixture { + org, + big, + small, + planted, + superseded, + only, + other_dims, + }) +} + +fn query() -> Vec { + vec![0.11, -0.42, 0.37, 0.05, -0.28, 0.19, 0.33] +} + +/// 六个问题的答案，一次问完；两种计划下各取一份来比 +#[derive(Debug, PartialEq)] +struct Answers { + nearest_in_big: Option, + small_limit_10: Vec, + other_dims_query: Vec, + big_top_10_has_superseded: bool, + big_top_10_has_other_dims: bool, + as_of_march_first: Option, +} + +async fn ask(pool: &PgPool, f: &Fixture) -> anyhow::Result { + use utopia_store::documents::vector_search; + let big_now = vector_search(pool, f.big, &query(), 10, None).await?; + let small_now = vector_search(pool, f.small, &query(), 10, None).await?; + let five = vector_search(pool, f.big, &[1.0, 0.0, 0.0, 0.0, 0.0], 10, None).await?; + let big_then = + vector_search(pool, f.big, &query(), 10, Some(t("2026-03-01T00:00:00Z"))).await?; + Ok(Answers { + nearest_in_big: big_now.first().copied(), + small_limit_10: small_now, + other_dims_query: five, + big_top_10_has_superseded: big_now.contains(&f.superseded), + big_top_10_has_other_dims: big_now.contains(&f.other_dims), + as_of_march_first: big_then.first().copied(), + }) +} + +#[tokio::test] +async fn the_answers_agree_with_and_without_the_index() -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = PgPool::connect(&url).await?; + let f = seed(&pool).await?; + + let run = async { + // ---- 一、没有索引：精确路径给出的答案 + vector_index::drop(&pool, Target::Chunks, DIMS).await?; + assert_eq!( + vector_index::status(&pool, Target::Chunks, DIMS).await?, + None + ); + let exact = ask(&pool, &f).await?; + assert_eq!(exact.nearest_in_big, Some(f.planted), "重合的那一条最近"); + assert_eq!( + exact.small_limit_10, + vec![f.only], + "小库只有一条，LIMIT 10 就回这一条——不多不少，也不串大库" + ); + assert_eq!( + exact.other_dims_query, + vec![f.other_dims], + "5 维的查询只找 5 维的" + ); + assert!(!exact.big_top_10_has_superseded, "顶掉的不是命中"); + assert!(!exact.big_top_10_has_other_dims, "7 维的查询看不见 5 维的"); + // 三月一日：顶掉的那条还活着，它与查询重合，所以它或活着的那条居首都对—— + // 两条向量一样，并列由 id 定。只钉「顶掉的那条回来了」这件事 + let then = utopia_store::documents::vector_search( + &pool, + f.big, + &query(), + 2, + Some(t("2026-03-01T00:00:00Z")), + ) + .await?; + assert!( + then.contains(&f.superseded), + "三月一日它还活着，带时刻的检索该看见它" + ); + + // ---- 二、建索引：任务本体，事务外，两次幂等 + let built = vector_index::build(&pool, Target::Chunks, DIMS).await?; + assert!(built.created, "第一次建"); + assert_eq!(built.name, "chunks_embedding_hnsw_7"); + assert_eq!( + vector_index::status(&pool, Target::Chunks, DIMS).await?, + Some(true) + ); + let again = vector_index::build(&pool, Target::Chunks, DIMS).await?; + assert!(!again.created, "第二次是空操作"); + let def: (String,) = sqlx::query_as( + "SELECT pg_get_indexdef(indexrelid) FROM pg_index + WHERE indexrelid = 'chunks_embedding_hnsw_7'::regclass", + ) + .fetch_one(&pool) + .await?; + assert!(def.0.contains("USING hnsw"), "{}", def.0); + assert!(def.0.contains("vector_cosine_ops"), "{}", def.0); + assert!( + def.0.contains("WHERE (vector_dims(embedding) = 7)"), + "{}", + def.0 + ); + + // ---- 三、有索引：同样六问，答案必须一样 + let indexed = ask(&pool, &f).await?; + assert_eq!(indexed, exact, "索引改了答案"); + Ok::<_, anyhow::Error>(()) + } + .await; + + let _ = vector_index::drop(&pool, Target::Chunks, DIMS).await; + sqlx::query("DELETE FROM knowledge_bases WHERE id = ANY($1)") + .bind(vec![f.big, f.small]) + .execute(&pool) + .await?; + sqlx::query("DELETE FROM organizations WHERE id = $1") + .bind(f.org) + .execute(&pool) + .await?; + run +} diff --git a/docs/decisions/0035-a-vector-index-is-built-by-a-job.md b/docs/decisions/0035-a-vector-index-is-built-by-a-job.md new file mode 100644 index 00000000..5a5f5528 --- /dev/null +++ b/docs/decisions/0035-a-vector-index-is-built-by-a-job.md @@ -0,0 +1,74 @@ +# 0035 · A vector index is built by a job + +- **Status**: implemented · a partial HNSW index per dimension on `chunks.embedding` and `entities.profile_embedding`, requested by the first write of that dimension and built by a `build_vector_index` job outside any transaction · `vector_search` and `nearest_typed_entities` write the dimension as a literal, cast both sides and set `hnsw.iterative_scan = relaxed_order` · type resolution gathers a batch's neighbours eight at a time, in order, and remembers descendant sets per batch (#512, #514) · dimensions above 2000 stay on the exact path +- **Written**: 2026-09-09 (conventions in the [README](README.md)) +- **Related**: the ingest migration said P1 scans sequentially and an HNSW index comes "at volume"; this is that note coming due. [0019](0019-the-second-clock-can-be-rewound.md) is why the record-axis filter stays on the query and the index accommodates it. [0016](0016-close-the-open-seams-before-cutting-new-ones.md) C2 is the loop that scanned the entity table once per subject. + +> `chunks.embedding` had no index of any kind, so every hybrid query computed cosine distance against every embedded chunk in the base and sorted the lot to take ten. `entities.profile_embedding` had none either, and type resolution asked it once per subject, sixty subjects a round, ten rounds a job. The column is `vector` with no `(N)`: the dimension follows the workspace's embedding model, which no migration knows. + +## Measured before deciding + +Sixty thousand chunks of 1024 dimensions in one table across three bases (50,000 / 10,000 / 20), random vectors, pgvector 0.8.6. + +| | | +|---|---| +| Today's query, base of 50k | seq scan, 195 ms | +| Today's query, base of 10k | `chunks_kb_idx` + sort, 39 ms | +| Expression HNSW index, serial build over 60k rows | 87 s, 469 MB | +| Rewritten query, base of 50k | 4.5 ms | +| HNSW forced on the 10k base, `iterative_scan = off` | 3 rows of the 24 asked for | +| same, `relaxed_order` | 24 rows, 7.9 ms | + +So the cost was about 4 ms per thousand chunks in the base being searched, per query. Nobody feels it at hundreds of chunks; a base of 100k pays 0.4 s per question on this leg, and type resolution paid it sixty times a round. + +Then again with the code in this record, on a copy of a bench base (531 real chunks and 3,269 real entity profiles at 1024 dims) with a synthetic base of 50,000 chunks written beside them: + +| | | +|---|---| +| The rewritten `vector_search` on the 50k base, no index | seq scan and top-N sort, 378.6 ms | +| same, index built | index scan, 6.0 ms | +| Serial build over 50,531 rows at the default `maintenance_work_mem` of 64 MB | 4 min 20 s; "graph no longer fits after 13,930 tuples" | +| same at 512 MB | 60.6 s, 394 MB | +| 200 real entity queries on a base of 1,415 profiles, HNSW forced, recall@10 against the exact top ten | 0.997 by distance; 0.917 by id, the gap being ties at the tenth place | +| 522 real chunk queries, each base about 1% of the index, HNSW forced, `scan_mem_multiplier` 1 | recall@10 0.705 with `relaxed_order`, 154 lists short, 32 ms a query; 0.075 without `relaxed_order` | +| same, `max_scan_tuples` 100,000 | unchanged: 0.705, 154 short | +| same, `scan_mem_multiplier` 4 | recall@10 1.0, no list short, 74 ms a query; 16 is the same | +| A real base of 229 chunks, HNSW forced, the query far from the base | 0 rows after 19,345 tuples, 105 ms | +| The same base with `chunks_kb_idx` back and the planner choosing | `chunks_kb_idx` and a sort, 1.5 ms; the 50k base beside it still takes the HNSW path, 3.8 ms | + +The forced lines are the planner's choice taken away (the `kb_id` index dropped on the copy): they are the floor, and the reason the scan settings below are not the defaults. + +## Decisions + +**A job, not a migration.** The dimension is only known when vectors are written, and `CREATE INDEX CONCURRENTLY` cannot run inside the transaction sqlx wraps migrations in, while a plain `CREATE INDEX` holds `ACCESS EXCLUSIVE` on `chunks` for the length of the build. So the first write of a dimension asks for an index (`vector_index::request`), which queues one `build_vector_index` job for that table and dimension unless one is already queued, and the job builds it `CONCURRENTLY IF NOT EXISTS` on a plain connection. An index that a failed build left invalid is dropped and rebuilt, because `IF NOT EXISTS` would otherwise take it for finished. Once the index is seen the process remembers it, so a write costs a set lookup; before that it costs one catalog query and one insert-unless-queued, for the minute or two the build takes. The build's session sets `maintenance_work_mem` to 512 MB and `max_parallel_maintenance_workers` to 0: at the default 64 MB the graph stops fitting after some fourteen thousand rows of 1024 dimensions and every row after that goes through disk (four minutes and twenty seconds for 50k rows against one minute), and a parallel build wants shared memory that Docker's default 64 MB `/dev/shm` does not have. + +**A partial expression index per dimension.** `USING hnsw ((embedding::vector(N)) vector_cosine_ops) WHERE vector_dims(embedding) = N`. The cast gives pgvector the dimension the column type lacks; the predicate keeps rows of another dimension out, so a workspace that changed its embedding model has two indexes and two populations rather than one build that fails. `vector_cosine_ops` matches the `<=>` the queries use. + +**The dimension is a literal in the SQL.** With `vector_dims(embedding) = $2` bound as a parameter the custom plan uses the partial index and the generic plan falls back to a sequential scan. sqlx's prepared statements switch to a generic plan after five executions, so the index would have stopped being used on the sixth query, silently. `same_dims` and `distance` format the integer in, and the `ORDER BY` expression is character for character the indexed one. + +**`hnsw.iterative_scan = relaxed_order` on every nearest-neighbour read.** HNSW takes `ef_search` candidates and applies the `WHERE` afterwards. On a table shared by tenants a small base holds few of those candidates, and `LIMIT 10` comes back with three rows or none: measured above, and the ordinary case rather than a corner. Iterative scan keeps walking until the limit is met, and it stops early on two conditions, both loosened here. The one that bound in measurement is memory: `hnsw.scan_mem_multiplier` caps the scan at that multiple of `work_mem`, and at the default of 1 (4 MB) the forced scan over real bases that are each a hundredth of the index came back short on 154 of 522 queries, recall 0.705, at 32 ms a query; at 4 every list filled, recall 1.0, at 74 ms a query, and 16 changed nothing more. The other is `hnsw.max_scan_tuples`, raised from 20,000 to 100,000; it did not bind on the copy (the same 0.705 at either value), and it is raised because of what each failure looks like: a short list is silent, a slow query is visible, and a hundred thousand tuples bound the slow case at about half a second. The planner keeps a base that small on the exact path when it has the choice (229 real chunks beside the 50k: `chunks_kb_idx` and a sort, 1.5 ms), so the forced figures are the floor rather than the expectation. All three settings are `SET LOCAL`, so each read runs in a transaction. A pgvector older than 0.8 has no such setting; the process probes once and leaves it unset, and the query is still correct. + +**The planner chooses the path.** With `chunks_kb_idx` present it picked index-plus-sort for the 20-row and 10k-row bases on its own and HNSW only for the 50k one. No application-side threshold: a threshold is a guess about the planner, and a wrong guess is slow on both sides. + +**Entities take the same mechanism.** `nearest_typed_entities` reads the subject's vector first so the dimension can be written into the SQL, then runs the same shaped query. It gained a dimension guard it never had: a base with two dimensions of profiles used to error on `<=>`. + +**The loop gathers before it reasons** (#514). The sixty neighbour queries of a batch are independent; they were serial because the loop was. `nearest_typed_for_each` runs them eight at a time with `buffered`, which yields in input order, and the per-subject reasoning stays in that order because adjudication downstream reads it. Eight is well under the pool of 32, which `db.rs` sizes for concurrent short queries. `descendants_of` is keyed on the coarse class, drawn from a small vocabulary, so a batch asks the same recursive query over and over; `DescendantsMemo` answers once per class per batch. A subject with no class has no descendants axis (0009) and stays out of the memo rather than sharing a key with a real class. + +**Above 2000 dimensions there is no index.** That is the HNSW limit for `vector`; `text-embedding-3-large` is 3072. Such a dimension is never requested and the query stays exact. `halfvec` reaches 4000 at another precision and waits for someone to need it. + +## Dead ends + +- **ivfflat.** Needs a representative sample at build time to pick its lists, and a corpus that grows past the sample loses recall quietly. For a system whose claim is traceable evidence, silent recall loss is the wrong failure mode. +- **A numbered migration.** Two branches each adding the next number merge cleanly and then neither runs; and see the transaction above. +- **Dropping the record-axis filter so the index applies cleanly.** Replay is the product (0019). The index accommodates the filter. +- **Sizing the pool to the batch.** Sixty concurrent scans would fit a pool of 64 and move the load onto Postgres; `db.rs` already argues against sizing the pool to the worker count. +- **Caching neighbours across subjects.** Every subject has a different query vector; there is nothing to share. + +## What the tests pin + +Every question is asked twice, without the index and with it, and the answers must agree: the nearest chunk, a small base beside a large one filling its `LIMIT`, tenancy, a second dimension in one table, a superseded chunk, a moment on the record axis. For entities: the batch comes back in input order at concurrency 1 and 8, with and without the index; a two-connection pool completes sixty subjects; the memo returns what the query returns and does not notice a class added mid-batch. A test asserts on answers only; a plan change stays quiet. + +## Open questions + +- **Recall on a corpus the planner sends to the index.** The entity figure above (0.997) is on real profiles with the index forced; the chunk figures are on real chunks in an adversarial table, a synthetic tenant whose vectors sit closer to every query than the query's own base does, and they reach 1.0 only with the scan memory raised. A real deployment with two large tenants of different subject matter is the case still unmeasured, and `hnsw.ef_search` (40 here; 200 cost 10 ms in the synthetic run and changed nothing on the copy) is the first knob if it disappoints. +- **A dimension that leaves.** When a workspace changes model, the old dimension's index stays until someone drops it. It is small harm and no mechanism yet. diff --git a/docs/decisions/README.md b/docs/decisions/README.md index faa5964a..bd27eede 100644 --- a/docs/decisions/README.md +++ b/docs/decisions/README.md @@ -56,6 +56,7 @@ The test for writing one: if someone (including us) looks at a piece of code in | 0030 | [A rule may read what a rule concluded](0030-a-rule-may-read-what-a-rule-concluded.md) | Cuts 1–2 · A rule could not read another rule's conclusion (`attribute_facts()` reads `facts`, subject scope is `entities.type_id`, one pass), so `A → B → C` has to be written as one flattened rule. The answer is the fixed point the axiom side already runs, **inside one call and in memory** — which answers 0013's objection instead of waiving it, since a run stays a pure function of the asserted facts. A derived typing enters as a **premise**, not a filter, so the conclusion narrows to the window the entity was actually a `B`; a rule may read its own conclusions because the derivable space is finite; `fact_derivations` gains `premise_derived_id` and the proof becomes a tree; retirement cascades for free because recompute is total. The page showing which rules feed which is next | | 0031 | [An event holds at the moment it names](0031-an-event-holds-at-the-moment-it-names.md) | Implemented (#486) · `Validity::under` normalises every write by the predicate's `temporal` (an event is one moment written at both ends, an eternal fact has no dates), `world_axis` and `read_span` read an event as the bucket it names and an undated event at no moment, an eternal fact at every moment · the prompt marks `[event]` / `[eternal]` and says what to write · no schema change, old rows read correctly · the panel's point rendering and the ontology hint are the UI cut | | 0032 | [A rule computes what it concludes](0032-a-rule-computes-what-it-concludes.md) | Cuts 1–2 · A rule's conclusion was a **constant** (`attribute_rules.conclude_value` is a JSONB literal), so a criterion about a number nobody wrote down — net pay, a margin, a ratio — gets computed elsewhere and typed back in as an assertion, losing the readings it stood on. An operand on either side becomes an expression tree over attributes: its premises are the readings it touched and its interval is their intersection, so retirement needs no new machinery. **Picked, not typed** — not because a parsed string would break the display (it would not; that reason is wrong and recorded as wrong) but because a picker cannot compose an expression over a predicate that does not exist, and because a text box grows a grammar on request. A value may be reached **across a relation** (a path has definite premises and an interval; a to-many path is just a cartesian product the evaluator already walks), while **aggregation is refused**: a `sum` asserts *these are all the readings*, a completeness claim this base cannot hold — and while recompute still retires it correctly, its proof cannot say why it was withdrawn, and nothing incremental could ever wake it. Row count, unit checking and division by zero are the costs | +| 0035 | [A vector index is built by a job](0035-a-vector-index-is-built-by-a-job.md) | Implemented · a partial HNSW index per dimension on `chunks.embedding` and `entities.profile_embedding`, requested by the first write of that dimension and built by a `build_vector_index` job outside any transaction · the two nearest-neighbour reads write the dimension as a literal, cast both sides and set `hnsw.iterative_scan = relaxed_order` · type resolution gathers its neighbours eight at a time in order and remembers descendant sets per batch (#512, #514) · dimensions above 2000 stay on the exact path | | 0033 | [RSS summaries are scoped to the source being listed](0033-rss-source-summaries-are-source-scoped.md) | Query implemented for #417 · source- and generation-scoped lateral aggregation preserves the flat API; nested `rss_full_content` public contract pending; baseline rows are intentionally outside the five work counts | ## Not a decision record