自建Qdrant迁移至AnalyticDB PostgreSQL版

前提条件

• 已创建Qdrant集群。

• 已安装Python环境，建议使用Python 3.8及以上版本。

• 已安装所需的Python库。

  pip install psycopg2
  pip install qdrant-client==1.6.0
  pip install pyaml
  pip install tqdm

迁移操作

步骤一：导出Qdrant

1. 准备好导出脚本export.py及其导出配置文件qdrant2csv.yaml，并创建输出目录，本文以output为例。

   导出脚本export.py如下。

   import yaml
   import json
   from qdrant_client import QdrantClient
   import os
   from enum import IntEnum
   from tqdm import tqdm
   
   with open("./qdrant2csv.yaml", "r") as f:
       config = yaml.safe_load(f)
   
   
   print("configuration:")
   print(config)
   
   qdrant_config = config["qdrant"]
   
   
   class DataType(IntEnum):
       ID = 1
       FLOAT_VECTOR = 2
       JSON = 3
   
   
   def data_convert_to_str(data, dtype, delimeter):
       if dtype == DataType.ID:
           return str(data)
       elif dtype == DataType.FLOAT_VECTOR:
           return "{" + ", ".join(str(x) for x in data) + "}"
       elif dtype == DataType.JSON:
           return str(data).replace(delimeter, f"\\{delimeter}").replace("\"", "\\\"")
       Exception(f"Unsupported DataType {dtype}")
   
   
   def csv_write_rows(datum, fd, fields_types, delimiter="|"):
       for data in datum:
           for i in range(len(data)):
               data[i] = data_convert_to_str(data[i], fields_types[i], delimiter)
           fd.write(delimiter.join(data) + "\n")
   
   
   def csv_write_header(headers, fd, delimiter="|"):
       fd.write(delimiter.join(headers) + "\n")
   
   
   def dump_collection(collection_name: str):
       results = []
       file_cnt = 0
       print("connecting to qdrant...")
       client = QdrantClient(**qdrant_config)
   
       export_config = config["export"]
       tmp_path = os.path.join(export_config["output_path"], collection_name)
       if not os.path.exists(tmp_path):
           os.mkdir(tmp_path)
   
       # fetch info of collection
       fields_meta_list = ["id bigint"]
       fields_types = [DataType.ID]
       headers = ["id"]
       collection = client.get_collection(collection_name)
       total_num = collection.points_count
       if isinstance(collection.config.params.vectors, dict):
           # multi vectors
           for vec_name in collection.config.params.vectors.keys():
               fields_types.append(DataType.FLOAT_VECTOR)
               fields_meta_list.append(f"{vec_name} real[]")
               headers.append(vec_name)
       else:
           # single vector
           fields_types.append(DataType.FLOAT_VECTOR)
           fields_meta_list.append("vector real[]")
           headers.append("vector")
   
       fields_types.append(DataType.JSON)
       fields_meta_list.append("payload json")
       headers.append("payload")
   
       fields_meta_str = ','.join(fields_meta_list)
       create_table_sql = f"CREATE TABLE {collection_name} " \
                          f" ({fields_meta_str});"
   
       with open(os.path.join(export_config["output_path"], collection_name, "create_table.sql"), "w") as f_d:
           f_d.write(create_table_sql)
   
       print(create_table_sql)
   
       def write_to_csv_file(col_names, data):
           if len(results) == 0:
               return
           nonlocal file_cnt
           assert(file_cnt <= 1e9)
           output_file_name = os.path.join(export_config["output_path"], collection_name, f"{str(file_cnt).zfill(10)}.csv")
           with open(output_file_name, "w", newline="") as csv_file:
               # write header
               csv_write_header(col_names, csv_file)
               # write data
               csv_write_rows(data, csv_file, fields_types)
               file_cnt += 1
               results.clear()
   
       offset_id = None
   
       with tqdm(total=total_num, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}") as pbar:
           while True:
               res = client.scroll(collection_name=collection_name,
                                   limit=1000,
                                   offset=offset_id,
                                   with_payload=True,
                                   with_vectors=True)
   
               records = res[0]
               for record in records:
                   # append id
                   record_list = [record.id]
                   # append vectors
                   if isinstance(record.vector, dict):
                       # multi vector
                       for vector_name in headers[1:-1]:
                           record_list.append(record.vector[vector_name])
                   else:
                       # single vector
                       record_list.append(record.vector)
                   # append payload
                   record_list.append(json.dumps(record.payload, ensure_ascii=False))
                   results.append(record_list)
   
                   if len(results) >= export_config["max_line_in_file"]:
                       write_to_csv_file(headers, data=results)
   
                   pbar.update(1)
   
               if len(res) == 0 or len(res[0]) == 0 or res[1] is None:
                   # finished
                   break
               else:
                   offset_id = res[1]
   
       write_to_csv_file(headers, data=results)
   
   
   for name in config["export"]["collections"]:
       dump_collection(name)
   

   导出配置文件qdrant2csv.yaml如下。

   qdrant:  # 连接Qdrant的配置项
       host: 'localhost'  # Qdrant的主机地址
       port: 6333        # Qdrant端口号，默认值：6433 
       grpc_port: 6434   # gRPC 接口的端口，默认值: 6334
       api_key: ''  # 用于在Qdrant Cloud中进行身份验证的API密钥 
       url: ''      # 主机名或者字符串"Optional[scheme], host, Optional[port], Optional[prefix]"
       location: '' # 填写memory，则以内存模式连接实例；填写普通的字符串，则等同于填写url配置项；如果不填，则会通过host+port连接实例        
   
   export:
      collections:
       - 'test_collection'
       - 'multi'                 # 填写所有需要导出的collection
     max_line_in_file: 40000     # 文件切分行数
     output_path: './output'     # 导出目标目录
   

