DuckDB也要加入向量数据库的战局?

小半 • 2024年2月22日下午12:34 • 微信精选 • 阅读 373

广告时间：

• 长期承接数据处理、分析、建模工作
• 软件开发工作
• 培训
• 大模型相关服务

DuckDB的引入固定长度数据类型Array

不知道大家在看DuckDB 0.10.0 发布了，带来了重大更新的时候有没有留意到DuckDB 引入了新的数据类型Array，这里的 Array 不同于Postgres的数组类型（Postgres的是可变长度的），DuckDB的 Array 列存储固定大小的数组。该列中的所有字段都必须具有相同的长度和相同的基础类型。数组通常用于存储数字数组，但可以包含任何统一的数据类型，包括 ARRAY、LIST 和 STRUCT 类型。

并且DuckDB特别强调Array可用于存储诸如词嵌入或图像嵌入之类的向量。另外DuckDB还对应的引入了向量的相似度函数array_cosine_similarity，下面再看下官方的例子

-- create sample data
CREATE TABLE x (i INT, v FLOAT[3]);
CREATE TABLE y (i INT, v FLOAT[3]);
INSERT INTO x VALUES (1, array_value(1.0::FLOAT, 2.0::FLOAT, 3.0::FLOAT));
INSERT INTO y VALUES (1, array_value(2.0::FLOAT, 3.0::FLOAT, 4.0::FLOAT));
-- compute cross product
SELECT array_cross_product(x.v, y.v)
FROM x, y
WHERE x.i = y.i;
-- compute cosine similarity
SELECT array_cosine_similarity(x.v, y.v)
FROM x, y
WHERE x.i = y.i;

下面再来一个网上的例子^[1]

Python的例子

Insert data from numpy arrays

import duckdb
import numpy as np

conn = duckdb.connect(database=":memory:", read_only=False)

conn.execute("CREATE TABLE data (id INTEGER, vector FLOAT4[768]);")


def normalize(vec: np.ndarray) -> np.ndarray:
    return vec / np.linalg.norm(vec)


# insert 3000 random vectors, each vector has 768 dimensions
for i in range(3000):
    vector = np.random.rand(768).astype("float32")
    # normalize vector before inserting
    norm_vector = normalize(vector)
    conn.execute("INSERT INTO data VALUES (?, ?)", (i, vector))

关于数据类型的注释：

• DuckDB FLOAT4：单精度浮点数（4 字节）来源^[2]
• Numpy float32：单精度浮点数类型/32 位精度浮点数类型。来源^[3]

所以 DuckDB的 FLOAT4 = numpy.float32

查询数据

我们需要显式地将 numpy 数组转换为 FLOAT4[768].

query_vector = query = np.random.rand(768).astype("float32")
norm_query_vector = normalize(query_vector)

result = conn.execute(
    """SELECT id, array_cosine_similarity(vector, $query_vector::FLOAT4[768]) as cosim 
    FROM data ORDER BY cosim DESC LIMIT 10""",
    {"query_vector": norm_query_vector},
).fetchdf()
result

这将给我们一个最相似的 10 个向量的表格，按余弦相似度排序。

id	cosim
462	0.780559
1652	0.778595
1494	0.777005
344	0.776584
2760	0.776564
522	0.776486
2446	0.775915
1614	0.775862
1313	0.775783
2129	0.775757

结束语

当然要作为一个Vector DB，至少还需要引入对近似最近邻进行索引和搜索（例如，使用一些 HNSW 实现或类似实现）。

引用链接

[1] 网上的例子: https://ricardoanderegg.com/posts/duckdb-vector-database/
[2] 来源: https://duckdb.org/docs/sql/data_types/numeric#floating-point-types
[3] 来源: https://numpy.org/devdocs/reference/arrays.scalars.html#numpy.single