kgrag/app_pkg/examples/search_vector.py
2026-06-30 13:35:52 +08:00

76 lines
2.1 KiB
Python

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("没有查询到任何结果")