DuckDB 新扩展：向量相似度搜索

前言

在当前的数据分析、机器学习和AI领域，向量搜索已成为数据库功能的重要组成部分。Postgres 有了 pgvector，Oracle 23ai 也新增了向量搜索功能，现在 DuckDB 也有插件了。

DuckDB v0.10.0 带来了很多重大更新，其中就有DuckDB也要加入向量数据库的战局的固定长度的Array数据类型^[1]。期间也有几篇向量搜索的尝试的文章，如使用 DuckDB + Ibis 进行 RAG 和在 DUCKDB 中构建人工智能驱动的搜索功能。

之前的探索只是序曲，现在，DuckDB 的 vss 扩展带来了真正的向量搜索能力。vss 引入了对 HNSW（ Hierarchical Navigable Small World，层次化可导航小世界）索引的支持，用以加速向量相似性搜索。

添加此数据类型（Array）的最初动机是为可以利用其子元素的位置语义的列表提供优化操作，并避免分支，因为所有列表的长度都相同。例如，在 NumPy 中进行的数组操作：堆叠、移动、乘法等。此外，还想提高与 Apache Arrow 的互操作性，因为以前 Arrow 的固定大小列表类型在导入 DuckDB 时会被转换为常规的可变大小列表，从而丢失了一些类型信息。

随着对 向量嵌入 和 语义相似性搜索 的兴趣日益增长，DuckDB 还为这种新的 ARRAY 类型引入了几个距离度量函数：array_distance^[2]、 array_inner_product^[3] 和 array_cosine_similarity^[4]。

对于那些尚未涉足词嵌入或向量搜索的读者，这是一次了解这一强大技术的机会，简而言之，这是一种用于将文档、图像、实体 – 数据表示为高维向量，然后在向量空间中搜索相似的向量，使用某种数学“距离”表达式来衡量相似性。这种技术被广泛应用于从自然语言处理到推荐系统和图像识别的各个领域，并且由于生成性 AI 的出现和预训练模型的可用性，最近它的受欢迎程度急剧上升。

这真的很让社区感到兴奋！虽然我们（本文的我们均指DuckDB Labs）最初公开表示我们不会为 DuckDB 添加向量相似性搜索索引，因为我们认为这超出了我们的范围，但我们对通过扩展支持自定义索引非常感兴趣。我个人从 DuckDB 的空间扩展^[5] 开始就一直想要插入一个“R-树”索引！所以当我们的一个客户项目演变成创建一个概念验证的自定义“HNSW”索引扩展时，我们决定尝试一下。于是……事情就这样发展下去了。

`vss` 扩展

现在，我们很高兴地宣布 DuckDB 的 vss 向量相似性搜索扩展的诞生了！虽然有些人可能会说我们在向量搜索派对上迟到了，但我们认为派对才刚刚开始！

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

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

CREATE TABLE embeddings (vec FLOAT[3]);
CREATE INDEX idx ON embeddings USING HNSW (vec);

这种索引类型不能像内置的 ART 索引^[6] 那样用来强制约束或唯一性，也不能用来加快连接或索引常规列。相反，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_similarity 和 array_inner_product。

• 默认是 l2sq，它使用欧几里得距离（array_distance）：

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

• 要使用余弦距离（array_cosine_similarity）：

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

• 要使用内积（array_inner_product）：

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

实现

vss 扩展基于 usearch^[7] 库，它提供了 HNSW 索引数据结构的灵活 C++ 实现，拥有非常令人印象深刻的性能基准。虽然我们目前只使用了 usearch 提供的所有功能和调整选项的子集，但我们很高兴探索如何在未来利用它的更多特性。到目前为止，我们很高兴它与 DuckDB 的开发理念非常吻合。就像 DuckDB 本身一样，usearch 用无外部依赖的可移植 C++11 编写，并在宽松的许可证下发布，这使得它非常容易集成到我们的扩展构建和分发管道中。

局限性

目前 HNSW 索引的一个主要限制是，它只能创建在内存数据库中，除非将配置参数 SET hnsw_enable_experimental_persistence = ⟨bool⟩ 设置为 true。如果这个参数没有设置，任何尝试在磁盘数据库中创建 HNSW 索引的操作都会产生错误信息。但如果设置了该参数，索引不仅会在内存中创建，还会在检查点过程中作为 DuckDB 数据库文件的一部分持久化到磁盘。在重启或加载带有持久化 HNSW 索引的数据库文件后，索引将在首次访问相关表时延迟加载回内存，这比从头开始重新创建索引要快得多。

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

我们正在积极解决这个和其他与索引持久性相关的问题，希望能够在 DuckDB v0.10.3^[8] 中解决，但目前我们建议只在内存数据库中使用 HNSW 索引。

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

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

结论

DuckDB 的 vss 扩展是一个新扩展，它为 DuckDB 中的固定大小列表列添加了创建 HNSW 索引的支持，加速了向量相似性搜索查询。该扩展目前可以在所有支持的平台上（包括 WASM！）的 DuckDB v0.10.2 上通过运行 INSTALL vss; LOAD vss 安装。vss 扩展为 DuckDB 扩展开辟了新天地，提供了自定义索引类型，我们很高兴能够改进并扩展这一功能。

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

感谢 DuckDB Labs 客户的赞助，使这项创新工作得以实现。现在，我们邀请每一位对向量搜索感兴趣的开发者和数据科学家，通过我们的文档和 GitHub 仓库深入了解 vss 扩展，并加入我们的社区，共同推动数据分析技术的进步。

💡：本文在原文 Vector Similarity Search in DuckDB^[10]的基础上翻译整理

引用链接

[1] Array数据类型: https://duckdb.org/docs/sql/data_types/array
[2] array_distance: https://duckdb.org/docs/sql/functions/array#array_distancearray1-array2
[3] array_inner_product: https://duckdb.org/docs/sql/functions/array#array_inner_productarray1-array2
[4] array_cosine_similarity: https://duckdb.org/docs/sql/functions/array#array_cosine_similarityarray1-array2
[5] 空间扩展: https://duckdb.org/docs/extensions/spatial
[6] ART 索引: https://duckdb.org/docs/sql/indexes
[7] usearch: https://github.com/unum-cloud/usearch
[8] DuckDB v0.10.3: https://duckdb.org/docs/dev/release_calendar#upcoming-releases
[9] vss 扩展文档: https://duckdb.org/docs/extensions/vss
[10] Vector Similarity Search in DuckDB: https://duckdb.org/2024/05/03/vector-similarity-search-vss.html