DuckDB 中的向量相似度搜索

Max Gabrielsson

2024-05-03 · 8 分钟

在 DuckDB v0.10.0 中，我们引入了 ARRAY 数据类型，它用于存储固定大小的列表，以补充现有的可变大小 LIST 数据类型。

添加此数据类型的最初动机是为列表提供优化操作，这些操作可以利用其子元素的定位语义并避免分支，因为所有列表都具有相同的长度。例如，想象一下你在 NumPy 中会进行的数组操作：堆叠、移位、乘法——应有尽有。此外，我们希望提高与 Apache Arrow 的互操作性，因为以前 Arrow 的固定大小列表类型在被摄取到 DuckDB 时会转换为常规的可变大小列表，从而丢失一些类型信息。

然而，随着对向量嵌入和语义相似度搜索的热情高涨，我们也为这种新的 ARRAY 类型偷偷加入了几个距离度量函数：array_distance、array_negative_inner_product 和 array_cosine_distance

如果您是今天幸运的 10,000 人之一，并且还没有听说过词嵌入或向量搜索，简单来说，它是一种将文档、图像、实体——即数据——表示为高维向量，然后在一个向量空间中搜索相似向量的技术，使用某种数学“距离”表达式来衡量相似度。这广泛应用于各种领域，从自然语言处理到推荐系统和图像识别，并且由于生成式 AI 的出现和预训练模型的可用性，最近人气激增。

这让社区非常兴奋！虽然我们（DuckDB Labs）最初公开表示不会向 DuckDB 添加向量相似度搜索索引，因为我们认为这超出了范围，但我们对通过扩展支持自定义索引普遍很感兴趣。哎呀，自从 DuckDB 空间扩展问世以来，我个人就一直在唠叨着想要插入一个“R 树”索引！所以当我们的一个客户项目演变为创建自定义“HNSW”索引扩展的概念验证时，我们决定试一试。然后……事情就这样一步步发展起来了。

快进到今天，我们很高兴地宣布 DuckDB 的 vss 向量相似度搜索扩展已可用！虽然有些人可能会说我们参加向量搜索派对迟了，但我们更愿意认为这场派对才刚刚开始！

那么，vss 里有什么呢？

向量相似度搜索 (VSS) 扩展

从表面上看，vss 似乎是一个相对较小的 DuckDB 扩展。它不提供任何新的数据类型、标量函数或复制函数，而只提供一种新的索引类型：HNSW（分层可导航小世界），这是一种基于图的索引结构，特别适合高维向量相似度搜索。

-- Create a table with an array column
CREATE TABLE embeddings (vec FLOAT[3]);

-- Create an HNSW index on the column
CREATE INDEX idx ON embeddings USING HNSW (vec);

这种索引类型不能像内置的 ART 索引那样用于强制约束或唯一性，也不能用于加速连接或索引常规列。相反，HNSW 索引仅适用于包含 FLOAT 元素的 ARRAY 类型列，并且仅用于加速计算常量 FLOAT ARRAY 与索引列中 FLOAT ARRAY 之间“距离”的查询，并按结果距离排序并返回前 n 个结果。也就是说，以下形式的查询：

SELECT *
FROM embeddings
ORDER BY array_distance(vec, [1, 2, 3]::FLOAT[3])
LIMIT 3;

的逻辑计划将被优化为对新的 HNSW 索引扫描操作符的投影，完全移除限制和排序。我们可以通过检查 EXPLAIN 输出进行验证

EXPLAIN
SELECT *
FROM embeddings
ORDER BY array_distance(vec, [1, 2, 3]::FLOAT[3])
LIMIT 3;

┌───────────────────────────┐
│         PROJECTION        │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│             #0            │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│         PROJECTION        │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│            vec            │
│array_distance(vec, [1.0, 2│
│         .0, 3.0])         │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│      HNSW_INDEX_SCAN      │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│   t1 (HNSW INDEX SCAN :   │
│            idx)           │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│            vec            │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│           EC: 3           │
└───────────────────────────┘

您可以向 HNSW 索引创建语句传递一个 metric 参数，以决定使用哪种距离度量。支持的度量包括 l2sq、cosine 和 inner_product，它们与三个内置距离函数：array_distance、array_cosine_distance 和 array_negative_inner_product 相对应。默认值为 l2sq，它使用欧几里得距离（array_distance）。

