from pymilvus import connections, Collection, db
import pandas as pd
from openpyxl import Workbook

# 连接 Milvus
connections.connect(
    uri="http://172.18.30.165:19530",
    token="root:Milvus"
)

# 使用数据库
db.using_database("XIAN")

# 获取集合
collection = Collection("ce_shi_3_an_zi_duan_token_fen_a1da59")
collection.load()  # 确保集合已加载

# 创建查询迭代器
iterator = collection.query_iterator(
    batch_size=10,  # 可以调整批量大小以提高效率
    expr="id > 0",
    output_fields=["*"]  # 指定要获取的字段
)

results = []

# 遍历所有结果
try:
    while True:
        # 获取下一批结果
        res = iterator.next()
        # print(res)
        
        # 如果没有更多结果，退出循环
        if not res:
            break
            
        # 将结果添加到列表中
        results.extend(res)
        print(f"已获取 {len(results)} 条记录")
        
finally:
    # 确保迭代器被关闭
    iterator.close()
# 将结果转换为DataFrame
df = pd.DataFrame(results)

# 检查是否有数据
if not df.empty:
    # 写入Excel文件
    excel_file = "milvus_query_results.xlsx"
    
    # 使用openpyxl引擎以获得更好的格式控制
    with pd.ExcelWriter(excel_file, engine='openpyxl') as writer:
        df.to_excel(writer, index=False, sheet_name='查询结果')
        
        # 获取工作簿和工作表对象以进行格式设置
        workbook = writer.book
        worksheet = writer.sheets['查询结果']
        
        # 设置列宽
        for column in worksheet.columns:
            max_length = 0
            column_letter = column[0].column_letter
            for cell in column:
                try:
                    if len(str(cell.value)) > max_length:
                        max_length = len(cell.value)
                except:
                    pass
            adjusted_width = (max_length + 2) * 1.2
            worksheet.column_dimensions[column_letter].width = adjusted_width
    
    print(f"结果已成功写入 {excel_file}")
else:
    print("没有查询到任何结果")