从Milvus迁移至openGauss DataVec
- 环境准备
- 迁移操作

从Milvus迁移至openGauss DataVec

本教程通过Python编程语言，支持将本地Milvus数据迁移至openGauss DataVec实例中

环境准备

已部署2.3 及以上版本的Milvus实例
已部署7.0.0-RC1 及以上版本的openGauss实例，容器部署参考容器镜像安装
已安装3.8 及以上版本的Python环境
已安装涉及的Python库

1. pip3 install psycopg2
2. pip3 install pymilvus==2.4.9
3. pip3 install numpy

迁移操作

参考如下迁移脚本milvus2datavec.py及配置文件config.ini，根据本地部署的Milvus与openGauss进行配置修改：

迁移配置文件config.ini如下：

1. [Milvus]
2. host = localhost
3. port = 19530

5. [openGauss]
6. user = postgres
7. password = xxxxxx
8. port = 5432
9. database = postgres

11. [Table]
12. milvus_collection_name = test
13. opengauss_table_name = test

15. [SparseVector]
16. # openGauss only support 1000 dimensions for sparsevec
17. default_dimension = 1000

19. [Output]
20. folder = output

22. [Migration]
23. cleanup_temp_files = true

milvus2datavec.py迁移脚本如下：

1. import psycopg2
2. import csv
3. from pymilvus import connections, Collection, utility
4. import configparser
5. import numpy as np
6. import os
7. import logging
8. from typing import List, Dict, Any, Optional, Union
9. from datetime import datetime

11. # Configure logging
12. logging.basicConfig(
13. level=logging.INFO,
14. format='%(asctime)s - %(levelname)s - %(message)s',
15. handlers=[logging.FileHandler('migration.log'), logging.StreamHandler()]
16. )
17. logger = logging.getLogger(__name__)

20. class MilvusToOpenGaussMigrator:
21. def __init__(self, config_file: str = 'config.ini'):
22. self.config = self._load_config(config_file)
23. self.csv_file_path = self._get_csv_file_path()
24. self.fields = []
25. self.MAX_WINDOW_SIZE = 16384  # Milvus default max query window
26. self.SPARSE_DIMENSION = self.config.getint('SparseVector', 'default_dimension', fallback=1000)
27. self.MAX_SPARSE_DIMENSION = 1000
28. self.milvus_version = None

30. def _load_config(self, config_file: str) -> configparser.ConfigParser:
31. """Load configuration file"""
32. config = configparser.ConfigParser()
33. try:
34. if not config.read(config_file):
35. raise FileNotFoundError(f"Config file {config_file} not found")
36. return config
37. except Exception as e:
38. logger.error(f"Failed to load config: {e}")
39. raise

41. def _get_csv_file_path(self) -> str:
42. """Generate CSV file path with timestamp"""
43. output_folder = self.config.get('Output', 'folder', fallback='output')
44. os.makedirs(output_folder, exist_ok=True)
45. milvus_collection = self.config.get('Table', 'milvus_collection_name')
46. timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
47. return os.path.join(output_folder, f"{milvus_collection}_{timestamp}.csv")

49. def _connect_milvus(self) -> Collection:
50. """Connect to Milvus"""
51. try:
52. connections.connect(
53. alias="default",
54. host=self.config.get('Milvus', 'host'),
55. port=self.config.get('Milvus', 'port')
56. )

58. self.milvus_version = utility.get_server_version()
59. logger.info(f"Connected to Milvus {self.milvus_version}")

61. collection_name = self.config.get('Table', 'milvus_collection_name')
62. collection = Collection(collection_name)
63. collection.load()
64. self.fields = [field.name for field in collection.schema.fields]
65. logger.info(f"Loaded collection: {collection_name}")
66. return collection
67. except Exception as e:
68. logger.error(f"Milvus connection failed: {e}")
69. raise

71. def _connect_opengauss(self) -> psycopg2.extensions.connection:
72. """Connect to openGauss"""
73. try:
74. conn = psycopg2.connect(
75. user=self.config.get('openGauss', 'user'),
76. password=self.config.get('openGauss', 'password'),
77. host=self.config.get('openGauss', 'host'),
78. port=self.config.get('openGauss', 'port'),
79. database=self.config.get('openGauss', 'database')
80. )
81. logger.info("Connected to openGauss")
82. return conn
83. except Exception as e:
84. logger.error(f"openGauss connection failed: {e}")
85. raise

87. def _process_sparse_vector(self, sparse_data: Union[dict, bytes, list], dimension: int) -> str:
88. """Convert to openGauss SPARSEVEC format: '{indice:value,...}/dim'"""
89. if sparse_data is None:
90. return "NULL"

92. try:
93. # Convert to {index:value} dict
94. if dimension is None or dimension <=0:
95. dimension = self.MAX_SPARSE_DIMENSION

97. sparse_dict = {}