2. 将导出脚本export.py、导出配置文件qdrant2csv.yaml及输出目录output存放至同一个目录下。目录层级如下。

   ├── export.py
   ├── qdrant2csv.yaml
   └── output

3. 根据Qdrant集群信息，修改qdrant2csv.yaml中的配置项。

4. 运行Python脚本，并查看输出结果。

   python export.py

   输出结果如下。

   .
   ├── export.py
   ├── qdrant2csv.yaml
   └── output
       ├── test_collection
       │   ├── 0000000000.csv
       │   ├── 0000000001.csv
       │   ├── 0000000002.csv
       │   └── create_table.sql
       └── multi
           ├── 0000000000.csv
           └── create_table.sql

步骤二：导入AnalyticDB PostgreSQL版向量数据库

1. 准备好导入脚本import.py、导入配置文件csv2adbpg.yaml及需要导入的数据data（即在导出步骤中得到的output目录）。

   导入脚本import.py如下。

   import psycopg2
   import yaml
   import glob
   import os
   
   if __name__ == "__main__":
       with open('csv2adbpg.yaml', 'r') as config_file:
           config = yaml.safe_load(config_file)
   
       print("current config:" + str(config))
   
       db_host = config['database']['host']
       db_port = config['database']['port']
       db_name = config['database']['name']
       schema_name = config['database']['schema']
       db_user = config['database']['user']
       db_password = config['database']['password']
       data_path = config['data_path']
   
       conn = psycopg2.connect(
           host=db_host,
           port=db_port,
           database=db_name,
           user=db_user,
           password=db_password,
           options=f'-c search_path={schema_name},public'
       )
   
       cur = conn.cursor()
   
       # check schema
       cur.execute("SELECT schema_name FROM information_schema.schemata WHERE schema_name = %s", (schema_name,))
       existing_schema = cur.fetchone()
       if existing_schema:
           print(f"Schema {schema_name} already exists.")
       else:
           # create schema
           cur.execute(f"CREATE SCHEMA {schema_name}")
           print(f"Created schema: {schema_name}")
   
       for table_name in os.listdir(data_path):
           table_folder = os.path.join(data_path, table_name)
           print(f"Begin Process table: {table_name}")
           if os.path.isdir(table_folder):
               create_table_file = os.path.join(table_folder, 'create_table.sql')
               with open(create_table_file, 'r') as file:
                   create_table_sql = file.read()
               try:
                   cur.execute(create_table_sql)
               except psycopg2.errors.DuplicateTable as e:
                   print(e)
                   conn.rollback()
                   continue
               print(f"Created table: {table_name}")
   
               cnt = 0
               csv_files = glob.glob(os.path.join(table_folder, '*.csv'))
               for csv_file in csv_files:
                   with open(csv_file, 'r') as file:
                       copy_command = f"COPY {table_name} FROM STDIN DELIMITER '|' HEADER"
                       cur.copy_expert(copy_command, file)
                   cnt += 1
                   print(f"Imported data from: {csv_file} | {cnt}/{len(csv_files)} file(s) Done")
   
           conn.commit()
           print(f"Finished import table: {table_name}")
           print('#'*60)
   
       cur.close()
       conn.close()
   

   导入配置文件csv2adbpg.yaml如下。

   database:
       host: "192.16.XX.XX"         # AnalyticDB PostgreSQL版实例的外网地址
       port: 5432                   # AnalyticDB PostgreSQL版实例的端口号
       name: "vector_database"      # 导入目标数据库名 
       user: "username"             # AnalyticDB PostgreSQL版实例的数据库账号
       password: ""                 # 账号密码
       schema: "public"             # 导入Schama名，若不存在则会自动创建
   
   data_path: "./data"            # 导入数据源
   

2. 将导入脚本import.py和导入配置文件csv2adbpg.yaml与需要导入的数据data存放在同一目录下。目录层级如下。

   .
   ├── csv2adbpg.yaml
   ├── data
   │   ├── test_collection
   │   │   ├── 0000000000.csv
   │   │   ├── 0000000001.csv
   │   │   ├── 0000000002.csv
   │   │   └── create_table.sql
   │   └── multi
   │       ├── 0000000000.csv
   │       └── create_table.sql
   └── import.py

3. 根据AnalyticDB PostgreSQL版实例信息，修改csv2adbpg.yaml文件中配置项。

4. 运行Python脚本。

   python import.py

5. 在AnalyticDB PostgreSQL版向量数据库中检查数据是否正常导入。

6. 重建所需要的索引。具体操作，请参见创建向量索引。

相关文档

更多关于Qdrant的介绍，请参见Qdrant产品文档。