CREATE INDEX l2sq_idx ON embeddings USING HNSW (vec)
WITH (metric = 'l2sq');

使用余弦距离（array_cosine_distance）

CREATE INDEX cos_idx ON embeddings USING HNSW (vec)
WITH (metric = 'cosine');

使用内积（array_negative_inner_product）

CREATE INDEX ip_idx ON embeddings USING HNSW (vec)
WITH (metric = 'ip');

实现

vss 扩展基于 usearch 库，该库提供了 HNSW 索引数据结构的灵活 C++ 实现，拥有令人印象深刻的性能基准。虽然我们目前只使用了 usearch 提供的所有功能和调优选项的一个子集，但我们很高兴能探索未来如何利用其更多功能。到目前为止，我们很满意它与 DuckDB 的开发理念非常契合。就像 DuckDB 本身一样，usearch 使用可移植的 C++11 编写，没有外部依赖，并在宽松的许可下发布，使其能够非常顺利地集成到我们的扩展构建和分发流程中。

限制

目前最大的限制是 HNSW 索引只能在内存数据库中创建，除非将配置参数 SET hnsw_enable_experimental_persistence = bool 设置为 true。如果未设置此参数，在磁盘支持的数据库中创建 HNSW 索引的任何尝试都将导致错误消息；但如果设置了该参数，索引不仅会在内存中创建，还会在检查点期间作为 DuckDB 数据库文件的一部分持久化到磁盘。重新启动或加载包含持久化 HNSW 索引的数据库文件后，每当首次访问关联表时，索引将惰性加载回内存，这比从头开始重新创建索引要快得多。

将此功能锁定在实验性标志之后的原因是，我们仍然存在一些与自定义索引持久化相关的已知问题，我们希望在默认启用之前解决这些问题。特别是，WAL 恢复尚未为自定义索引正确实现，这意味着如果在 HNSW 索引表中存在未提交更改时发生崩溃或数据库意外关闭，可能会导致数据丢失或索引损坏。虽然通过先单独启动 DuckDB，加载 vss 扩展，然后 ATTACH 数据库文件（这确保了 HNSW 索引功能在 WAL 回放期间可用）可以在技术上从意外关闭中手动恢复，但您不应在生产工作负载中依赖此方法。

我们正在积极解决与索引持久化相关的这个问题和其他问题，希望它们能进入 DuckDB v0.10.3，但目前我们建议仅在内存数据库中使用 HNSW 索引。

然而，在运行时，就像 ART 索引一样，HNSW 索引必须能够完全载入 RAM。并且 HNSW 在运行时分配的内存是在 DuckDB 内存管理系统“外部”分配的，这意味着它不会遵守 DuckDB 的 memory_limit 配置参数。

到目前为止，HNSW 索引的另一个当前限制是它只支持数组元素的 FLOAT（32 位单精度浮点数）类型，并且只支持与三个内置距离函数 array_distance、array_negative_inner_product 和 array_cosine_distance 相对应的距离度量。但这同样是我们计划在不久的将来扩展的功能，因为它更多的是“我们还没来得及做”的限制，而不是技术限制。

结论

DuckDB 的 vss 扩展是一个新扩展，它增加了在 DuckDB 中为固定大小列表列创建 HNSW 索引的支持，从而加速向量相似度搜索查询。通过运行 INSTALL vss; LOAD vss，该扩展目前可以在 DuckDB v0.10.2 的所有支持平台（包括 Wasm！）上安装。vss 扩展通过提供自定义索引类型，为 DuckDB 扩展开辟了新领域，我们很高兴能在此功能的基础上进行完善和扩展。

虽然我们仍在努力解决上述一些限制，特别是与持久化（和性能）相关的限制，但我们仍然非常希望分享 vss 扩展的这个早期版本，因为我们相信这将为社区带来许多酷炫的机会。因此，请务必查看 vss 扩展文档以获取有关如何使用此扩展的更多信息！

这项工作得到了 DuckDB Labs 客户的赞助！如果您对特定功能的类似工作感兴趣，请联系 DuckDB Labs。此外，我们也非常欢迎贡献者！请通过 Discord 或在 vss 扩展 GitHub 仓库上联系 DuckDB Labs 团队，以了解最新进展。