跳到主要内容
Open In Colab GitHub Repository

使用 VectorETL 高效地将数据加载到 Milvus

在本教程中，我们将探讨如何使用 VectorETL 高效地将数据加载到 Milvus 中，VectorETL 是一个专为向量数据库设计的轻量级 ETL 框架。VectorETL 简化了从各种来源提取数据、使用 AI 模型将其转换为向量嵌入以及将其存储在 Milvus 中进行快速和可扩展检索的过程。在本教程结束时，您将拥有一个工作的 ETL 流水线，可以轻松地集成和管理向量搜索系统。让我们开始吧！

准备工作​

依赖和环境​

$ pip install --upgrade vector-etl pymilvus

如果您使用 Google Colab，为了启用刚安装的依赖项，您可能需要重启运行时（点击屏幕顶部的"Runtime"菜单，并从下拉菜单中选择"Restart session"）。

VectorETL 支持多种数据源，包括 Amazon S3、Google Cloud Storage、本地文件等。您可以在这里查看支持的数据源的完整列表。在本教程中，我们将重点介绍 Amazon S3 作为数据源示例。

我们将从 Amazon S3 加载文档。因此，您需要准备 AWS_ACCESS_KEY_ID 和 AWS_SECRET_ACCESS_KEY 作为环境变量，以安全地访问您的 S3 存储桶。此外，我们将使用 OpenAI 的 text-embedding-ada-002 嵌入模型为数据生成嵌入。您还应该准备 api key OPENAI_API_KEY 作为环境变量。

import os

os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
os.environ["AWS_ACCESS_KEY_ID"] = "your-aws-access-key-id"
os.environ["AWS_SECRET_ACCESS_KEY"] = "your-aws-secret-access-key"

工作流程​

定义数据源（Amazon S3）​

在这种情况下，我们从 Amazon S3 存储桶中提取文档。VectorETL 允许我们指定存储桶名称、文件路径以及我们正在处理的数据类型。

source = {
    "source_data_type": "Amazon S3",
    "bucket_name": "my-bucket",
    "key": "path/to/files/",
    "file_type": ".csv",
    "aws_access_key_id": os.environ["AWS_ACCESS_KEY_ID"],
    "aws_secret_access_key": os.environ["AWS_SECRET_ACCESS_KEY"],
}

配置嵌入模型（OpenAI）​

一旦我们设置了数据源，我们需要定义嵌入模型，该模型将把我们的文本数据转换为向量嵌入。在这里，我们在这个示例中使用 OpenAI 的 text-embedding-ada-002。

embedding = {
    "embedding_model": "OpenAI",
    "api_key": os.environ["OPENAI_API_KEY"],
    "model_name": "text-embedding-ada-002",
}

将 Milvus 设置为目标数据库​

我们需要将生成的嵌入存储在 Milvus 中。在这里，我们使用 Milvus Lite 定义我们的 Milvus 连接参数。

target = {
    "target_database": "Milvus",
    "host": "./milvus.db",  # 如果使用 Zilliz Cloud，则为 os.environ["ZILLIZ_CLOUD_PUBLIC_ENDPOINT"]
    "api_key": "",  # 如果使用 Zilliz Cloud，则为 os.environ["ZILLIZ_CLOUD_TOKEN"]
    "collection_name": "my_collection",
    "vector_dim": 1536,  # text-embedding-ada-002 为 1536
}

对于 host 和 api_key：

  • 将 host 设置为本地文件，例如 ./milvus.db，并将 api_key 留空是最方便的方法，因为它会自动利用 Milvus Lite 将所有数据存储在此文件中。

  • 如果您有大规模数据，可以在 docker 或 kubernetes 上设置性能更高的 Milvus 服务器。在此设置中，请使用服务器 uri，例如 http://localhost:19530，作为您的 host 并将 api_key 留空。

  • 如果您想使用 Zilliz Cloud，Milvus 的完全托管云服务，请调整 host 和 api_key，它们对应于 Zilliz Cloud 中的公共端点和 API 密钥。

指定嵌入的列​

现在，我们需要指定 CSV 文件中的哪些列应该转换为嵌入。这确保只处理相关的文本字段，从而优化效率和存储。

embed_columns = ["col_1", "col_2", "col_3"]

创建并执行 VectorETL 流水线​

在所有配置就绪后，我们现在初始化 ETL 流水线，设置数据流并执行它。

from vector_etl import create_flow

flow = create_flow()
flow.set_source(source)
flow.set_embedding(embedding)
flow.set_target(target)
flow.set_embed_columns(embed_columns)

# 执行流程
flow.execute()

通过遵循本教程，我们已经成功构建了一个端到端的 ETL 流水线，使用 VectorETL 将文档从 Amazon S3 移动到 Milvus。VectorETL 在数据源方面非常灵活，因此您可以根据您的特定应用需求选择您喜欢的任何数据源。凭借 VectorETL 的模块化设计，您可以轻松扩展此流水线以支持其他数据源和嵌入模型，使其成为 AI 和数据工程工作流程的强大工具！