使用 SQL 开始向量搜索

TiDB 扩展了 MySQL 语法以支持向量搜索，并引入了向量数据类型和多个向量函数。

本文将展示如何使用 SQL 语句在 TiDB 中进行向量搜索。在本文中，你将使用 MySQL 命令行客户端完成以下任务：

  • 连接到 TiDB 集群
  • 创建向量表
  • 存储向量嵌入
  • 执行向量搜索查询

前置需求

为了能够顺利完成本文中的操作，你需要提前：

如果你还没有 TiDB 集群，可以按照以下方式创建：

参考部署本地测试集群或部署正式集群，创建本地集群。

快速开始

第 1 步：连接到平凯数据库集群

在本地部署的集群启动后，在终端中执行你的集群连接命令：

以下为 macOS 上的连接命令示例：

mysql --comments --host 127.0.0.1 --port 4000 -u root

第 2 步：创建向量表

创建表时，你可以使用 VECTOR 数据类型声明指定列为向量列。

例如，要创建一个带有三维 VECTOR 列的 embedded_documents 表，可以使用 MySQL CLI 执行以下 SQL 语句：

USE test;
CREATE TABLE embedded_documents (
    id        INT       PRIMARY KEY,
    -- document 列存储文档的原始内容
    document  TEXT,
    -- embedding 列存储文档的向量表示
    embedding VECTOR(3)
);

预期输出如下：

Query OK, 0 rows affected (0.27 sec)

第 3 步：向表中插入向量

向 embedded_documents 表中插入三行，每一行包含数据和数据的向量嵌入：

INSERT INTO embedded_documents
VALUES
    (1, 'dog', '[1,2,1]'),
    (2, 'fish', '[1,2,4]'),
    (3, 'tree', '[1,0,0]');

预期输出如下：

Query OK, 3 rows affected (0.15 sec)
Records: 3  Duplicates: 0  Warnings: 0

第 4 步：查询向量表

要验证上一步中的三行数据是否已正确插入，可以查询 embedded_documents 表：

SELECT * FROM embedded_documents;

预期输出如下：

+----+----------+-----------+
| id | document | embedding |
+----+----------+-----------+
|  1 | dog      | [1,2,1]   |
|  2 | fish     | [1,2,4]   |
|  3 | tree     | [1,0,0]   |
+----+----------+-----------+
3 rows in set (0.15 sec)

第 5 步：执行向量搜索查询

与全文搜索类似，在使用向量搜索时，你需要提供搜索词。

在本例中，搜索词是“一种会游泳的动物”，假设其对应的向量是 [1,2,3]。在实际应用中，你需要使用嵌入模型将用户的搜索词转换为向量。

执行以下 SQL 语句后，TiDB 会计算 [1,2,3] 与表中各向量之间的余弦距离 (vec_cosine_distance)，然后对这些距离进行排序并输出表中最接近搜索向量（余弦距离最小）的前三个文档。

SELECT id, document, vec_cosine_distance(embedding, '[1,2,3]') AS distance
FROM embedded_documents
ORDER BY distance
LIMIT 3;

预期输出如下：

+----+----------+---------------------+
| id | document | distance            |
+----+----------+---------------------+
|  2 | fish     | 0.00853986601633272 |
|  1 | dog      | 0.12712843905603044 |
|  3 | tree     |  0.7327387580875756 |
+----+----------+---------------------+
3 rows in set (0.15 sec)

搜索结果中的三个词按它们与查询向量的距离排序：距离越小，对应的 document 越相关。

因此，从输出结果来看，会游泳的动物很可能是一条鱼 (fish)，或者是一只有游泳天赋的狗 (dog)。

另请参阅

  • 向量数据类型
  • 向量搜索索引
下载 PDF