通过 SQL 快速上手 TiDB + AI

TiDB 扩展了 MySQL 语法以支持向量搜索，并引入了新的向量数据类型以及若干向量函数。

本文档演示了如何仅使用 SQL 语句快速上手 TiDB 向量搜索。你将学习如何使用 MySQL 命令行客户端完成以下操作：

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

前置条件

要完成本文档中的步骤，你需要：

• 在本地安装 MySQL 命令行客户端（MySQL CLI）。
• 一个 TiDB 集群。

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

• （推荐）创建一个 TiDB Cloud Starter 实例。
• 部署本地测试 TiDB Self-Managed 集群或部署生产环境 TiDB Self-Managed 集群。

快速上手

步骤 1. 连接到 TiDB

根据你选择的 TiDB 部署方式连接到 TiDB。

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,
    -- Column to store the original content of the document.
    document  TEXT,
    -- Column to store the vector representation of the document.
    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. 执行向量搜索查询

与全文搜索类似，用户在使用向量搜索时会向应用提供搜索词。

本例中，搜索词为 “a swimming animal”，其对应的向量嵌入假定为 [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 与搜索向量的相关性越高。

因此，根据输出，最有可能的“swimming animal”是 fish，其次可能是会游泳的 dog。

另请参阅

• 向量数据类型
• 向量搜索索引