在pgvector中使用HNSW索引实现自定义评分权重的最佳实践

背景介绍

在向量数据库应用中，我们经常需要对搜索结果进行自定义加权处理。pgvector作为PostgreSQL的向量扩展，提供了高效的HNSW索引支持，但在实现自定义评分权重时会遇到一些性能挑战。

问题分析

当我们需要根据文档的重要性(is_important字段)对向量相似度结果进行加权时，直接使用CASE WHEN语句会导致HNSW索引失效。例如：

SELECT id, 
    (CASE WHEN is_important THEN 1.2 ELSE 1 END * (embedding <#> search_vector)) AS weighted_score
FROM Document
ORDER BY weighted_score
LIMIT 50

这种写法虽然逻辑正确，但无法利用HNSW索引的优势，导致全表扫描，在400万条记录的情况下性能极差。

解决方案探索

子查询方案

第一种尝试是使用子查询先获取近似结果，再应用权重：

SELECT id, 
    (CASE WHEN isImportant THEN 1.2 ELSE 1 END * (summaryEmbedding <#> search_vector)) AS weighted_score
FROM Document
WHERE id IN (
    SELECT id
    FROM Document
    ORDER BY (summaryEmbedding <#> search_vector) ASC
    LIMIT 128
)
ORDER BY weighted_score ASC
LIMIT 50

这种方案利用了索引，但存在召回率问题，可能过滤掉实际加权后应该排在前面的重要文档。

合并查询方案

第二种方案是分别查询重要和非重要文档，然后合并结果：

1. 先查询重要文档(加权1.2倍)
2. 再查询非重要文档(加权1倍)
3. 合并结果并排序

虽然看起来需要多次查询，但实际测试表明：

• 性能与子查询方案相当甚至更好
• 召回率显著提高
• 可以利用索引加速每部分查询

性能优化建议

1. 索引参数调优：对于400万量级数据，m=16可能比m=32更合适，能显著提高搜索速度
2. 并行查询：合并方案可以并行执行各部分查询
3. 结果集大小：适当调整子查询中的LIMIT值(如128)平衡召回率和性能

实现示例

-- 合并方案实现
(SELECT id, 1.2 * (embedding <#> search_vector) AS score
 FROM Document WHERE is_important = true
 ORDER BY embedding <#> search_vector
 LIMIT 100)

UNION ALL

(SELECT id, 1.0 * (embedding <#> search_vector) AS score
 FROM Document WHERE is_important = false
 ORDER BY embedding <#> search_vector
 LIMIT 100)

ORDER BY score
LIMIT 50

结论

在pgvector中实现自定义评分权重时，合并查询方案相比子查询方案具有明显优势。它不仅保持了良好的性能，还显著提高了结果质量。开发者应根据实际数据分布和业务需求，选择合适的LIMIT值和权重系数，并通过测试确定最佳的HNSW索引参数。