Google Spanner数据库KNN向量查询耗时超30秒的优化咨询

问题背景 #

现有约6000条记录的文档数据库，迁移至Google Spanner后执行RAG查询耗时超30秒，执行计划显示时间主要消耗在ml.predict及“merge distributed union”环节。

现有查询语句 #

SELECT a.id, a.title, cosine_distance(embeddings,
  (
    SELECT embeddings.VALUES
    FROM   ml.predict(model embeddingsmodel,
      (
        SELECT 'What is the meaning of life?' AS content)))) as distance
from `documents` as a
order by distance

嵌入模型定义 #

CREATE MODEL EmbeddingsModel
INPUT(content STRING(MAX))
OUTPUT(
  embeddings
    STRUCT<
      statistics STRUCT<truncated BOOL, token_count FLOAT64>,
      values ARRAY<FLOAT64>>
)
REMOTE OPTIONS (
  endpoint = '//aiplatform.googleapis.com/projects/<myproject-id>/locations/<mylocation>/publishers/google/models/textembedding-gecko@003'
);

文档表结构 #

CREATE TABLE
  content (id STRING(36) DEFAULT (GENERATE_UUID()),
    text STRING(MAX),
    title STRING(MAX),
    embeddings ARRAY<FLOAT64>,
    )
PRIMARY KEY
  (id);

优化方案 #

1. 预计算查询嵌入向量，避免重复调用模型 #

当前查询的嵌套写法可能导致ml.predict被多次触发（分布式执行场景下），改为先单独计算一次查询向量，再关联所有文档计算距离，减少远程模型调用次数：

WITH query_embedding AS (
  SELECT embeddings.values AS vec
  FROM ml.predict(model embeddingsmodel, (SELECT 'What is the meaning of life?' AS content))
)
SELECT a.id, a.title, cosine_distance(a.embeddings, q.vec) AS distance
FROM `documents` a, query_embedding q
ORDER BY distance

2. 为embeddings列创建向量索引 #

Spanner支持向量索引加速相似性查询，替代全表扫描计算余弦距离，大幅降低计算开销：

CREATE VECTOR INDEX idx_content_embeddings
ON content(embeddings)
OPTIONS (
  distance_type = 'COSINE',
  dimension = 768 -- 匹配textembedding-gecko@003的默认输出维度，若有调整需同步修改
);

3. 限制返回结果数量 #

如果业务仅需Top N条最相关结果，添加LIMIT子句避免全表排序：

WITH query_embedding AS (
  SELECT embeddings.values AS vec
  FROM ml.predict(model embeddingsmodel, (SELECT 'What is the meaning of life?' AS content))
)
SELECT a.id, a.title, cosine_distance(a.embeddings, q.vec) AS distance
FROM `documents` a, query_embedding q
ORDER BY distance
LIMIT 10; -- 根据实际需求调整数量

配合向量索引使用时，索引可直接返回Top N结果，无需全表排序。

4. 调整Spanner实例配置 #

检查Spanner实例的节点数量与机器类型，若配置过低会导致分布式merge环节瓶颈。可临时增加节点数或升级机器类型测试性能，根据实际负载调整长期配置。

5. 优化模型调用的网络延迟 #

确保AI Platform模型部署区域与Spanner实例在同一区域，减少跨区域网络开销。若当前endpoint的location与Spanner不同，修改REMOTE OPTIONS中的区域参数。

内容的提问来源于stack exchange，提问作者Paul Vincent Craven