在 SQLAlchemy 中使用平凯数据库向量搜索

本文档将展示如何使用 SQLAlchemy 与 TiDB 向量搜索进行交互，以及如何存储向量和执行向量搜索查询。

前置需求

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

  • 在你的机器上安装 Python 3.8 或更高版本
  • 在你的机器上安装 Git
  • 准备一个 TiDB 集群

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

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

运行示例应用程序

你可以通过以下步骤快速了解如何在 SQLAlchemy 中使用 TiDB 向量搜索。

第 1 步：克隆示例代码仓库

将 tidb-vector-python 仓库克隆到本地：

git clone https://github.com/pingcap/tidb-vector-python.git

第 2 步：创建虚拟环境

为你的项目创建虚拟环境：

cd tidb-vector-python/examples/orm-sqlalchemy-quickstart
python3 -m venv .venv
source .venv/bin/activate

第 3 步：安装所需的依赖

安装示例项目所需的依赖：

pip install -r requirements.txt

你也可以直接为项目安装以下依赖项：

pip install pymysql python-dotenv sqlalchemy tidb-vector

第 4 步：配置环境变量

对于本地部署的 TiDB，请在 Python 项目的根目录下新建一个 .env 文件，将以下内容复制到 .env 文件中，并根据集群的启动参数修改环境变量值为 TiDB 实际对应的值：

TIDB_DATABASE_URL=mysql+pymysql://<USERNAME>:<PASSWORD>@<HOST>:<PORT>/<DATABASE>
# 例如：TIDB_DATABASE_URL="mysql+pymysql://root@127.0.0.1:4000/test"

如果你在本机运行 TiDB，<HOST> 默认为 127.0.0.1。<PASSWORD> 初始密码为空，若你是第一次启动集群，则无需带上此字段。

以下为各参数的解释：

  • <USERNAME>：连接 TiDB 集群的用户名。
  • <PASSWORD>：连接 TiDB 集群的密码。
  • <HOST>：TiDB 集群的主机。
  • <PORT>：TiDB 集群的端口。
  • <DATABASE>：要连接的数据库名称。

第 5 步：运行示例应用程序

python sqlalchemy-quickstart.py

输出示例：

Get 3-nearest neighbor documents:
  - distance: 0.00853986601633272
    document: fish
  - distance: 0.12712843905603044
    document: dog
  - distance: 0.7327387580875756
    document: tree
Get documents within a certain distance:
  - distance: 0.00853986601633272
    document: fish
  - distance: 0.12712843905603044
    document: dog

示例代码片段

你可以参考以下示例代码片段来完成自己的应用程序开发。

创建向量表

连接到平凯数据库集群

import os
import dotenv

from sqlalchemy import Column, Integer, create_engine, Text
from sqlalchemy.orm import declarative_base, Session
from tidb_vector.sqlalchemy import VectorType

dotenv.load_dotenv()

tidb_connection_string = os.environ['TIDB_DATABASE_URL']
engine = create_engine(tidb_connection_string)

定义向量列

创建一个表格，其中包含一个向量数据类型的 embedding 列，用于存储三维向量。

Base = declarative_base()

class Document(Base):
    __tablename__ = 'sqlalchemy_demo_documents'
    id = Column(Integer, primary_key=True)
    content = Column(Text)
    embedding = Column(VectorType(3))

存储包含向量的 document

with Session(engine) as session:
   session.add(Document(content="dog", embedding=[1, 2, 1]))
   session.add(Document(content="fish", embedding=[1, 2, 4]))
   session.add(Document(content="tree", embedding=[1, 0, 0]))
   session.commit()

搜索近邻向量

可以选择使用余弦距离 (CosineDistance) 函数，查询与向量 [1, 2, 3] 语义最接近的前 3 个 document。

with Session(engine) as session:
   distance = Document.embedding.cosine_distance([1, 2, 3]).label('distance')
   results = session.query(
      Document, distance
   ).order_by(distance).limit(3).all()

搜索一定距离内的向量

可以选择使用余弦距离 (CosineDistance) 函数，查询与向量 [1, 2, 3] 的余弦距离小于 0.2 的向量。

with Session(engine) as session:
    distance = Document.embedding.cosine_distance([1, 2, 3]).label('distance')
    results = session.query(
        Document, distance
    ).filter(distance < 0.2).order_by(distance).limit(3).all()

另请参阅

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