Python SDK对接向量数据库

Python SDK对接向量数据库

本文介绍如何使用Python语言调用openGauss向量数据库

环境准备

限制：如果数据库非OM工具安装，建议python版本为3.11及以上如果数据库是OM工具安装，建议python版本为3.6-3.10

在线安装

1. pip3 install psycopg2

注意这里安装的是pypi上的psycopg2的包，不包含多向量查询的特性。如果要使用多向量查询的特性，参考离线安装。 - 离线安装 1）下载适配好openGauss的psycopg2包，下载链接：gitcode官网。 2）进入openGauss-connector-python-psycopg2根目录，执行

1. sh build.sh -bd /data/compile/openGauss-server/dest/ -v 5.0.0

-bd: 指定openGauss数据库构建结果目录 -v: 指定构建包的版本号。不指定则默认为5.0.0 编译完成后的驱动，在 output 目录下，解压安装包后，会得到两个目录 lib 和 psycopg2。 3) 将 psycopg2 目录拷贝到 Python 解释器的 site-packages 下（需要先pip3 uninstall psycopg2），对于lib文件夹需要设置环境变量。

1. echo "export LD_LIBRARY_PATH=[/path/to/lib]:$LD_LIBRARY_PATH" >> ~/.bashrc 2. source ~/.bashrc

基本操作

1.连接数据库

1. import psycopg2
2. from psycopg2 import sql
3. import numpy as np
4. from typing import List

6. def create_connection(dbname:str, user:str, password:str, host:str, port:int):
7. conn = psycopg2.connect(
8. dbname = dbname,
9. user = user,
10. password = password,
11. host = host,
12. port = port
13. )
14. cursor = conn.cursor()
15. return conn, cursor

2.创建表

1. def create_table(conn, cursor, table_name:str, dim:int):
2. cursor.execute(
3. sql.SQL(
4. "CREATE TABLE IF NOT EXISTS public.{table_name} (id BIGINT PRIMARY KEY, embedding vector({dim}));"
5. ).format(table_name = sql.Identifier(table_name), dim = sql.Literal(dim))
6. )
7. conn.commit()

3.创建索引

1. def create_index(conn, cursor, table_name:str, index_name:str):
2. cursor.execute(
3. sql.SQL(
4. """
5. CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name}
6. USING hnsw (embedding vector_l2_ops);
7. """
8. ).format(index_name = sql.Identifier(index_name), table_name = sql.Identifier(table_name))
9. )
10. conn.commit()

4.插入/删除/更新数据

批量插入

1. def insert(conn, cursor, table_name:str, embeddings:List[List[float]], ids:List[int]):
2. data = list(zip(ids, embeddings))
3. cursor.executemany(
4. sql.SQL("INSERT INTO public.{table_name} (id, embedding) VALUES(%s, %s);")
5. .format(table_name = sql.Identifier(table_name)), data
6. )
7. conn.commit()

删除

1. def delete(conn, cursor, table_name:str, ids:List[int]):
2. cursor.execute(
3. sql.SQL(
4. "DELETE FROM public.{table_name} WHERE id IN ({ids});"
5. ).format(table_name = sql.Identifier(table_name), ids = sql.SQL(',').join(map(sql.Literal, ids)))
6. )
7. delete_count = cursor.rowcount
8. conn.commit()
9. return delete_count

更新

1. def update(conn, cursor, table_name:str, id:int, vector:List[List[float]]):
2. cursor.execute(
3. sql.SQL(
4. "UPDATE public.{table_name} SET embedding = %s WHERE id = %s;"
5. ).format(table_name = sql.Identifier(table_name)), (vector ,id)
6. )
7. conn.commit()

5.查询

1. # 串行批量查询
2. def select(conn, cursor, table_name:str, queries:List[List[float]], topk:int):
3. ids = []
4. for emb in queries:
5. cursor.execute(
6. sql.SQL(
7. "SELECT * FROM public.{table_name} ORDER BY embedding <-> %s::vector LIMIT %s::int;"
8. ).format(table_name = sql.Identifier(table_name)), (emb, topk)
9. )
10. conn.commit()
11. result = cursor.fetchall()
12. ids.append([int(i[0]) for i in result])
13. return ids

6.删除表

1. def drop_table(conn, cursor, table_name:str):
2. cursor.execute(
3. sql.SQL(
4. "DROP TABLE IF EXISTS public.{table_name};"
5. ).format(table_name = sql.Identifier(table_name))
6. )
7. conn.commit()

7.关闭连接

1. def close_connection(conn, cursor):
2. conn.close()
3. cursor.close()

8.多向量并发查询

多向量召回支持在单次搜索请求中同时提交多个查询向量，openGauss将并行对查询向量进行搜索，并返回多组结果。

函数名

1. execute_multi_search(dbconfig, conn_pool_mgr, sql_template, argslist, scan_params, max_workers)

输入参数

dbconfig:数据库连接配置，包含user、password、dbname、host、port
conn_pool_mgr：连接池管理对象，可以自定义设置，当其为None时，函数内部会自行创建
sql_template:查询语句，要求是单条查询语句(select为首单词)、包含向量操作符（<->/<=>/<#>/<+>/<~>/<%>）
argslist：查询参数，需要元组列表的格式，要求不为空
scan_params：需要通过set设置的参数（如hnsw_ef_search、nprobes）
max_workers:连接池最大连接数，和数据库最大连接数（由参数max_connections设置）有关，一般来说，连接池最大连接数要小于数据库最大连接数，但是数据库对于管理员用户的连接限制会略超过max_connections设置。

输出参数

查询结果，形式为[[(1, '[1,2,3]'),(2, '[2,2,2]')], [],...]，表示n个查询向量对应的limit个结果。

使用案例

1. from psycopg2.extras import execute_multi_search, init_conn_pool, close_conn_pool
2. sql_template = "SELECT * FROM test_table1 ORDER BY embedding <-> %s LIMIT %s;"
3. scan_params = {"enable_seqscan": "off", "hnsw_ef_search" : 40}
4. dbconfig = {'user': 'yourusername', 'password': 'xxxxxx', 'host': 'yourhost', 'dbname': 'yourdbname', 'port' : 5432}
5. argslist = [('[1,1,1]', 1), ('[2,2,3]', 2)]

7. conn_pool_mgr = init_conn_pool(dbconfig, 2, scan_params)
8. res = execute_multi_search(dbconfig, conn_pool_mgr, sql_template, argslist, scan_params, 2)
9. close_conn_pool(conn_pool_mgr)

注意：如果并发中的数据库配置密码输入错误，可能会报ERROR: The account has been locked，这时候需要用超级用户登录为当前用户解锁，具体命令为：ALTER ROLE username ACCOUNT UNLOCK;

用例

1. conn, cursor = create_connection("testdb", "test_user", YourPassword, "localhost", 5432)
2. create_table(conn, cursor, "test_table1", 3)
3. create_index(conn, cursor, "test_table1", "idx_test1")
4. insert(conn, cursor, "test_table1", [[1.2, 3, 5], [4.3, 5.2, 1]], [0, 1])
5. delete(conn, cursor, "test_table1", [0])
6. update(conn, cursor, "test_table1", 1, [1, 3, 3])
7. select(conn, cursor, "test_table1", [[1, 2, 2], [3, 5, 1]], 2)
8. drop_table(conn, cursor, "test_table1")
9. close_connection(conn, cursor)