99. if isinstance(sparse_data, dict):
100. sparse_dict = {
101. int(k+1): float(v)
102. for k, v in sparse_data.items()
103. }
104. else:
105. raise ValueError(f"Unsupported format: {type(sparse_data)}")

107. if not sparse_dict:
108. return "{}/" + str(dimension)

110. try:
111. # Sort by index to ensure consistent output

113. sorted_items = sorted(sparse_dict.items(), key=lambda x: x[0])
114. entries = ",".join(f"{k}:{v}" for k, v in sorted_items)
115. return "{" + entries + "}/" + str(dimension)
116. except Exception as sort_error:
117. logger.warning(f"Sorting failed, using unsorted vector: {sort_error}")
118. entries = ",".join(f"{k}:{v}" for k, v in sparse_dict.items())
119. return "{" + entries + "}/" + str(dimension)

121. except Exception as e:
122. logger.error(f"Sparse vector conversion failed: {e}")
123. return "NULL"

125. def _process_field_value(self, value: Any, field_type: str, dimension: Optional[int] = None) -> str:
126. """Convert field value to CSV string"""
127. if value is None:
128. return "NULL"
129. elif field_type == "SPARSE_FLOAT_VECTOR":
130. return self._process_sparse_vector(value, dimension or self.SPARSE_DIMENSION)
131. elif isinstance(value, (list, np.ndarray)):
132. return "[" + ",".join(str(x) for x in value) + "]"
133. elif isinstance(value, dict):
134. return json.dumps(value)
135. else:
136. return str(value)

138. def _create_opengauss_table(self, conn: psycopg2.extensions.connection, collection: Collection) -> None:
139. """Create table with SPARSEVEC columns"""
140. table_name = self.config.get('Table', 'opengauss_table_name')
141. cursor = conn.cursor()

143. try:
144. # Check if table exists
145. cursor.execute(f"SELECT EXISTS(SELECT 1 FROM pg_tables WHERE tablename = '{table_name}');")
146. if cursor.fetchone()[0]:
147. logger.warning(f"Table {table_name} exists")
148. return

150. # Build CREATE TABLE statement
151. columns = []
152. for field in collection.schema.fields:
153. dim = field.dim if hasattr(field, 'dim') else None
154. pg_type = self._milvus_to_opengauss_type(field.dtype.name, dim)
155. columns.append(f"{field.name} {pg_type}")

157. # Add primary key if exists
158. pk_fields = [f.name for f in collection.schema.fields if f.is_primary]
159. if pk_fields:
160. columns.append(f"PRIMARY KEY ({', '.join(pk_fields)})")

162. create_sql = f"CREATE TABLE {table_name} ({', '.join(columns)});"
163. cursor.execute(create_sql)
164. conn.commit()
165. logger.info(f"Created table: {table_name}")
166. except Exception as e:
167. conn.rollback()
168. logger.error(f"Table creation failed: {e}")
169. raise
170. finally:
171. cursor.close()

173. def _query_milvus_data(self, collection: Collection, limit: int, offset: int) -> List[Dict]:
174. """Query data with sparse vector support"""
175. try:
176. # Get field metadata
177. field_meta = {f.name: {"type": f.dtype.name, "dim": getattr(f, 'dim', None)}
178. for f in collection.schema.fields}

180. # Adjust query window
181. if offset + limit > self.MAX_WINDOW_SIZE:
182. new_limit = self.MAX_WINDOW_SIZE - offset
183. if new_limit <= 0:
184. return []
185. limit = new_limit
186. logger.warning(f"Adjusted limit to {limit}")

188. # Query data
189. results = collection.query(
190. expr="",
191. output_fields=self.fields,
192. limit=limit,
193. offset=offset
194. )

196. # Process results
197. processed = []
198. for row in results:
199. processed_row = {}
200. for field in self.fields:
201. meta = field_meta[field]
202. value = row.get(field)
203. processed_row[field] = self._process_field_value(
204. value, meta["type"], meta["dim"])
205. processed.append(processed_row)

207. return processed
208. except Exception as e:
209. logger.error(f"Query failed: {e}")
210. raise

212. def _export_to_csv_chunked(self, collection: Collection) -> List[str]:
213. """Export data to CSV chunks"""
214. file_paths = []
215. total_rows = 0
216. offset = 0

218. try:
219. # Flush and get row count
220. collection.flush()
221. total_count = collection.num_entities
222. logger.info(f"Total rows to export: {total_count}")

224. while total_rows < total_count:
225. # Create chunk file
226. chunk_id = len(file_paths) + 1
227. chunk_file = f"{os.path.splitext(self.csv_file_path)[0]}_part{chunk_id}.csv"
228. file_paths.append(chunk_file)

230. with open(chunk_file, 'w', newline='') as f:
231. writer = csv.DictWriter(f, fieldnames=self.fields)
232. writer.writeheader()

