向量相似度搜索扩展有什么新功能？
TL;DR: DuckDB 离成为向量数据库又近了一步！本文将展示向量搜索扩展中实现的新性能优化。
在上一篇博客文章中，我们介绍了 DuckDB 向量相似度搜索 (VSS) 扩展。尽管该扩展仍处于实验阶段，但我们认为深入了解自首次发布以来我们一直在开发的一些新功能和改进的细节会很有趣。
索引速度改进
如前所述，在已填充的表上创建 HNSW（分层可导航小世界）索引比先创建索引再插入表中效率更高。这是因为如果预先知道总行数，则更容易预测索引的大小，从而可以将工作分成足够大的块以分配给多个线程。然而，在最初的版本中，这种工作分配有点过于粗粒度，因为我们只会为表中的每个行组（默认约 120,000 行）调度一个额外的辅助线程。
我们现在在索引创建管道中引入了一个额外的缓冲步骤，从而实现了更细粒度的工作分配、更智能的内存分配以及更少的辅助线程之间的争用。这使得在拥有大量可用线程的环境中构建 HNSW 索引时，CPU 利用率大大提高，并显著加速，无论底层表的大小如何。
这项更改的另一个好处是，我们现在可以在构建索引时显示进度条，当您仍然需要等待索引创建完成时（尽管现在更好地利用了系统资源！），这是一个很好的改进。
新的距离函数
在 VSS 的初始版本中，我们支持三种不同的距离函数：array_distance、array_cosine_similarity 和 array_inner_product。然而，只有 array_distance 函数才真正是“距离”函数，因为它在向量相似时返回接近 0 的结果，而在向量不相似时返回接近 1 的结果，这与例如 array_cosine_similarity 在向量完全相同时返回 1 的情况形成对比。哎呀！
为了弥补这一点，我们引入了两个新的距离函数
array_cosine_distance，等同于1 - array_cosine_simililarityarray_negative_inner_product，等同于-array_inner_product
现在这些函数将通过 HNSW 索引加速，从而使所有支持的度量指标的查询模式和排序保持一致，无论您是否使用 HNSW 索引。此外，如果您使用例如 cosine 度量的 HNSW，并编写一个使用 1 - array_cosine_similarity 作为排名标准的 top-k 样式查询，优化器应该能够将表达式规范化为 array_cosine_distance 并也为此函数使用索引。
为了完整性，我们还为动态大小的 LIST 数据类型添加了等效的距离函数（前缀为 list_ 而不是 array_），并将 <=> 二进制运算符更改为 array_cosine_distance 的别名，以匹配 PostgreSQL 的 pgvector 扩展的语义。
索引加速“Top-K”聚合
自上次以来，DuckDB 核心的另一个很酷的进展是，DuckDB 现在为 min_by 和 max_by 聚合函数（及其别名 arg_min 和 arg_max）提供了额外的重载。这些新的重载接受一个可选的第三个参数 n，该参数指定要保留的 top-k（或 top-n）元素的数量，并将它们输出到一个已排序的 LIST 值中。示例如下：
-- Create a table with some example data
CREATE OR REPLACE TABLE vecs AS
SELECT
row_number() OVER () AS id,
[a, b, c]::FLOAT[3] AS vec
FROM
range(1,4) AS x(a), range(1,4) AS y(b), range(1,4) AS z(c);
-- Find the top 3 rows with the vector closest to [2, 2, 2]
SELECT
arg_min(vecs, array_distance(vec, [2, 2, 2]::FLOAT[3]), 3)
FROM
vecs;
[{'id': 14, 'vec': [2.0, 2.0, 2.0]}, {'id': 13, 'vec': [2.0, 1.0, 2.0]}, {'id': 11, 'vec': [1.0, 2.0, 2.0]}]
当然，VSS 扩展现在包含优化器规则，当排序输入是引用已索引向量列的距离函数时，可以使用 HNSW 索引来加速这些 top-k 聚合，类似于我们在上一篇博客文章中讨论的 SELECT a FROM b ORDER BY array_distance(a.vec, query_vec) LIMIT k 查询模式。这些新的重载允许您以更简洁易读的方式表达相同的查询，同时仍然避免对底层表进行全表扫描和排序（只要该表具有匹配的 HNSW 索引）。
索引加速 LATERAL 连接
在对 VSS 的初始版本进行一些基准测试后，我们意识到，尽管我们 HNSW 索引上的索引查找非常快（这要归功于其所基于的 USearch 库！），但与其他解决方案相比，使用 DuckDB 一次搜索单个向量会产生大量的延迟。造成这种情况的原因有很多且很细微，但我们想明确一点，我们选择的 HNSW 实现 USearch 在这里并不是瓶颈，因为性能分析显示只有大约 2% 的运行时实际花费在 usearch 内部。
相反，大多数按查询开销来自于 DuckDB 尚未针对“点查询”进行优化，即那些只真正获取和处理单行的查询。由于 DuckDB 基于向量化执行引擎，最小的工作单元不是 1 行，而是 2,048 行，并且因为我们期望处理大量数据，我们通常倾向于预先花费大量时间来优化查询计划并预分配大型缓冲区和缓存，以便一旦开始执行，一切都尽可能高效。但是当实际工作集如此之小时，许多这些工作就变得不必要了。例如，如果您知道结果中只有少数几行，那么检查并散列常量 768 长度查询向量的每个元素以尝试查找公共子表达式真的值得吗？
虽然我们对未来如何改进这种情况有一些想法，但我们现在决定采取另一种方法，转而专注于我们的优势而不是劣势。那就是，处理大量数据！因此，我们不尝试优化“1:N”查询，即“给定一个嵌入，给我最接近的 N 个嵌入”，而是专注于“N:M”查询，“给定所有 N 个嵌入，将它们分别与最接近的 M 个嵌入配对”。这会是什么样子呢？当然，那就是一个 LATERAL 连接！
基本上，我们现在能够利用 HNSW 索引来加速 LATERAL 连接，其中“内部”查询看起来就像我们通常针对的 top-k 样式查询，例如
SELECT a
FROM b
ORDER BY array_distance(a.vec, query_vec)
LIMIT k;
但是 query_vec 数组现在引用了一个“外部”连接表。唯一的要求是内部表在其向量列上有一个匹配距离函数的 HNSW 索引。示例如下：
-- Set the random seed for reproducibility
SELECT setseed(0.42);
-- Create some example tables
CREATE TABLE queries AS
SELECT
i AS id,
[random(), random(), random()]::FLOAT[3] AS embedding
FROM generate_series(1, 10_000) r(i);
CREATE TABLE items AS
SELECT
i AS id,
[random(), random(), random()]::FLOAT[3] AS embedding
FROM generate_series(1, 10_000) r(i);
-- Collect the 5 closest items to each query embedding
SELECT queries.id AS id, list(inner_id) AS matches
FROM queries, LATERAL (
SELECT
items.id AS inner_id,
array_distance(queries.embedding, items.embedding) AS dist
FROM items
ORDER BY dist
LIMIT 5
)
GROUP BY queries.id;
在我的配备 Apple M3 Pro 芯片和 36 GB 内存的 MacBook 上执行此操作大约需要 10 秒。
如果我们 EXPLAIN 这个查询计划，我们会看到许多高级操作符
PRAGMA explain_output = 'optimized_only';
EXPLAIN ...
普通查询计划（操作符和预期基数）
┌───────────────────────────┐
│ PROJECTION │
│ ──────────────────── │
│ ~5000000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ HASH_GROUP_BY │
│ ──────────────────── │
│ ~5000000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ PROJECTION │
│ ──────────────────── │
│ ~10000000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ PROJECTION │
│ ──────────────────── │
│ ~10000000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ RIGHT_DELIM_JOIN │
│ ──────────────────── │
│ ~10000000 Rows ├──────────────┐
└─────────────┬─────────────┘ │
┌─────────────┴─────────────┐┌─────────────┴─────────────┐
│ SEQ_SCAN ││ HASH_JOIN │
│ ──────────────────── ││ ──────────────────── │
│ ~10000 Rows ││ ~10000000 Rows ├──────────────┐
└───────────────────────────┘└─────────────┬─────────────┘ │
┌─────────────┴─────────────┐┌─────────────┴─────────────┐
│ PROJECTION ││ DUMMY_SCAN │
│ ──────────────────── ││ │
│ ~10000000 Rows ││ │
└─────────────┬─────────────┘└───────────────────────────┘
┌─────────────┴─────────────┐
│ FILTER │
│ ──────────────────── │
│ ~10000000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ PROJECTION │
│ ──────────────────── │
│ ~50000000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ WINDOW │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ PROJECTION │
│ ──────────────────── │
│ ~50000000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ CROSS_PRODUCT ├──────────────┐
└─────────────┬─────────────┘ │
┌─────────────┴─────────────┐┌─────────────┴─────────────┐
│ SEQ_SCAN ││ DELIM_SCAN │
│ ──────────────────── ││ ──────────────────── │
│ ~10000 Rows ││ ~5000 Rows │
└───────────────────────────┘└───────────────────────────┘
虽然这个计划看起来非常复杂，但这些操作符中最令人担忧的是计划底部的 CROSS_PRODUCT，它会使预期基数暴增，这表明我们正在做大量可能不想做的工作。然而，如果我们在 items 表上使用以下方式创建 HNSW 索引：
CREATE INDEX my_hnsw_idx ON items USING HNSW(embedding);
并重新运行 EXPLAIN，我们得到的是这个计划：
带 HNSW 索引的查询计划（操作符和预期基数）
┌───────────────────────────┐
│ PROJECTION │
│ ──────────────────── │
│ ~50000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ HASH_GROUP_BY │
│ ──────────────────── │
│ ~50000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ PROJECTION │
│ ──────────────────── │
│ ~50000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ PROJECTION │
│ ──────────────────── │
│ ~50000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ PROJECTION │
│ ──────────────────── │
│ ~50000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ HNSW_INDEX_JOIN │
│ ──────────────────── │
│ ~50000 Rows │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│ SEQ_SCAN │
│ ──────────────────── │
│ ~10000 Rows │
└───────────────────────────┘
我们可以看到这个计划被大大简化了，但最重要的是，新的 HNSW_INDEX_JOIN 操作符取代了之前的 CROSS_PRODUCT 节点，并且估计基数从 5,000,000 降至 50,000！现在执行此查询大约需要 0.15 秒。这几乎是 66 倍的加速！
此优化最近才添加到 VSS 扩展中，因此如果您已经为 DuckDB v1.1.2 安装了 vss，请运行以下命令以获取最新版本：
UPDATE EXTENSIONS (vss);
结论
各位，本次更新就到这里！我们希望您喜欢 DuckDB 向量相似度搜索扩展的这次更新。虽然本次更新重点介绍了新功能和改进，例如更快的索引、额外的距离函数和更多的优化器规则，但我们仍在努力改进上一篇博客文章中提到的一些限制。我们希望很快能分享更多关于自定义索引和基于索引的优化的内容！如果您有任何问题或反馈，请随时通过 duckdb-vss GitHub 仓库或 DuckDB Discord 与我们联系。期待与您再见！
