An embedding wider than 2000 dims gets a halfvec index that picks the candidates, and they are still scored and ordered at full precision

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)

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<Near> = 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)
 }
 

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),
     ))

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<Postgres>) {
     }
 }
 
-/// 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<Built> {
     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<Buil
         conn.execute(
             format!(
                 "CREATE INDEX CONCURRENTLY IF NOT EXISTS {name} ON {table} \
-                 USING hnsw (({col}::vector({dims})) vector_cosine_ops) WHERE {pred}",
+                 USING hnsw (({col}::{ty}({dims})) {ty}_cosine_ops) WHERE {pred}",
                 table = target.table(),
                 col = target.column(),
                 pred = same_dims(target.column(), dims),
+                ty = index_type(dims),
             )
             .as_str(),
         )
@@ -381,6 +404,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn a_dimension_wider_than_vector_orders_by_halfvec_and_scores_by_vector() {
+        assert_eq!(
+            index_order("c.embedding", 2, 2000),
+            "c.embedding::vector(2000) <=> $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] {

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:?}"

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<f32> {
+    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<Fixture> {
+    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<f32>)> = (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<i64> {
+    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
+}

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;