From e90d57808cf904ff0e6de1cd81ea0bee31f32bd9 Mon Sep 17 00:00:00 2001 From: Bogdan Szabo Date: Tue, 27 Oct 2026 11:23:16 +0100 Subject: [PATCH] An embedding wider than 2000 dims gets a halfvec index that picks the candidates, and they are still scored and ordered at full precision Signed-off-by: Bogdan Szabo --- crates/utopia-store/src/documents.rs | 3 +- crates/utopia-store/src/name_vectors.rs | 9 +- crates/utopia-store/src/resolution.rs | 3 +- crates/utopia-store/src/vector_index.rs | 57 ++++- .../tests/a_vector_index_is_built_by_a_job.rs | 2 +- ...e_embedding_is_ranked_at_full_precision.rs | 207 ++++++++++++++++++ crates/utopia-store/tests/store/main.rs | 1 + 7 files changed, 268 insertions(+), 14 deletions(-) create mode 100644 crates/utopia-store/tests/store/a_wide_embedding_is_ranked_at_full_precision.rs diff --git a/crates/utopia-store/src/documents.rs b/crates/utopia-store/src/documents.rs index d257013e0..81d314a20 100644 --- a/crates/utopia-store/src/documents.rs +++ b/crates/utopia-store/src/documents.rs @@ -1918,7 +1918,7 @@ pub async fn vector_search( SELECT id, {distance} AS distance FROM chunks c WHERE c.kb_id = $1 AND c.embedding IS NOT NULL AND {live} AND {same_dims}{only} - ORDER BY {distance} + ORDER BY {order} LIMIT $3 ) SELECT id FROM nearest ORDER BY {resort}", @@ -1930,6 +1930,7 @@ pub async fn vector_search( }, same_dims = crate::vector_index::same_dims("c.embedding", dims), distance = crate::vector_index::distance("c.embedding", 2, dims), + order = crate::vector_index::index_order("c.embedding", 2, dims), resort = crate::vector_index::RESORT, )) .bind(kb_id) diff --git a/crates/utopia-store/src/name_vectors.rs b/crates/utopia-store/src/name_vectors.rs index 14700c80a..92d68c7db 100644 --- a/crates/utopia-store/src/name_vectors.rs +++ b/crates/utopia-store/src/name_vectors.rs @@ -161,6 +161,7 @@ pub async fn nearest(pool: &PgPool, kb_id: Uuid, query: &[f32], k: usize) -> App return Ok(Vec::new()); } let distance = vector_index::distance("v.embedding", 2, dims); + let order = vector_index::index_order("v.embedding", 2, dims); let same_dims = vector_index::same_dims("v.embedding", dims); let sql = format!( // 外层次序照 `vector_index::RESORT` 的规矩（`distance + 0, id`），列名限定到 CTE： @@ -169,7 +170,7 @@ pub async fn nearest(pool: &PgPool, kb_id: Uuid, query: &[f32], k: usize) -> App SELECT v.fact_id AS id, v.entity_id, ({distance}) AS distance FROM name_vectors v WHERE v.kb_id = $1 AND {same_dims} - ORDER BY {distance} + ORDER BY {order} LIMIT $3 ) SELECT n.entity_id, n.id AS fact_id, @@ -184,13 +185,17 @@ pub async fn nearest(pool: &PgPool, kb_id: Uuid, query: &[f32], k: usize) -> App ORDER BY n.distance + 0, n.id LIMIT $4", ); + // 规矩 3：有了索引，小库在共享表上不开 iterative_scan 会回不满 + let mut tx = pool.begin().await?; + vector_index::relaxed_order(pool, &mut tx).await?; let rows: Vec = sqlx::query_as(&sql) .bind(kb_id) .bind(Vector::from(query.to_vec())) .bind((k * 4) as i64) .bind(k as i64) - .fetch_all(pool) + .fetch_all(&mut *tx) .await?; + tx.commit().await?; Ok(rows) } diff --git a/crates/utopia-store/src/resolution.rs b/crates/utopia-store/src/resolution.rs index b7f34250a..355b124e4 100644 --- a/crates/utopia-store/src/resolution.rs +++ b/crates/utopia-store/src/resolution.rs @@ -3104,13 +3104,14 @@ pub async fn nearest_typed_entities( 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 {same_dims} - ORDER BY {distance} + ORDER BY {order} LIMIT $4 ) -- 次序在外层再排一遍，并列由实体 id 定（`vector_index::RESORT`，#652） SELECT canonical_name, type_id, key, distance, same_document FROM nearest ORDER BY {resort}", distance = crate::vector_index::distance("e.profile_embedding", 3, dims), + order = crate::vector_index::index_order("e.profile_embedding", 3, dims), resort = crate::vector_index::RESORT, same_dims = crate::vector_index::same_dims("e.profile_embedding", dims), )) diff --git a/crates/utopia-store/src/vector_index.rs b/crates/utopia-store/src/vector_index.rs index e1be34e71..e91159733 100644 --- a/crates/utopia-store/src/vector_index.rs +++ b/crates/utopia-store/src/vector_index.rs @@ -11,8 +11,9 @@ //! 1. **维度写成字面量。** `vector_dims(col) = $2` 绑参数时自定义计划能用索引， //! generic plan 退回顺扫；sqlx 的预处理语句跑五次就切 generic，第六次起索引 //! 悄悄失效。用 [`same_dims`] / [`distance`] 把整数写进 SQL。 -//! 2. **两侧都 cast 到 `vector(N)`。** 索引建在表达式 `col::vector(N)` 上， -//! `ORDER BY` 必须一字不差地写同一个表达式。 +//! 2. **`ORDER BY` 一字不差地写索引的表达式。** 用 [`index_order`]：索引建在 +//! `col::vector(N)`（超过 [`VECTOR_MAX_DIMS`] 时是 `col::halfvec(N)`）上， +//! 写成别的样子索引就用不上。选出来的距离用 [`distance`]，永远是全精度。 //! 3. **`hnsw.iterative_scan = relaxed_order`。** HNSW 先取 `ef_search` 个候选再过 //! WHERE；一张按 kb 分租的表上，小库的行在候选里占不到几个，`LIMIT 10` 会回 //! 三行甚至零行（实测 1 万行的库在 6 万行的表上：关着回 3/24，开着回满）。 @@ -64,10 +65,24 @@ async fn unlock_build(mut conn: sqlx::pool::PoolConnection) { } } -/// pgvector 的 HNSW 对 `vector` 类型的上限。超过的维度（text-embedding-3-large -/// 是 3072）不建索引，查询照常走精确路径。`halfvec` 能到 4000，但那是另一种 -/// 精度，等有人用到再说 -pub const MAX_DIMS: usize = 2000; +/// pgvector 的 HNSW 对 `vector` 类型的上限 +pub const VECTOR_MAX_DIMS: usize = 2000; + +/// 能建索引的上限：`halfvec` 的。超过 [`VECTOR_MAX_DIMS`] 的维度（text-embedding-3-large +/// 是 3072）索引存半精度，它只负责挑出候选；候选的距离与最终次序仍按 `vector` 全精度 +/// 算（[`distance`] 加外层 [`RESORT`]），所以半精度只可能影响「谁进候选」，不改分数。 +/// 没有索引时这种维度每一问都扫全库（3072 维、二十多万个名字向量，一问几十秒）。 +/// 再宽的不建索引，查询照常走精确路径 +pub const MAX_DIMS: usize = 4000; + +/// 索引里存的类型：能用 `vector` 就用，放不下才降到 `halfvec` +fn index_type(dims: usize) -> &'static str { + if dims > VECTOR_MAX_DIMS { + "halfvec" + } else { + "vector" + } +} /// 哪一列。两张表同一套机制，名字不同而已 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -122,11 +137,18 @@ pub fn same_dims(column: &str, dims: usize) -> String { format!("vector_dims({column}) = {dims}") } -/// `<=>` 两侧都 cast 到字面维度（规矩 1、2）；`param` 是查询向量的参数号 +/// 全精度的距离：选出来、给外层排的那一列。`param` 是查询向量的参数号 pub fn distance(column: &str, param: usize, dims: usize) -> String { format!("{column}::vector({dims}) <=> ${param}::vector({dims})") } +/// `ORDER BY` 写的表达式，与索引一字不差（规矩 1、2）。不超过 [`VECTOR_MAX_DIMS`] 时 +/// 就是 [`distance`]；超过时是半精度，只用来让索引挑候选 +pub fn index_order(column: &str, param: usize, dims: usize) -> String { + let ty = index_type(dims); + format!("{column}::{ty}({dims}) <=> ${param}::{ty}({dims})") +} + /// 近邻查询的外层次序（#652）。里层照规矩 2 写 `ORDER BY {distance} LIMIT n` 让索引 /// 接住，套进 `MATERIALIZED` 的 CTE（列名 `distance`、`id`），外层再按这个排： /// - `relaxed_order` 下索引回的次序只是大致按距离，外层补一次真排序 @@ -243,7 +265,7 @@ pub async fn build_as_owner( 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}" + "{dims} dims: HNSW on halfvec holds up to {MAX_DIMS}" ))); } let name = index_name(target, dims); @@ -273,10 +295,11 @@ pub async fn build(pool: &PgPool, target: Target, dims: usize) -> AppResult $2::vector(2000)" + ); + assert_eq!( + index_order("c.embedding", 2, 3072), + "c.embedding::halfvec(3072) <=> $2::halfvec(3072)" + ); + assert_eq!( + distance("c.embedding", 2, 3072), + "c.embedding::vector(3072) <=> $2::vector(3072)" + ); + } + #[test] fn a_target_round_trips_through_the_payload() { for t in [Target::Chunks, Target::EntityProfiles] { 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 index 33f9235dd..0f9b91039 100644 --- 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 @@ -106,7 +106,7 @@ async fn the_first_write_of_a_dimension_queues_one_build() -> anyhow::Result<()> 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"); + .expect_err("HNSW on halfvec holds up to 4000 dims"); assert!( matches!(err, utopia_core::AppError::Validation(_)), "是 Validation，分发层据此判 Terminal：{err:?}" diff --git a/crates/utopia-store/tests/store/a_wide_embedding_is_ranked_at_full_precision.rs b/crates/utopia-store/tests/store/a_wide_embedding_is_ranked_at_full_precision.rs new file mode 100644 index 000000000..0ccfeefca --- /dev/null +++ b/crates/utopia-store/tests/store/a_wide_embedding_is_ranked_at_full_precision.rs @@ -0,0 +1,207 @@ +//! 比 `vector` 的 HNSW 上限还宽的嵌入（text-embedding-3-large 是 3072 维）也有索引： +//! 索引存 `halfvec`，只负责挑候选；距离与次序仍按 `vector` 全精度算。 +//! +//! 两条分块的距离只差在半精度看不见的地方：半精度下并列（并列由 id 定，id 小的那条 +//! 离得远），全精度下近的那条必须排第一。外层若按半精度的距离排，就会把远的排上来。 + +use pgvector::Vector; +use sqlx::postgres::PgPoolOptions; +use sqlx::{Executor, PgPool}; +use utopia_store::vector_index::{self, Target, JOB_KIND, VECTOR_MAX_DIMS}; +use uuid::Uuid; + +/// 这个测试独占的维度：刚好放不进 `vector` 的索引 +const DIMS: usize = VECTOR_MAX_DIMS + 1; +const CROWD: usize = 300; + +fn lcg(seed: &mut u64) -> f32 { + *seed = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((*seed >> 33) as f32 / (1u64 << 31) as f32) - 0.5 +} + +/// 前两维给定、其余为零的向量 +fn head(a: f32, b: f32) -> Vec { + let mut v = vec![0.0; DIMS]; + v[0] = a; + v[1] = b; + v +} + +struct Fixture { + org: Uuid, + big: Uuid, + small: Uuid, + /// 全精度下离查询近的那条 + nearer: Uuid, + /// 半精度下与 `nearer` 一样远，id 更小 + farther: Uuid, + only: 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()); + // `farther` 先取 id：并列时它排前面 + let (farther, nearer, only) = (Uuid::now_v7(), Uuid::now_v7(), Uuid::now_v7()); + let tag = Uuid::now_v7(); + + sqlx::query("INSERT INTO organizations (id, name) VALUES ($1, 'wide-test')") + .bind(org) + .execute(pool) + .await?; + sqlx::query("INSERT INTO workspaces (id, org_id, name) VALUES ($1, $2, 'wide-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) + VALUES ($1, $2, $3, $4, 'ready')", + ) + .bind(doc) + .bind(kb) + .bind(format!("{name}.md")) + .bind(format!("sha-{tag}-{name}")) + .execute(pool) + .await?; + } + + let mut tx = pool.begin().await?; + let mut seed = 0x31de_u64; + let mut rows: Vec<(Uuid, Uuid, Uuid, Vec)> = (0..CROWD) + .map(|_| { + let v = (0..DIMS).map(|_| lcg(&mut seed)).collect(); + (Uuid::now_v7(), big, doc_big, v) + }) + .collect(); + // 1.0004 与 1.0 在半精度里是同一个数（1 附近的间距是 2^-10，1.0004 不到一半，舍成 1），在全精度里不是 + rows.push((farther, big, doc_big, head(1.0, 1.0004))); + rows.push((nearer, big, doc_big, head(1.0, 1.0))); + rows.push((only, small, doc_small, head(0.2, 1.0))); + for (seq, (id, kb, doc, v)) in rows.into_iter().enumerate() { + sqlx::query( + "INSERT INTO chunks (id, kb_id, document_id, seq, text, embedding) + VALUES ($1, $2, $3, $4, 'wide', $5)", + ) + .bind(id) + .bind(kb) + .bind(doc) + .bind(seq as i32) + .bind(Vector::from(v)) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + Ok(Fixture { + org, + big, + small, + nearer, + farther, + only, + }) +} + +async fn index_scans(pool: &PgPool) -> anyhow::Result { + pool.execute("SELECT pg_stat_force_next_flush()").await?; + pool.execute("SELECT pg_stat_clear_snapshot()").await?; + let name = vector_index::index_name(Target::Chunks, DIMS); + Ok(sqlx::query_scalar( + "SELECT coalesce(sum(idx_scan), 0)::int8 FROM pg_stat_user_indexes WHERE indexrelname = $1", + ) + .bind(name) + .fetch_one(pool) + .await?) +} + +#[tokio::test] +async fn a_wide_embedding_is_indexed_and_ranked_at_full_precision() -> anyhow::Result<()> { + let Some(url) = utopia_store::test_db::url() else { + return Ok(()); + }; + let pool = PgPool::connect(&url).await?; + // 一条连接、不许顺扫也不许排序：规划器只剩索引这条路，统计也都记在这条连接上 + let indexed = PgPoolOptions::new() + .max_connections(1) + .after_connect(|conn, _| { + Box::pin(async move { + conn.execute("SET enable_seqscan = off; SET enable_sort = off") + .await?; + Ok(()) + }) + }) + .connect(&url) + .await?; + let f = seed(&pool).await?; + let forget_jobs = || { + sqlx::query("DELETE FROM jobs WHERE kind = $1 AND payload->>'dims' = $2") + .bind(JOB_KIND) + .bind(DIMS.to_string()) + .execute(&pool) + }; + + let run = async { + // 上一次跑留下的都清掉，从「什么都没有」开始 + vector_index::drop(&pool, Target::Chunks, DIMS).await?; + forget_jobs().await?; + let queued = vector_index::request(&pool, Target::Chunks, DIMS).await?; + assert!(queued.is_some(), "{DIMS} 维也排一条建索引任务"); + + let built = vector_index::build(&pool, Target::Chunks, DIMS).await?; + assert!(built.created); + let def: String = sqlx::query_scalar( + "SELECT pg_get_indexdef(indexrelid) FROM pg_index WHERE indexrelid = $1::regclass", + ) + .bind(&built.name) + .fetch_one(&pool) + .await?; + assert!(def.contains("halfvec_cosine_ops"), "{def}"); + + let before = index_scans(&indexed).await?; + let top = + utopia_store::documents::vector_search(&indexed, f.big, &head(1.0, 0.0), 2, None, None) + .await?; + let small = utopia_store::documents::vector_search( + &indexed, + f.small, + &head(1.0, 0.0), + 10, + None, + None, + ) + .await?; + let after = index_scans(&indexed).await?; + + assert!(after > before, "检索走了半精度索引"); + assert_eq!(top, vec![f.nearer, f.farther], "次序按全精度距离"); + assert_eq!(small, vec![f.only], "大库旁边的小库也回满"); + Ok::<_, anyhow::Error>(()) + } + .await; + + let _ = vector_index::drop(&pool, Target::Chunks, DIMS).await; + forget_jobs().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/crates/utopia-store/tests/store/main.rs b/crates/utopia-store/tests/store/main.rs index c5c72edf5..3ef2bf29f 100644 --- a/crates/utopia-store/tests/store/main.rs +++ b/crates/utopia-store/tests/store/main.rs @@ -115,6 +115,7 @@ mod a_token_is_the_person_but_not_all_of_them; mod a_trimmed_description_is_not_stale; mod a_typed_fact_is_as_confirmed_as_its_statement; mod a_viewer_never_sees_a_credential; +mod a_wide_embedding_is_ranked_at_full_precision; mod a_wrong_time_can_be_corrected; mod adopting_an_iri_adopts_the_shape; mod an_admin_grants_system_admin;