[VectorDB] 밀버스(milvus) DB 데이터 csv 추출 튜토리얼

안녕하세요? 오랜만에 밀버스 데이터 베이스에서 csv파일로 데이터를 추출하는 코드를 작성해보겠습니다.

---

필요한 모듈 설치

pip install pymilvus

 

Milvus와 연결 후 데이터를 CSV로 저장하는 코드입니다.

 

필요한 모듈 임포트

from pymilvus import ( Collection, connections )
import json
import csv

 

MilvusManager 선언

class MilvusManager:
    def __init__(self):
        self.dbname = "default"
        self.connection_args = {
            "uri": "주소",  # Milvus 서버 주소와 포트
            "db_name": self.dbname
        }

        self.collection_name = "test_sparse"  # 데이터가 저장된 컬렉션 이름

 

연결(connect) 선언

def connect(self):
	try:
		connections.connect(
            alias=self.dbname,
            **self.connection_args
        )
        if connections.has_connection(self.dbname):
            print("Milvus와 연결되었습니다.")
            return True
        else:
            print("Milvus 연결 실패: 연결이 확인되지 않았습니다.")
            return False

    except Exception as e:
            print(f"Milvus 연결 실패: {e}")
            return False

 

export_data_to_csv 선언

def export_data_to_csv(self, output_file="milvus_data.csv"):
    if not self.connect():
        print("Milvus 연결 실패: 연결이 확인되지 않았습니다.")
        return

    try:
        # 컬렉션 연결 및 데이터 조회
        collection = Collection(name=self.collection_name, using=self.dbname)

        # `pk` 필드를 제외한 나머지 필드로 데이터 조회
        output_fields = ["text", "text_vector", "sparse_vector"]
        query = collection.query(expr="", output_fields=output_fields, limit=10000)
        # limit 값은 최대로 가져올 수를 지정합니다.

        # CSV 파일로 저장
        with open(output_file, mode='w', newline='') as file:
        # 인코딩 에러가 발생할시 아래의 utf-8로 구동할것.
        # with open(output_file, mode='w', newline='', encoding='utf-8', errors='replace') as file:
            writer = csv.writer(file, escapechar='\\', quoting=csv.QUOTE_MINIMAL, quotechar='"')  # 따옴표를 사용하여 필드를 감쌈

            # CSV 헤더 작성 (필드 순서에 맞춰 작성)
            writer.writerow(output_fields)

            # 각 데이터 행을 CSV로 작성
            for item in query:
                # sparse_vector를 JSON 문자열로 변환
                item['sparse_vector'] = json.dumps(item['sparse_vector'])  # 딕셔너리를 JSON 문자열로 변환
                row = [item[field] for field in output_fields]
                writer.writerow(row)

        print(f"데이터가 '{output_file}' 파일로 저장되었습니다.")
    except Exception as e:
        print(f"데이터 추출 중 오류 발생: {e}")

 

코드 실행

if __name__ == "__main__":
    manager = MilvusManager()
    manager.connect()
    manager.export_data_to_csv("저장할 파일 경로")
    # ex) /workspace/milvus_data.csv

 

주요 코드 설명

1. Milvus 연결: MilvusManager 클래스는 Milvus에 연결 및 데이터를 CSV로 추출 하는 기능을 제공합니다.

connect() 메소를 토앻 Milvus서버와 연결을 시도, 연결을 성공할시 True값을 반환합니다.

 

2. 컬렉션(Collection)과 데이터 조회(search): export_data_to_csv() 메소드는 Milvus에 연결된 후, 데이터를 조회는 합니다.

이때 output으로 지정한 text, text_vector, sparse_vector필드만 추출하며, 갯수 조정이 필요할시 limit값을 추가 설정하면 됩니다.

 

3. csv 파일 저자: 데이터를 CSV 파일로 저장할 때, sparse_vector필드는 JSON 문자열로 변환한 후 저장합니다.

이전에 그냥 저장했다가 데이터 삽입할때 값이 안들어가지는 오류가 발생하여 JSON 문자열로 변환하는 코드를 추가하였습니다.

 

인코딩 문제

• 만약 CSV 파일을 저장하는 도중 인코딩 에러가 발생할 시 UTF-8 인코딩을 지정하고, error='replace'옵션을 추가하여 인코딩 문제를 해결 할 수 있습니다.
  원하시는 분은 주석해제 후 사용하시면 됩니다.
• 데이터크기 : Milvus에서 추출하는 데이터가 매우 클 수 있으며, limit 값을 적절히 설정하시면 되겠습니다.
  
  
  

오늘은 Milvus 데이터를 csv파일로 추출 후 저장하는 코드를 작성해봤습니다.

데이터를 좀 더 확실하고 정확하게 판단하기 위해 csv 파일로 추출하는 방법도 필요하다고 판단하여

이렇게 코드를 작성했습니다.

 

위 코드와 과정을 사용하여 여러분들의 Milvus 데이터를 쉽게 CSV로 추출 및 변경해보세요!