DuckDB 向量相似度搜索带来哪些重大更新？

✍️: Max Gabrielsson

🔗：https://duckdb.org/2024/10/23/whats-new-in-the-vss-extension.html

在 上一篇博客文章 中，我们介绍了 DuckDB 的 向量相似度搜索 (VSS) 扩展[1]。尽管该扩展仍处于实验阶段，我们认为详细说明一些自首次发布以来的新增功能和改进会很有趣。

索引速度提升

如之前文档所述，针对已有数据表创建 HNSW（Hierarchical Navigable Small World graphs，即分层可导航小世界）索引，要比先创建索引再插入数据高效得多。这是因为如果能提前知道表的总行数，就能更准确地预测索引大小，从而将任务分配到多个线程。然而，早期版本中的任务分配过于粗放，每 行组[2] （默认大约 120,000 行）仅安排了一个额外线程。

我们现在引入了一个新的缓冲步骤，使任务分配更为细化，内存分配更加智能，并减少线程间的争抢。这显著提高了 CPU 利用率，无论底层表的大小如何，在多线程环境中构建 HNSW 索引的速度有了明显提升。

此外，这一改进还允许在索引创建时显示进度条，即便索引构建仍需一定时间，但资源利用率已经大幅提升。

新的距离函数

在 VSS 的最初版本中，我们提供了三种距离函数：
array_distance[3]，
array_cosine_similarity[4]，以及 array_inner_product[5]。不过，严格来说，只有 array_distance
 是真正的“距离”函数，它在向量相似时返回接近 0 的值，在不相似时返回接近 1。而例如 array_cosine_similarity
 返回的是向量完全相同时的 1。

为了解决这一问题，我们引入了两个新距离函数：

• • array_cosine_distance
  ，等效于 1 - array_cosine_simililarity
  

• • array_negative_inner_product
  ，等效于 -array_inner_product
  

现在，这些函数都能使用 HNSW 索引进行加速，保证无论使用何种度量标准，查询模式和排序都能保持一致。此外，如果您使用的是 cosine
 度量，并且通过 1 - array_cosine_similarity
 编写 top-k 风格的查询，优化器会将其自动规范为 array_cosine_distance
 并使用索引加速查询。

为了保持一致性，我们还为 LIST 数据类型[6] 添加了相应的距离函数（以 list_
 作为前缀），并将 <=>
 二元操作符改为 array_cosine_distance
 的别名，与 PostgreSQL 的 pgvector 扩展[7] 相匹配。

索引加速的 Top-K
 聚合

自上次更新以来，DuckDB 核心功能中又新增了对 min_by[8] 和 max_by[9] 聚合函数（及其别名 arg_min
 和 arg_max
）的额外支持。

这些新函数增加了一个可选参数 n
，用于指定保留前 n
 个元素，并将它们输出为排序的 LIST
。以下是一个示例：

-- 创建一个包含示例数据的表
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);

-- 查找与向量 [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 聚合。这样，您可以使用更简洁易读的方式表达查询，同时仍然避免对底层表进行全表扫描和排序（只要该表具有匹配的 HNSW 索引）。

索引加速的 LATERAL
 联接

在对 VSS 初始版本的基准测试中，我们发现虽然 HNSW 索引查找速度非常快（得益于其基于的 USearch[10] 库），但 DuckDB 在处理单个向量搜索时的延迟较高。问题不在于 USearch 实现，而在于 DuckDB 并未针对单点查询进行优化（即只处理一行的查询）。因为 DuckDB 基于向量化执行引擎，最小的工作单元是 2,048 行，我们通常会花费大量时间来优化查询计划和预分配内存缓存，以便在实际执行时提高效率。然而，当查询的数据集非常小时，这些优化显得多余。

因此，我们采取了另一种方法，专注于 DuckDB 的长处：处理大规模数据。与其优化“1:N”查询（即“给定一个嵌入，返回最接近的 N 个嵌入”），我们转而优化“N:M”查询（即“给定 N 个嵌入，找到每个嵌入最接近的 M 个匹配项”）。这就是 LATERAL 联接[11] 的用武之地。

我们现在可以使用 HNSW 索引加速 LATERAL
 联接，其中“内层”查询类似于常见的 top-k 风格查询，例如：

SELECT a
FROM b
ORDER BY array_distance(a.vec, query_vec)
LIMIT k;

但其中的 query_vec
 数组现在是“外层”联接表的引用。唯一的要求是内层表的向量列上存在与距离函数匹配的 HNSW 索引。以下是一个示例：

-- 设置随机种子以确保结果可复现
SELECT setseed(0.42);

-- 创建一些示例表
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);

-- 收集与每个查询嵌入最接近的 5 个 items
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;

在配备 36 GB 内存的 Apple M3 Pro 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！执行时间也从 10 秒缩短到约 0.15 秒，速度提升了 66 倍！

这个优化最近才被添加到 VSS 扩展中，如果您已经为 DuckDB v1.1.2 安装了 vss
，可以通过以下命令获取最新版本：

UPDATE EXTENSIONS (vss);

结论

这次更新就到这里！我们希望您喜欢 DuckDB 向量相似度搜索扩展的新功能介绍。这次更新侧重于新特性和优化，如更快的索引构建、额外的距离函数和优化规则，但我们仍在努力改进之前文章中提到的一些限制。希望不久后能分享更多与自定义索引及基于索引的优化相关的内容！如果您有任何问题或反馈，请在 duckdb-vss GitHub 仓库[12] 或 DuckDB Discord[13] 上联系我们。期待您的加入！

引用链接

[1]
 向量相似度搜索 (VSS) 扩展: https://duckdb.org/docs/extensions/vss
[2]
 行组: https://duckdb.org/docs/internals/storage#row-groups
[3]
 array_distance: https://duckdb.org/docs/sql/functions/array#array_distancearray1-array2
[4]
 array_cosine_similarity: https://duckdb.org/docs/sql/functions/array#array_cosine_similarityarray1-array2
[5]
 array_inner_product: https://duckdb.org/docs/sql/functions/array#array_inner_productarray1-array2
[6]
 LIST 数据类型: https://duckdb.org/docs/sql/data_types/list
[7]
 pgvector 扩展: https://github.com/pgvector/pgvector
[8]
 min_by: https://duckdb.org/docs/sql/functions/aggregates#min_byarg-val-n
[9]
 max_by: https://duckdb.org/docs/sql/functions/aggregates#max_byarg-val-n
[10]
 USearch: https://github.com/unum-cloud/usearch
[11]
 LATERAL 联接: https://duckdb.org/docs/sql/query_syntax/from#lateral-joins
[12]
 duckdb-vss GitHub 仓库: https://github.com/duckdb/duckdb_vss
[13]
 DuckDB Discord: https://discord.duckdb.org/