优化向量查询

• 一、背景介绍
• 二、原始 SQL 查询
• 三、优化思路
• 四、优化后的 SQL 查询
  • 新 SQL 的优化点：
• 五、代码实现说明
• 六、总结

本文档介绍了如何通过预先存储向量到数据表中来优化向量查询操作。我们将对比旧的查询方式与优化后的新查询方式，并详细解释新 SQL 的优化点，以及如何在代码中利用这些优化。

---

一、背景介绍

在原始场景中，需要对大规模文本数据进行向量相似度查询。传统方法是将输入向量直接嵌入到 SQL 查询中进行逐行比较，这会导致以下问题：

• 计算开销大：每次查询时都需要将输入文本转换为向量，并将该向量与数据库中每条记录的向量进行相似度计算。
• 性能瓶颈：当数据量大或向量维度高时，查询时间可能非常长。

为了解决上述问题，我们采取了预先计算并存储文本向量的优化方案，将文本向量存储在单独的表中，并在查询时通过向量 ID 直接引用已存储的向量。

---

二、原始 SQL 查询

原始的 SQL 查询通过 LATERAL 子句动态计算输入向量，并与数据库中存储的向量进行比较：

--# rumi_rmp_professor.name_vector_search
SELECT
  id, first_name, last_name, name, department, department_id,
  avg_difficulty_rounded, avg_rating_rounded, num_ratings, school_id,
  would_take_again_count, would_take_again_percent_rounded, source_url,
  (1 - (name_vector <=> input_vector)) AS name_similarity
FROM
  rumi_rmp_professor,
  LATERAL (
    VALUES (
      ?::VECTOR(3072)
    )
  ) AS input(input_vector)
WHERE
  school_id = ?
  AND name_vector IS NOT NULL
  AND (1 - (name_vector <=> input_vector)) > 0.8
ORDER BY
  (1 - (name_vector <=> input_vector)) DESC
LIMIT 1;

问题点：

• 每次查询都需要构造一个新的输入向量 input_vector，计算开销大。
• 向量相似度的计算在 SELECT、WHERE 和 ORDER BY 子句中多次重复执行。

---

三、优化思路

优化的主要思路包括：

1. 预先存储向量：

   • 将文本向量预先计算好，并存储在独立的表（如 rumi_embedding）中，以便后续直接引用，无需每次重新计算。

2. 通过向量 ID 查询：

   • 在查询时，通过向量 ID 获取预先存储的向量，而不是重新计算输入向量。

3. 减少重复计算：

   • 使用子查询将相似度计算提取出来，避免在 WHERE 和 ORDER BY 中重复计算相似度。

---

四、优化后的 SQL 查询

优化后的 SQL 利用子查询预先计算相似度，并通过向量 ID 获取存储的向量：

--# rumi_rmp_professor.name_vector_search_by_vector_id
SELECT
  sub.id,
  sub.first_name,
  sub.last_name,
  sub.name,
  sub.department,
  sub.department_id,
  sub.avg_difficulty_rounded,
  sub.avg_rating_rounded,
  sub.num_ratings,
  sub.school_id,
  sub.would_take_again_count,
  sub.would_take_again_percent_rounded,
  sub.source_url,
  sub.name_similarity
FROM (
  SELECT
    p.id,
    p.first_name,
    p.last_name,
    p.name,
    p.department,
    p.department_id,
    p.avg_difficulty_rounded,
    p.avg_rating_rounded,
    p.num_ratings,
    p.school_id,
    p.would_take_again_count,
    p.would_take_again_percent_rounded,
    p.source_url,
    (1 - (p.name_vector <=> e.v)) AS name_similarity
  FROM
    rumi_rmp_professor p
  JOIN
    rumi_embedding e ON e.id = ?         -- 传入 vector_id
  WHERE
    p.school_id = ?
    AND p.name_vector IS NOT NULL
) AS sub
WHERE
  sub.name_similarity > 0.8
ORDER BY
  sub.name_similarity DESC
LIMIT 1;

新 SQL 的优化点：

1. 向量预先存储：

   • 使用 rumi_embedding 表存储所有已计算好的向量，通过向量 ID (e.id = ?) 获取目标向量，而不是在查询时动态计算输入向量。

2. 消除重复计算：

   • 将向量相似度的计算 (1 - (p.name_vector <=> e.v)) 放在子查询中，只计算一次，并赋值给临时列 name_similarity。
   • 外层查询使用已计算的 name_similarity 进行过滤和排序，避免多次重复计算。

3. 减少向量构造开销：

   • 通过预先存储向量，只需查询数据库获取现有向量，而不再每次通过 LATERAL VALUES 动态构造新向量，从而节省了构造向量的时间和资源。

这些优化使得查询性能从原始的 3461 毫秒减少到 2979 毫秒，显著提高了查询效率。

---

五、代码实现说明

以下是 Java 代码片段，用于获取文本对应的向量 ID 并执行优化后的查询：

public Long getVectorId(String text) {
  // 根据文本生成 MD5 值
  String md5 = Md5Utils.getMD5(text);
  // 查询预先存储的向量记录
  String sql = String.format("select id from %s where md5=? and m=?", TableNames.rumi_embedding);
  Long id = Db.queryLong(sql, md5, OpenAiModels.text_embedding_3_large);
  if (id != null) {
    return id;
  }

  // 如果未找到则生成向量
  float[] embeddingArray;
  synchronized (vectorLock) {
    embeddingArray = OpenAiClient.embeddingArray(text, OpenAiModels.text_embedding_3_large);
  }

  // 将浮点数组转换为字符串形式
  String string = Arrays.toString(embeddingArray);
  id = SnowflakeIdUtils.id();

  // 转换为 PGobject 格式的向量
  PGobject pgVector = PgVectorUtils.getPgVector(string);
  Row saveRecord = new Row()
    .set("t", text)
    .set("v", pgVector)
    .set("id", id)
    .set("md5", md5)
    .set("m", OpenAiModels.text_embedding_3_large);

  // 保存新的向量记录到数据库
  synchronized (writeLock) {
    Db.save("rumi_embedding", saveRecord);
  }
  return id;
}

public Row searchName(String name, Long schoolId) {
  // 获取文本的向量 ID
  Long vectorId = Aop.get(VectorService.class).getVectorId(name);
  // 获取优化后的 SQL 查询模板
  String professorSql = SqlTemplates.get("rumi_rmp_professor.name_vector_search_by_vector_id");
  // 执行查询
  Row professor = Db.findFirst(professorSql, vectorId, schoolId);
  return professor;
}

代码说明：

1. 获取向量 ID (getVectorId)：

   • 首先根据输入文本生成 MD5 值，并在 rumi_embedding 表中查询是否已存在对应的向量 ID。
   • 如果存在，则直接返回 ID；否则，调用 OpenAI 的接口生成新的文本向量，并将其存储到 rumi_embedding 表中。

2. 使用向量 ID 进行查询 (searchName)：

   • 调用 getVectorId 获取文本对应的向量 ID。
   • 通过优化后的 SQL 查询模板 rumi_rmp_professor.name_vector_search_by_vector_id，使用向量 ID 和学校 ID 进行查询，返回匹配的教授记录。

---

六、总结

通过预先存储向量、使用向量 ID 查询和消除重复计算，新 SQL 优化了原有的向量查询过程，显著减少了查询时间。同时，在代码实现中，我们通过合理的同步和锁机制保证了向量的生成和存储过程的线程安全，进一步提高了系统的稳定性和性能。