234. chunk_rows = 0
235. while chunk_rows < 100000:  # 100K rows per chunk
236. batch = self._query_milvus_data(
237. collection,
238. limit=min(50000, 100000 - chunk_rows),
239. offset=offset
240. )
241. if not batch:
242. break

244. writer.writerows(batch)
245. chunk_rows += len(batch)
246. total_rows += len(batch)
247. offset += len(batch)

249. if total_rows % 10000 == 0:
250. logger.info(f"Exported {total_rows}/{total_count} rows")

252. logger.info(f"Created chunk {chunk_id}: {chunk_rows} rows")

254. return file_paths
255. except Exception as e:
256. # Cleanup failed export
257. for f in file_paths:
258. try:
259. os.remove(f)
260. except:
261. pass
262. logger.error(f"Export failed: {e}")
263. raise

265. def _import_to_opengauss(self, conn: psycopg2.extensions.connection, file_paths: List[str]) -> None:
266. """Import CSV data to openGauss"""
267. table_name = self.config.get('Table', 'opengauss_table_name')
268. cursor = conn.cursor()

270. try:
271. # Prepare for bulk import
272. cursor.execute(f"TRUNCATE TABLE {table_name};")
273. conn.commit()

275. total_rows = 0
276. for i, csv_file in enumerate(file_paths, 1):
277. with open(csv_file, 'r') as f:
278. # Use COPY for bulk load
279. copy_sql = f"""
280. COPY {table_name} ({', '.join(self.fields)})
281. FROM STDIN WITH (FORMAT CSV, HEADER, NULL 'NULL');
282. """
283. cursor.copy_expert(copy_sql, f)
284. conn.commit()

286. rows_imported = cursor.rowcount
287. total_rows += rows_imported
288. logger.info(f"Imported {rows_imported} rows from {csv_file}")

290. logger.info(f"Total imported: {total_rows} rows")
291. except Exception as e:
292. conn.rollback()
293. logger.error(f"Import failed: {e}")
294. raise
295. finally:
296. cursor.close()

298. def run_migration(self) -> None:
299. """Execute full migration workflow"""
300. start_time = datetime.now()
301. logger.info("Starting migration")

303. try:
304. # Step 1: Connect to Milvus
305. milvus_collection = self._connect_milvus()

307. # Step 2: Export data to CSV
308. csv_files = self._export_to_csv_chunked(milvus_collection)

310. # Step 3: Connect to openGauss and create table
311. opengauss_conn = self._connect_opengauss()
312. self._create_opengauss_table(opengauss_conn, milvus_collection)

314. # Step 4: Import to openGauss
315. self._import_to_opengauss(opengauss_conn, csv_files)

317. # Cleanup
318. if self.config.getboolean('Migration', 'cleanup_temp_files', fallback=True):
319. for f in csv_files:
320. try:
321. os.remove(f)
322. except Exception as e:
323. logger.warning(f"Failed to delete {f}: {e}")

325. logger.info(f"Migration completed in {datetime.now() - start_time}")
326. except Exception as e:
327. logger.error(f"Migration failed: {e}")
328. raise
329. finally:
330. if 'opengauss_conn' in locals():
331. opengauss_conn.close()
332. connections.disconnect("default")

334. @staticmethod
335. def _milvus_to_opengauss_type(milvus_type: str, dim: Optional[int] = None) -> str:
336. """Map Milvus types to openGauss types"""
337. type_map = {
338. "Int64": "BIGINT",
339. "Int32": "INTEGER",
340. "Int16": "SMALLINT",
341. "Int8": "SMALLINT",
342. "Float": "REAL",
343. "Double": "DOUBLE PRECISION",
344. "Bool": "BOOLEAN",
345. "VarChar": "VARCHAR",
346. "String": "TEXT",
347. "Json": "JSONB",
348. "FLOAT_VECTOR": f"VECTOR({dim})" if dim else "VECTOR",
349. "BINARY_VECTOR": f"BIT({dim})" if dim else "BIT",
350. "SPARSE_FLOAT_VECTOR": "SPARSEVEC"
351. }
352. return type_map.get(milvus_type, "TEXT")

355. if __name__ == "__main__":
356. try:
357. migrator = MilvusToOpenGaussMigrator()
358. migrator.run_migration()
359. except Exception as e:
360. logger.error(f"Migration failed: {e}")
361. exit(1)

将迁移脚本、配置文件放置在同一目录下，层级如下：

1. ├── milvus2datavec.py
2. └── config.ini

运行迁移脚本，并查看输出结果：

1. python3 milvus2datavec.py

登录openGauss，查看数据是否完成迁移：

4.1 进入容器：

1. $ docker exec -it <CONTAINER ID> bash

4.2 登录omm超级用户：

1. $ su omm 2. $ gsql -d postgres -p 5432

4.3 查看迁移表数据量：

1. $ select count(*) from test;