Milvus 数据批量导入实战:Python代码解析

发布于:2025-03-06 ⋅ 阅读:(13) ⋅ 点赞:(0)

1 引言

在处理大规模数据的存储和检索时,向量数据库逐渐成为一种热门的解决方案。Milvus 作为一款高性能的向量数据库,在人工智能、机器学习等领域有着广泛的应用。本文将介绍如何使用 Python 代码将数据批量导入到 Milvus 数据库中,通过实际的代码示例来帮助大家理解导入过程和相关的技术要点。

2 代码功能概述

我们的代码主要实现了从本地文件读取数据,并将其批量导入到 Milvus 数据库的功能。代码涉及到命令行参数的解析、与 Milvus 服务器的连接、数据处理以及数据插入等操作。

3 代码解析

3.1 导入必要的库

import sys
import json
import argparse
import pandas as pd
from pymilvus import connections, FieldSchema, DataType, Collection

我们导入了sys用于系统相关的操作,json用于处理 JSON 格式的数据,argparse用于解析命令行参数,pandas用于数据处理,pymilvus库则用于与 Milvus 数据库进行交互。

3.2 数据导入函数

def import_data(data_list, collection_name, fields):
    df = pd.DataFrame(data_list)
    entities = []
    for field in fields:
        if 'id' == field.name:
            continue
        entities.append(df[field.name].to_list())

    collection = Collection(name=collection_name)
    insert_result = collection.insert(entities)
    collection.flush()
    return insert_result.insert_count

import_data函数接受数据列表、集合名称和字段定义作为参数。它首先将数据列表转换为pandas的DataFrame,然后根据字段定义,将除id字段外的数据提取出来组成entities列表。接着连接到指定的集合,插入数据并刷新,最后返回插入的行数。

3.3 命令行参数解析函数

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--collection", type=str, help="milvus collection名称", required=True)
    parser.add_argument("--host", type=str, help="milvus服务器IP", required=True)
    parser.add_argument("--port", type=int, default=19530, help="milvus服务器端口")
    parser